Combine McCode runs as one run of all their rays, weighted by Ncount

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):

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)
 
 

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'))