Merge pull request #255 from scipp/fix-broken-chunking

Fix broken chunking when computing Tof lookup table

Author: nvaytet

src/ess/reduce/time_of_flight/lut.py

@@ -4,6 +4,7 @@
 Utilities for computing time-of-flight lookup tables from neutron simulations.
 """
 
+import math
 from dataclasses import dataclass
 from typing import NewType
 
@@ -330,12 +331,11 @@ def make_tof_lookup_table(
     ndist = len(distance_bins) - 1
     max_size = 2e7
     total_size = ndist * len(simulation.time_of_arrival)
-    nchunks = total_size / max_size
-    chunk_size = int(ndist / nchunks) + 1
+    nchunks = math.ceil(total_size / max_size)
+    chunk_size = math.ceil(ndist / nchunks)
     pieces = []
-    for i in range(int(nchunks) + 1):
+    for i in range(nchunks):
         dist_edges = distance_bins[i * chunk_size : (i + 1) * chunk_size + 1]
-
         pieces.append(
             _compute_mean_tof_in_distance_range(
                 simulation=simulation,

tests/time_of_flight/lut_test.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2025 Scipp contributors (https://github.com/scipp)
+import pytest
+import scipp as sc
+
+from ess.reduce import time_of_flight
+from ess.reduce.time_of_flight import TofLookupTableWorkflow
+
+sl = pytest.importorskip("sciline")
+
+
+def test_lut_workflow_computes_table():
+    wf = TofLookupTableWorkflow()
+    wf[time_of_flight.DiskChoppers] = {}
+    wf[time_of_flight.SourcePosition] = sc.vector([0, 0, 0], unit='m')
+    wf[time_of_flight.NumberOfSimulatedNeutrons] = 100_000
+    wf[time_of_flight.SimulationSeed] = 60
+    wf[time_of_flight.PulseStride] = 1
+
+    lmin, lmax = sc.scalar(25.0, unit='m'), sc.scalar(35.0, unit='m')
+    dres = sc.scalar(0.1, unit='m')
+    tres = sc.scalar(333.0, unit='us')
+
+    wf[time_of_flight.LtotalRange] = lmin, lmax
+    wf[time_of_flight.DistanceResolution] = dres
+    wf[time_of_flight.TimeResolution] = tres
+
+    table = wf.compute(time_of_flight.TimeOfFlightLookupTable)
+
+    assert table.coords['distance'].min() < lmin
+    assert table.coords['distance'].max() > lmax
+    assert table.coords['event_time_offset'].max() == sc.scalar(1 / 14, unit='s').to(
+        unit=table.coords['event_time_offset'].unit
+    )
+    assert sc.isclose(table.coords['distance_resolution'], dres)
+    # Note that the time resolution is not exactly preserved since we want the table to
+    # span exactly the frame period.
+    assert sc.isclose(table.coords['time_resolution'], tres, rtol=sc.scalar(0.01))
+
+
+def test_lut_workflow_computes_table_in_chunks():
+    wf = TofLookupTableWorkflow()
+    wf[time_of_flight.DiskChoppers] = {}
+    wf[time_of_flight.SourcePosition] = sc.vector([0, 0, 0], unit='m')
+    # Lots of neutrons activates chunking
+    wf[time_of_flight.NumberOfSimulatedNeutrons] = 1_000_000
+    wf[time_of_flight.SimulationSeed] = 61
+    wf[time_of_flight.PulseStride] = 1
+
+    lmin, lmax = sc.scalar(25.0, unit='m'), sc.scalar(35.0, unit='m')
+    dres = sc.scalar(0.1, unit='m')
+    tres = sc.scalar(250.0, unit='us')
+
+    wf[time_of_flight.LtotalRange] = lmin, lmax
+    wf[time_of_flight.DistanceResolution] = dres
+    wf[time_of_flight.TimeResolution] = tres
+
+    table = wf.compute(time_of_flight.TimeOfFlightLookupTable)
+
+    assert table.coords['distance'].min() < lmin
+    assert table.coords['distance'].max() > lmax
+    assert table.coords['event_time_offset'].max() == sc.scalar(1 / 14, unit='s').to(
+        unit=table.coords['event_time_offset'].unit
+    )
+    assert sc.isclose(table.coords['distance_resolution'], dres)
+    # Note that the time resolution is not exactly preserved since we want the table to
+    # span exactly the frame period.
+    assert sc.isclose(table.coords['time_resolution'], tres, rtol=sc.scalar(0.01))
+
+
+def test_lut_workflow_pulse_skipping():
+    wf = TofLookupTableWorkflow()
+    wf[time_of_flight.DiskChoppers] = {}
+    wf[time_of_flight.SourcePosition] = sc.vector([0, 0, 0], unit='m')
+    wf[time_of_flight.NumberOfSimulatedNeutrons] = 100_000
+    wf[time_of_flight.SimulationSeed] = 62
+    wf[time_of_flight.PulseStride] = 2
+
+    lmin, lmax = sc.scalar(55.0, unit='m'), sc.scalar(65.0, unit='m')
+    dres = sc.scalar(0.1, unit='m')
+    tres = sc.scalar(250.0, unit='us')
+
+    wf[time_of_flight.LtotalRange] = lmin, lmax
+    wf[time_of_flight.DistanceResolution] = dres
+    wf[time_of_flight.TimeResolution] = tres
+
+    table = wf.compute(time_of_flight.TimeOfFlightLookupTable)
+
+    assert table.coords['event_time_offset'].max() == 2 * sc.scalar(
+        1 / 14, unit='s'
+    ).to(unit=table.coords['event_time_offset'].unit)
+
+
+def test_lut_workflow_non_exact_distance_range():
+    wf = TofLookupTableWorkflow()
+    wf[time_of_flight.DiskChoppers] = {}
+    wf[time_of_flight.SourcePosition] = sc.vector([0, 0, 0], unit='m')
+    wf[time_of_flight.NumberOfSimulatedNeutrons] = 100_000
+    wf[time_of_flight.SimulationSeed] = 63
+    wf[time_of_flight.PulseStride] = 1
+
+    lmin, lmax = sc.scalar(25.0, unit='m'), sc.scalar(35.0, unit='m')
+    dres = sc.scalar(0.33, unit='m')
+    tres = sc.scalar(250.0, unit='us')
+
+    wf[time_of_flight.LtotalRange] = lmin, lmax
+    wf[time_of_flight.DistanceResolution] = dres
+    wf[time_of_flight.TimeResolution] = tres
+
+    table = wf.compute(time_of_flight.TimeOfFlightLookupTable)
+
+    assert table.coords['distance'].min() < lmin
+    assert table.coords['distance'].max() > lmax
+    assert sc.isclose(table.coords['distance_resolution'], dres)