Merge pull request #561 from scipp/multi-distance-frames

Fix propagation of chopper cascade frames to an array of distances

Author: nvaytet

docs/user-guide/wfm/wfm-time-of-flight.ipynb

@@ -312,21 +312,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "threshold = 0.01  # Select part of the pulse where signal is above 1% of max counts\n",
-    "\n",
-    "t = one_pulse.hist(time=300)\n",
-    "tsel = t.data > threshold * t.data.max()\n",
-    "pulse_times = sc.midpoints(t.coords[\"time\"])[tsel]\n",
-    "\n",
-    "w = one_pulse.hist(wavelength=300)\n",
-    "wsel = w.data > threshold * w.data.max()\n",
-    "pulse_wavs = sc.midpoints(w.coords[\"wavelength\"])[wsel]\n",
+    "time_min = sc.scalar(0.0, unit='ms')\n",
+    "time_max = sc.scalar(3.4, unit='ms')\n",
+    "wavs_min = sc.scalar(0.01, unit='angstrom')\n",
+    "wavs_max = sc.scalar(10.0, unit='angstrom')\n",
     "\n",
     "frames = chopper_cascade.FrameSequence.from_source_pulse(\n",
-    "    time_min=pulse_times.min(),\n",
-    "    time_max=pulse_times.max(),\n",
-    "    wavelength_min=pulse_wavs.min(),\n",
-    "    wavelength_max=pulse_wavs.max(),\n",
+    "    time_min=time_min,\n",
+    "    time_max=time_max,\n",
+    "    wavelength_min=wavs_min,\n",
+    "    wavelength_max=wavs_max,\n",
     ")"
    ]
   },
@@ -398,8 +393,8 @@
     "\n",
     "workflow[unwrap.PulsePeriod] = sc.reciprocal(ess_beamline.source.frequency)\n",
     "workflow[unwrap.PulseStride | None] = None\n",
-    "workflow[unwrap.SourceTimeRange] = pulse_times.min(), pulse_times.max()\n",
-    "workflow[unwrap.SourceWavelengthRange] = pulse_wavs.min(), pulse_wavs.max()\n",
+    "workflow[unwrap.SourceTimeRange] = time_min, time_max\n",
+    "workflow[unwrap.SourceWavelengthRange] = wavs_min, wavs_max\n",
     "workflow[unwrap.Choppers] = choppers\n",
     "\n",
     "workflow[unwrap.Ltotal] = Ltotal\n",

src/scippneutron/tof/chopper_cascade.py

@@ -56,7 +56,7 @@ class Subframe:
     """
 
     def __init__(self, time: sc.Variable, wavelength: sc.Variable):
-        if time.sizes != wavelength.sizes:
+        if {dim: time.sizes.get(dim) for dim in wavelength.sizes} != wavelength.sizes:
             raise sc.DimensionError(
                 f'Inconsistent dims or shape: {time.sizes} vs {wavelength.sizes}'
             )
@@ -106,22 +106,33 @@ def propagate_by(self, distance: sc.Variable) -> Subframe:
 
     @property
     def start_time(self) -> sc.Variable:
-        """The start time of the subframe."""
-        return self.time.min()
+        """The start time of the subframe for each of the distances in self.time."""
+        # The `self.time` may have an additional dimension for distance, compared to
+        # `self.wavelength`, and we need to keep that dimension in the output.
+        out = self.time
+        for dim in self.wavelength.dims:
+            out = out.min(dim)
+        return out
 
     @property
     def end_time(self) -> sc.Variable:
-        """The end time of the subframe."""
-        return self.time.max()
+        """The end time of the subframe for each of the distances in self.time."""
+        # The `self.time` may have an additional dimension for distance, compared to
+        # `self.wavelength`, and we need to keep that dimension in the output.
+        out = self.time
+        for dim in self.wavelength.dims:
+            out = out.max(dim)
+        return out
 
     @property
     def start_wavelength(self) -> sc.Variable:
-        """The start wavelength of the subframe."""
+        """The start wavelength of the subframe for each of the distances in
+        self.time"""
         return self.wavelength.min()
 
     @property
     def end_wavelength(self) -> sc.Variable:
-        """The end wavelength of the subframe."""
+        """The end wavelength of the subframe for each of the distances in self.time."""
         return self.wavelength.max()
 
 
@@ -206,19 +217,14 @@ def bounds(self) -> sc.DataGroup:
 
     def subbounds(self) -> sc.DataGroup:
         """
-        The bounds of the subframes, defined as the union over subframes.
-
-        This is not the same as the bounds of the individual subframes, but defined as
-        the union of all subframes. Subframes that overlap in time are "merged" into a
-        single subframe.
-
-        This function is to some extent experimental: It is not clear if taking the
-        union of overlapping subframes has any utility in practice, since this may
-        simply indicate a problem with the chopper cascade. Attempts to handle this
-        automatically may be misguided.
+        The bounds of the individual subframes, stored as a DataGroup.
         """
-        starts = [subframe.start_time for subframe in self.subframes]
-        ends = [subframe.end_time for subframe in self.subframes]
+        starts = sc.concat(
+            [subframe.start_time for subframe in self.subframes], dim='subframe'
+        )
+        ends = sc.concat(
+            [subframe.end_time for subframe in self.subframes], dim='subframe'
+        )
         # Given how time-propagation and chopping works, the min wavelength is always
         # given by the same vertex as the min time, and the max wavelength by the same
         # vertex as the max time. Thus, this check should generally always pass.
@@ -228,48 +234,21 @@ def subbounds(self) -> sc.DataGroup:
                 'Subframes must be regular, i.e., min/max time and wavelength must '
                 'coincide.'
             )
-        wav_starts = [subframe.start_wavelength for subframe in self.subframes]
-        wav_ends = [subframe.end_wavelength for subframe in self.subframes]
-
-        @dataclass
-        class Bound:
-            start: sc.Variable
-            end: sc.Variable
-            wav_start: sc.Variable
-            wav_end: sc.Variable
-
-        bounds = [
-            Bound(start, end, wav_start, wav_end)
-            for start, end, wav_start, wav_end in zip(
-                starts, ends, wav_starts, wav_ends, strict=True
-            )
-        ]
-        bounds = sorted(bounds, key=lambda x: x.start)
-        current = bounds[0]
-        merged_bounds = []
-        for bound in bounds[1:]:
-            # If start is before current end, merge
-            if bound.start <= current.end:
-                current = Bound(
-                    current.start,
-                    max(current.end, bound.end),
-                    current.wav_start,
-                    max(current.wav_end, bound.wav_end),
-                )
-            else:
-                merged_bounds.append(current)
-                current = bound
-        merged_bounds.append(current)
-        time_bounds = [
-            sc.concat([bound.start, bound.end], dim='bound') for bound in merged_bounds
-        ]
-        wav_bounds = [
-            sc.concat([bound.wav_start, bound.wav_end], dim='bound')
-            for bound in merged_bounds
-        ]
+        wav_starts = sc.concat(
+            [subframe.start_wavelength for subframe in self.subframes], dim='subframe'
+        )
+        wav_ends = sc.concat(
+            [subframe.end_wavelength for subframe in self.subframes], dim='subframe'
+        )
+
+        time = sc.concat([starts, ends], dim='bound')
+        wavelength = sc.concat([wav_starts, wav_ends], dim='bound')
+        time_dims = list(set(time.dims) - {'subframe', 'bound'})
+        wavelength_dims = list(set(wavelength.dims) - {'subframe', 'bound'})
+
         return sc.DataGroup(
-            time=sc.concat(time_bounds, dim='subframe'),
-            wavelength=sc.concat(wav_bounds, dim='subframe'),
+            time=time.transpose([*time_dims, 'subframe', 'bound']),
+            wavelength=wavelength.transpose([*wavelength_dims, 'subframe', 'bound']),
         )
 
 

tests/tof/chopper_cascade_test.py

@@ -391,6 +391,38 @@ def test_frame_sequence_propagate_to_returns_new_sequence_with_added_propagated_
     assert result2[2] == result[1].propagate_to(distance * 2)
 
 
+def test_frame_sequence_propagate_to_array_of_distances(
+    source_frame_sequence: chopper_cascade.FrameSequence,
+) -> None:
+    frames = source_frame_sequence
+    distances = sc.array(dims=['distance'], values=[1.5, 1.7], unit='m')
+    result = frames.propagate_to(distances)
+    assert len(frames) == 1
+    assert len(result) == 2
+    assert result[1] == frames[0].propagate_to(distances)
+    subframe = result[1].subframes[0]
+    assert subframe.start_time.sizes["distance"] == 2
+    assert subframe.end_time.sizes["distance"] == 2
+
+
+def test_frame_sequence_propagate_to_array_of_distances_subbounds(
+    source_frame_sequence: chopper_cascade.FrameSequence,
+) -> None:
+    frames = source_frame_sequence
+    distances = sc.array(dims=['distance'], values=[1.5, 1.7], unit='m')
+    result = frames.propagate_to(distances)
+    subbounds = result[-1].subbounds()
+    subframe = result[1].subframes[0]
+    assert_identical(subbounds['time']['subframe', 0]['bound', 0], subframe.start_time)
+    assert_identical(subbounds['time']['subframe', 0]['bound', 1], subframe.end_time)
+    assert_identical(
+        subbounds['wavelength']['subframe', 0]['bound', 0], subframe.start_wavelength
+    )
+    assert_identical(
+        subbounds['wavelength']['subframe', 0]['bound', 1], subframe.end_wavelength
+    )
+
+
 def test_frame_sequence_chop_returns_new_sequence_with_added_chopped_frames(
     source_frame_sequence: chopper_cascade.FrameSequence,
 ) -> None:

tests/tof/wfm_test.py

@@ -149,7 +149,7 @@ def test_compute_wavelengths_from_wfm(disk_choppers, npulses):
     # computed correctly
     workflow[unwrap.SourceTimeRange] = (
         sc.scalar(0.0, unit='ms'),
-        sc.scalar(4.0, unit='ms'),
+        sc.scalar(3.4, unit='ms'),
     )
     workflow[unwrap.SourceWavelengthRange] = (
         sc.scalar(0.01, unit='angstrom'),