diff --git a/src/mccode_antlr/loader/datfile.py b/src/mccode_antlr/loader/datfile.py index 7a99353..0a247f5 100644 --- a/src/mccode_antlr/loader/datfile.py +++ b/src/mccode_antlr/loader/datfile.py @@ -81,33 +81,55 @@ def _get(name, idx): return I_arr, I_err_arr, N_arr, extras -def _combine_values(var: list[str], a, b): - """Two files' values, (nv, ...) each, as the values of one longer run. - - Intensity and counts add and the intensity's error adds in quadrature. Every other - variable -- the axis column a 1-D file starts with, say ``Mean`` -- must agree, - since it describes the bins rather than what landed in them. - """ - import numpy as np - +def _standard_indices(var: list[str]) -> tuple: + """Where ``I``, ``I_err`` and ``N`` are among a file's variables, or None.""" def index(name, fallback): if name in var: return var.index(name) return fallback if fallback < len(var) else None + return index('I', 0), index('I_err', 1), index('N', 2) + + +def _combine_values(var: list[str], a, b, rays_a: float, rays_b: float): + """Two runs' values, (nv, ...) each, as the values of one run of both runs' rays. - i, e, n = index('I', 0), index('I_err', 1), index('N', 2) + McCode normalises each run's intensity to its own ray count, so the combined + intensity is the ray-count-weighted mean of the two, and its error is the weighted + errors added in quadrature, over the total. Counts add. Every other variable -- the + axis column a 1-D file starts with, say ``Mean`` -- must agree, since it describes + the bins rather than what landed in them. + """ + import numpy as np + i, e, n = _standard_indices(var) + total = rays_a + rays_b out = a.copy() - for k in (i, n): - if k is not None: - out[k] = a[k] + b[k] + if i is not None: + out[i] = (rays_a * a[i] + rays_b * b[i]) / total if e is not None: - out[e] = np.sqrt(a[e] ** 2 + b[e] ** 2) + out[e] = np.sqrt((rays_a * a[e]) ** 2 + (rays_b * b[e]) ** 2) / total + if n is not None: + out[n] = a[n] + b[n] for k, name in enumerate(var): if k not in (i, e, n) and not np.array_equal(a[k], b[k]): raise RuntimeError(f'Cannot combine files whose {name!r} values differ') return out +def _g(x) -> str: + """A number as McCode prints one in a header.""" + return f'{float(x):g}' + + +def _moments(axis, weight) -> tuple: + """The weighted mean and spread of ``axis``, as McCode's ``statistics`` gives them.""" + import numpy as np + total = weight.sum() + if not total: + return 0.0, 0.0 + mean = (axis * weight).sum() / total + return mean, np.sqrt(max((axis ** 2 * weight).sum() / total - mean ** 2, 0.0)) + + @dataclass(eq=False) class DatFileCommon: """Base class for McCode output data files. @@ -247,13 +269,45 @@ def safe_to_combine(self, other) -> bool: return False return self._limits is None or self.metadata.get(self._limits) == other.metadata.get(self._limits) + def rays(self) -> int: + """The run's ray count, ``Ncount``, which its intensities are normalised to.""" + if 'Ncount' not in self.metadata: + raise RuntimeError(f'{self.source} has no Ncount, so how much it should count ' + f'for in a combination is unknown') + return int(self.metadata['Ncount']) + def __add__(self, other): + """The two runs as one run of both runs' rays. + + Each run's intensities are normalised to its own ray count, so they are averaged, + weighted by those counts, not added; see `_combine_values`. The header totals -- + ``values``, ``signal`` and ``statistics`` -- are recomputed from the result. + """ if not self.safe_to_combine(other): raise RuntimeError(f'Cannot combine these two {type(self).__name__} instances') - return type(self)(self.source, - combine_scan_dicts(self.metadata, other.metadata), - combine_scan_dicts(self.parameters, other.parameters), - _combine_values(self.variables, self.values, other.values)) + combined = type(self)(self.source, + combine_scan_dicts(self.metadata, other.metadata), + combine_scan_dicts(self.parameters, other.parameters), + _combine_values(self.variables, self.values, other.values, + self.rays(), other.rays())) + combined.metadata.update({k: v for k, v in combined.header_totals().items() + if k in combined.metadata}) + return combined + + def header_totals(self) -> dict: + """The header's ``values`` line -- total intensity, its error and total counts -- + and, where the file has bins, ``signal``, from these values.""" + import numpy as np + i, e, n = _standard_indices(self.variables) + intensity = self.values[i] if i is not None else np.zeros(self.values.shape[1:]) + error = self.values[e] if e is not None else np.zeros_like(intensity) + counts = self.values[n] if n is not None else np.zeros_like(intensity) + totals = {'values': f'{_g(intensity.sum())} {_g(np.sqrt((error ** 2).sum()))} ' + f'{_g(counts.sum())}'} + if intensity.ndim: + totals['signal'] = (f'Min={_g(intensity.min())}; Max={_g(intensity.max())}; ' + f'Mean={_g(intensity.mean())};') + return totals # ------------------------------------------------------------------ # Misc @@ -379,6 +433,17 @@ def dim_metadata(self) -> list[dict]: lower_limit, upper_limit = [float(x) for x in self['xlimits'].split()] return [dim_metadata(self.values.shape[1], self['xlabel'], lower_limit, upper_limit)] + def header_totals(self) -> dict: + """As for every file, plus ``statistics``: the intensity-weighted mean and spread of + the axis column, ``X0`` and ``dX``.""" + totals = super().header_totals() + i, e, n = _standard_indices(self.variables) + axes = [k for k in range(len(self.variables)) if k not in (i, e, n)] + if axes and i is not None: + mean, spread = _moments(self.values[axes[0]], self.values[i]) + totals['statistics'] = f'X0={_g(mean)}; dX={_g(spread)};' + return totals + def print_data(self, file): # one line per bin for row in self.values.T: @@ -450,6 +515,24 @@ def dim_metadata(self) -> list[dict]: return [dim_metadata(nx, self['xlabel'], lower_x, upper_x), dim_metadata(ny, self['ylabel'], lower_y, upper_y)] + def header_totals(self) -> dict: + """As for every file, plus ``statistics``: the intensity-weighted mean and spread of + each axis, ``X0``, ``dX``, ``Y0`` and ``dY``, at the bin centres McCode uses.""" + import numpy as np + totals = super().header_totals() + i, _, _ = _standard_indices(self.variables) + if i is None or 'xylimits' not in self.metadata: + return totals + lower_x, upper_x, lower_y, upper_y = [float(x) for x in self['xylimits'].split()] + intensity = self.values[i] + ny, nx = intensity.shape + x = lower_x + (np.arange(nx) + 0.5) * (upper_x - lower_x) / nx + y = lower_y + (np.arange(ny) + 0.5) * (upper_y - lower_y) / ny + x0, dx = _moments(x, intensity.sum(axis=0)) + y0, dy = _moments(y, intensity.sum(axis=1)) + totals['statistics'] = f'X0={_g(x0)}; dX={_g(dx)}; Y0={_g(y0)}; dY={_g(dy)};' + return totals + def print_data(self, file): labels = 'Data', 'Errors', 'Events' for i, name in enumerate(self.variables): diff --git a/tests/test_datfile_scipp.py b/tests/test_datfile_scipp.py index 2abee97..d141028 100644 --- a/tests/test_datfile_scipp.py +++ b/tests/test_datfile_scipp.py @@ -193,9 +193,10 @@ def test_add_combines_correctly(self): dat1 = self._load() dat2 = self._load() # independent load combined = dat1 + dat2 - self.assertAlmostEqual(float(combined.dataset['I'].value), 2 * 1234.5, places=3) - # Independent measurements: variances add - expected_var = 2 * (6.78 ** 2) + # Two runs of the same ray count, each normalised to it: the intensity is their + # mean, and the variance of a mean of two independent runs is half of one + self.assertAlmostEqual(float(combined.dataset['I'].value), 1234.5, places=3) + expected_var = (6.78 ** 2) / 2 self.assertAlmostEqual(float(combined.dataset['I'].variance), expected_var, places=3) @@ -274,11 +275,12 @@ def test_add_propagates_variance(self): dat1 = self._load() dat2 = self._load() # independent load combined = dat1 + dat2 - expected_var = 2 * np.array([0.1, 0.2, 0.3, 0.4, 0.5]) ** 2 + # equal ray counts: the mean intensity, with half the variance of one run + expected_var = np.array([0.1, 0.2, 0.3, 0.4, 0.5]) ** 2 / 2 np.testing.assert_allclose( combined.dataset['I'].variances, expected_var, rtol=1e-5 ) - np.testing.assert_allclose(combined.dataset['I'].values, [2, 4, 6, 8, 10]) + np.testing.assert_allclose(combined.dataset['I'].values, [1, 2, 3, 4, 5]) class TestDatFile2D(unittest.TestCase): @@ -326,9 +328,10 @@ def test_add_propagates_variance(self): dat1 = self._load() dat2 = self._load() # independent load combined = dat1 + dat2 - expected_I = 2 * np.array([[10, 20, 30], [40, 50, 60]], dtype=float) + # equal ray counts: the mean intensity, with half the variance of one run + expected_I = np.array([[10, 20, 30], [40, 50, 60]], dtype=float) np.testing.assert_allclose(combined.dataset['I'].values, expected_I) - expected_var = 2 * np.array([[1, 2, 3], [4, 5, 6]], dtype=float) ** 2 + expected_var = np.array([[1, 2, 3], [4, 5, 6]], dtype=float) ** 2 / 2 np.testing.assert_allclose(combined.dataset['I'].variances, expected_var) diff --git a/tests/test_datfile_without_scipp.py b/tests/test_datfile_without_scipp.py index ec5fda6..64e5760 100644 --- a/tests/test_datfile_without_scipp.py +++ b/tests/test_datfile_without_scipp.py @@ -91,19 +91,51 @@ def test_the_scipp_view_says_what_it_needs(no_scipp, tmp_path): def test_runs_combine_without_scipp(no_scipp, tmp_path): - """Intensity and counts add, errors add in quadrature, and the axis column stays.""" + """Two runs of the same instrument combine as one run of all their rays. + + McCode normalises each run's intensity to its own ray count, so intensities are + averaged, weighted by Ncount; errors are the weighted errors in quadrature over the + total; counts add; the axis column stays. + """ from mccode_antlr.loader.datfile import read_mccode_dat, write_combined_mccode_dats one = _write(tmp_path / 'one.dat', ncount=1000, seed=1, directory='one') two = _write(tmp_path / 'two.dat', ncount=3000, seed=2, directory='two') write_combined_mccode_dats([one, two], tmp_path / 'both.dat') both = read_mccode_dat(tmp_path / 'both.dat') np.testing.assert_allclose(both['L'], [1, 2, 3]) - np.testing.assert_allclose(both['I'], [2, 4, 6]) - np.testing.assert_allclose(both['I_err'], np.sqrt(2) * np.array([0.3, 0.4, 0.5])) + np.testing.assert_allclose(both['I'], [1, 2, 3]) + np.testing.assert_allclose(both['I_err'], np.sqrt(10) / 4 * np.array([0.3, 0.4, 0.5])) np.testing.assert_allclose(both['N'], [20, 40, 60]) assert both['Ncount'] == '4000' +def test_runs_are_weighted_by_their_ray_counts(no_scipp, tmp_path): + """A run of three times the rays counts three times as much, and the header's + totals describe the combined histogram, not the first run's.""" + from mccode_antlr.loader.datfile import read_mccode_dat + one = read_mccode_dat(_write(tmp_path / 'one.dat', ncount=1000)) + two = read_mccode_dat(_write(tmp_path / 'two.dat', ncount=3000, seed=2, directory='two')) + two.values[1] = [5.0, 6.0, 7.0] # a different intensity in each bin + both = one + two + np.testing.assert_allclose(both['I'], [(1 + 3 * 5) / 4, (2 + 3 * 6) / 4, (3 + 3 * 7) / 4]) + total, error, events = (float(x) for x in both['values'].split()) + assert total == pytest.approx(both['I'].sum(), rel=1e-5) + assert error == pytest.approx(np.sqrt((both['I_err'] ** 2).sum()), rel=1e-5) + assert events == pytest.approx(120) + mean = (both['L'] * both['I']).sum() / both['I'].sum() + assert both['statistics'].startswith(f'X0={mean:g};') + assert both['signal'] == f"Min={both['I'].min():g}; Max={both['I'].max():g}; Mean={both['I'].mean():g};" + + +def test_a_run_without_a_ray_count_is_not_combined(no_scipp, tmp_path): + from mccode_antlr.loader.datfile import read_mccode_dat + one = read_mccode_dat(_write(tmp_path / 'one.dat')) + two = read_mccode_dat(_write(tmp_path / 'two.dat', seed=2, directory='two')) + del two.metadata['Ncount'] + with pytest.raises(RuntimeError, match='Ncount'): + one + two + + def test_files_with_different_bins_do_not_combine(no_scipp, tmp_path): from mccode_antlr.loader.datfile import read_mccode_dat one = read_mccode_dat(_write(tmp_path / 'one.dat'))