fix: convert unlimited TOA to expected range (#209)

* fix: convert unlimited TOA to expected range

* fix: add masking in monitor correction

* fix: use dim coord not 'tof'

* test: add masks to expected

* test: loaded monitor has correct 'tof' range

Author: jokasimr

src/ess/dream/io/geant4.py

@@ -196,6 +196,27 @@ def _to_edges(centers: sc.Variable) -> sc.Variable:
     )
 
 
+def _rebin_to_tofrange(da):
+    '''Rebins a monitor TOA histogram to the range of values
+    expected from a real instrument at ESS.
+    That is, to the range [0, 1/14] s.
+
+    Strategy:
+        1. Create a grid that alingns with the pulse times.
+        2. Rebin TOA on that grid.
+        3. Fold pulses into new dimension and sum over pulses.
+    '''
+    dim = da.dim
+    unit = da.coords[dim].unit
+    period = (1.0 / sc.scalar(14.0, unit='Hz')).to(unit=unit)
+    N = sc.sum(da.coords[dim] < period).value
+    K = (da.coords[dim].max() // period + 1).to(dtype='int')
+    grid = sc.linspace(dim, sc.scalar(0.0, unit=unit), K * period, K * (N - 1) + 1)
+    out = da.rebin({dim: grid}).fold(dim, sizes={'_': K, dim: N - 1}).sum('_')
+    out.coords[dim] = sc.linspace(dim, sc.scalar(0.0, unit=unit), period, N)
+    return out
+
+
 def load_mcstas_monitor(
     file_path: MonitorFilename[RunType],
     position: CaveMonitorPosition,
@@ -221,14 +242,14 @@ def load_mcstas_monitor(
     tof, counts, err = np.loadtxt(file_path, usecols=(0, 1, 2), unpack=True)
 
     tof = _to_edges(sc.array(dims=["tof"], values=tof, unit="us"))
+    data = sc.DataArray(
+        sc.array(dims=["tof"], values=counts, variances=err**2, unit="counts"),
+        coords={"tof": tof},
+    )
+    data = _rebin_to_tofrange(data)
     return NeXusComponent[CaveMonitor, RunType](
         sc.DataGroup(
-            data=sc.DataArray(
-                sc.array(dims=["tof"], values=counts, variances=err**2, unit="counts"),
-                coords={
-                    "tof": tof,
-                },
-            ),
+            data=data,
             position=position,
         )
     )

src/ess/powder/correction.py

@@ -59,11 +59,15 @@ def normalize_by_monitor_histogram(
     )
     lut = sc.lookup(norm, dim="wavelength")
     if detector.bins is None:
-        result = (
-            detector / lut[sc.midpoints(detector.coords['wavelength'], dim='dspacing')]
-        )
+        c = lut[sc.midpoints(detector.coords['wavelength'], dim='dspacing')]
+        result = detector / c
+        result.masks['monitor_intensity_is_zero'] = c == sc.scalar(0.0, unit=c.unit)
     else:
-        result = detector.bins / lut
+        c = lut(detector.bins.coords['wavelength'])
+        result = detector / c
+        result.bins.masks['monitor_intensity_is_zero'] = c == sc.scalar(
+            0.0, unit=c.unit
+        )
     return ScaledCountsDspacing[RunType](result)
 
 

tests/dream/io/geant4_test.py

@@ -10,8 +10,19 @@
 import scipp.testing
 import scippnexus as snx
 
-from ess.dream import DreamGeant4ProtonChargeWorkflow, data, load_geant4_csv
-from ess.powder.types import Filename, NeXusComponent, NeXusDetectorName, SampleRun
+from ess.dream import DreamGeant4ProtonChargeWorkflow, data, io, load_geant4_csv
+from ess.powder.types import (
+    CaveMonitorPosition,
+    Filename,
+    NeXusComponent,
+    NeXusDetectorName,
+    SampleRun,
+)
+
+
+@pytest.fixture(scope="module")
+def monitor_file_path():
+    return data.simulated_monitor_diamond_sample()
 
 
 @pytest.fixture(scope="module")
@@ -184,3 +195,15 @@ def test_geant4_in_workflow(file_path, file):
     expected = load_geant4_csv(file)["instrument"]["mantle"]["events"]
     expected.coords["detector"] = sc.scalar("mantle")
     sc.testing.assert_identical(detector, expected)
+
+
+def test_geant4_monitor_in_expected_range(monitor_file_path):
+    dummy_position = CaveMonitorPosition(sc.vector([0, 0, 0.0], unit='m'))
+    mon = io.geant4.load_mcstas_monitor(monitor_file_path, dummy_position)['data']
+    assert (
+        mon.coords['tof'] >= sc.scalar(0.0, unit='s').to(unit=mon.coords['tof'].unit)
+    ).all()
+    assert (
+        mon.coords['tof']
+        <= sc.scalar(1 / 14.0, unit='s').to(unit=mon.coords['tof'].unit)
+    ).all()

tests/powder/correction_test.py

@@ -347,7 +347,12 @@ def test_normalize_by_monitor_histogram_expected_results():
         monitor=WavelengthMonitor[SampleRun, CaveMonitor](monitor),
         uncertainty_broadcast_mode=UncertaintyBroadcastMode.fail,
     )
-    expected = ScaledCountsDspacing[SampleRun](detector / monitor.data)
+    # Simple way to get all false array same shape as detector data.
+    # Is there a better way?
+    all_false = sc.isinf(detector.bins.data)
+    expected = ScaledCountsDspacing[SampleRun](
+        detector / monitor.data
+    ).bins.assign_masks({'monitor_intensity_is_zero': all_false})
     sc.testing.assert_identical(normalized, expected)
 
 
@@ -362,12 +367,16 @@ def test_normalize_by_monitor_histogram_ignores_monitor_values_out_of_range():
             'wavelength': sc.array(dims=['wavelength'], values=[0.0, 3, 4], unit='Å')
         },
     )
+    # Simple way to get all false array same shape as detector data.
+    all_false = sc.isinf(detector.bins.data)
     normalized = normalize_by_monitor_histogram(
         CountsWavelength[SampleRun](detector),
         monitor=WavelengthMonitor[SampleRun, CaveMonitor](monitor),
         uncertainty_broadcast_mode=UncertaintyBroadcastMode.fail,
     )
-    expected = ScaledCountsDspacing[SampleRun](detector / sc.scalar(4.0, unit='counts'))
+    expected = ScaledCountsDspacing[SampleRun](
+        detector / sc.scalar(4.0, unit='counts')
+    ).bins.assign_masks({'monitor_intensity_is_zero': all_false})
     sc.testing.assert_identical(normalized, expected)