From 62b5376ac36e8edf50ccf4322de18df929acfaa7 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Sun, 22 Mar 2026 21:52:55 +0100
Subject: [PATCH 01/15] clone chopper cascade code from scippneutron

---
 src/tof/optimization.py | 577 ++++++++++++++++++++++++++++++++++++++++
 src/tof/utils.py        |  14 +
 2 files changed, 591 insertions(+)
 create mode 100644 src/tof/optimization.py

diff --git a/src/tof/optimization.py b/src/tof/optimization.py
new file mode 100644
index 0000000..3260906
--- /dev/null
+++ b/src/tof/optimization.py
@@ -0,0 +1,577 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+# @author Simon Heybrock
+"""
+Compute result of applying a chopper cascade to a neutron pulse at a time-of-flight
+neutron source.
+
+See :py:class:`FrameSequence` for the main entry point.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import scipp as sc
+
+from .chopper import Chopper
+from .utils import wavelength_to_inverse_speed
+
+# from ..chopper import DiskChopper
+
+
+# def wavelength_to_inverse_velocity(wavelength):
+#     h = sc.constants.h
+#     m_n = sc.constants.m_n
+#     return (wavelength * m_n / h).to(unit='s/m')
+
+
+def propagate_times(
+    time: sc.Variable, wavelength: sc.Variable, distance: sc.Variable
+) -> sc.Variable:
+    """
+    Propagate a neutron frame by a distance.
+
+    Parameters
+    ----------
+    time:
+        Time of the neutron frame.
+    wavelength:
+        Wavelength of the neutron frame.
+    distance:
+        Distance to propagate. Can be a range of distances.
+
+    Returns
+    -------
+    :
+        Propagated time.
+    """
+    inverse_velocity = wavelength_to_inverse_speed(wavelength)
+    return time + (distance * inverse_velocity).to(unit=time.unit, copy=False)
+
+
+class Subframe:
+    """
+    Neutron "subframe" at a time-of-flight neutron source, described as the corners of a
+    polygon (initially a rectangle) in time and wavelength.
+    """
+
+    def __init__(self, time: sc.Variable, wavelength: sc.Variable):
+        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}'
+            )
+        self.time = time.to(unit='s', copy=False)
+        self.wavelength = wavelength.to(unit='angstrom', copy=False)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Subframe):
+            return NotImplemented
+        # Using sc.identical can lead to flaky behavior (different frames might have
+        # been obtained via a different operation order), so we use allclose instead.
+        same_time = sc.allclose(
+            self.time,
+            other.time,
+            rtol=sc.scalar(1.0e-12),
+            atol=sc.scalar(1.0e-16, unit=self.time.unit),
+        )
+        same_wavelength = sc.allclose(
+            self.wavelength,
+            other.wavelength,
+            rtol=sc.scalar(1.0e-12),
+            atol=sc.scalar(1.0e-16, unit=self.wavelength.unit),
+        )
+        return same_time and same_wavelength
+
+    def is_regular(self) -> bool:
+        """
+        Return True if the subframe is regular, i.e., if the vertex with the minimum
+        wavelength also has the minimum time, and the vertex with the maximum wavelength
+        also has the maximum time.
+        """
+        min_time = self.time == self.time.min()
+        min_wavelength = self.wavelength == self.wavelength.min()
+        max_time = self.time == self.time.max()
+        max_wavelength = self.wavelength == self.wavelength.max()
+        coinciding_min = min_time & min_wavelength
+        coinciding_max = max_time & max_wavelength
+        return coinciding_min.any() and coinciding_max.any()
+
+    def propagate_by(self, distance: sc.Variable) -> Subframe:
+        """
+        Propagate subframe by a distance.
+
+        Parameters
+        ----------
+        distance:
+            Distance to propagate. Note that this is a difference, not an absolute
+            value, in contrast to the distance in :py:meth:`Frame.propagate_to`.
+
+        Returns
+        -------
+        :
+            Propagated subframe.
+        """
+        return Subframe(
+            time=propagate_times(self.time, self.wavelength, distance),
+            wavelength=self.wavelength,
+        )
+
+    @property
+    def start_time(self) -> sc.Variable:
+        """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 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 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 for each of the distances in self.time."""
+        return self.wavelength.max()
+
+
+@dataclass
+class Frame:
+    """
+    A frame of neutrons, created from a single neutron pulse, potentially chopped into
+    subframes by choppers.
+    """
+
+    distance: sc.Variable
+    subframes: list[Subframe]
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Frame):
+            return NotImplemented
+        # Note: we don't use sc.identical here because it can fail on the dtype even
+        # if the values are equal.
+        same_distance = (
+            np.array_equal(self.distance.values, other.distance.values)
+            and self.distance.unit == other.distance.unit
+        )
+        return same_distance and all(
+            self_sub == other_sub
+            for self_sub, other_sub in zip(self.subframes, other.subframes, strict=True)
+        )
+
+    def propagate_to(self, distance: sc.Variable) -> Frame:
+        """
+        Compute new frame by propagating to a distance.
+
+        Parameters
+        ----------
+        distance:
+            New distance.
+
+        Returns
+        -------
+        :
+            Propagated frame.
+        """
+        delta = distance.to(unit=self.distance.unit, copy=False) - self.distance
+        subframes = [subframe.propagate_by(delta) for subframe in self.subframes]
+        return Frame(distance=distance, subframes=subframes)
+
+    def chop(self, chopper: Chopper) -> Frame:
+        """
+        Compute a new frame by applying a chopper.
+
+        A frame is a polygon in time and wavelength. Its initial shape is distorted
+        by propagation to the chopper. The chopper then cuts off the parts of the frame
+        that is outside of the chopper opening. Here we apply and algorithm that
+        computes a new polygon that is the intersection of the frame and the chopper
+        opening.
+
+        In practice a chopper may have multiple openings, so a frame may be chopped into
+        a number of subframes.
+
+        Parameters
+        ----------
+        chopper:
+            Chopper to apply.
+
+        Returns
+        -------
+        :
+            Chopped frame.
+        """
+        distance = chopper.distance.to(unit=self.distance.unit, copy=False)
+        if distance < self.distance:
+            raise ValueError(
+                f'Chopper distance {distance} is smaller than frame distance '
+                f'{self.distance}'
+            )
+        frame = self.propagate_to(distance)
+
+        # A chopper can have multiple openings, call _chop for each of them. The result
+        # is the union of the resulting subframes.
+        chopped = Frame(distance=frame.distance, subframes=[])
+        open_times, close_times = (t.to(unit='s') for t in chopper.open_close_times())
+        for subframe in frame.subframes:
+            # for open, close in zip(chopper.time_open, chopper.time_close, strict=True):
+            for open, close in zip(open_times, close_times, strict=True):
+                if (tmp := _chop(subframe, open, close_to_open=True)) is not None:
+                    if (tmp := _chop(tmp, close, close_to_open=False)) is not None:
+                        chopped.subframes.append(tmp)
+        return chopped
+
+    def bounds(self) -> sc.DataGroup:
+        """The bounds of the frame, i.e., the global min and max time and wavelength."""
+        start = sc.reduce([sub.start_time for sub in self.subframes]).min()
+        end = sc.reduce([sub.end_time for sub in self.subframes]).max()
+        wav_start = sc.reduce([sub.start_wavelength for sub in self.subframes]).min()
+        wav_end = sc.reduce([sub.end_wavelength for sub in self.subframes]).max()
+        return sc.DataGroup(
+            time=sc.concat([start, end], dim='bound'),
+            wavelength=sc.concat([wav_start, wav_end], dim='bound'),
+        )
+
+    def subbounds(self) -> sc.DataGroup:
+        """
+        The bounds of the individual subframes, stored as a DataGroup.
+        """
+        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.
+        # Exceptions may be subframes that have been created manually.
+        if not all(subframe.is_regular() for subframe in self.subframes):
+            raise NotImplementedError(
+                'Subframes must be regular, i.e., min/max time and wavelength must '
+                'coincide.'
+            )
+        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=time.transpose([*time_dims, 'subframe', 'bound']),
+            wavelength=wavelength.transpose([*wavelength_dims, 'subframe', 'bound']),
+        )
+
+
+@dataclass
+class FrameSequence:
+    """
+    A sequence of frames, created from a single neutron pulse, potentially chopped into
+    subframes by choppers.
+
+    It is recommended to use the :py:meth:`from_source_pulse` constructor to create a
+    frame sequence from a source pulse. Then, a chopper cascade can be applied using
+    :py:meth:`chop`.
+    """
+
+    frames: list[Frame]
+
+    @staticmethod
+    def from_source_pulse(
+        time_min: sc.Variable,
+        time_max: sc.Variable,
+        wavelength_min: sc.Variable,
+        wavelength_max: sc.Variable,
+    ):
+        """
+        Initialize a frame sequence from min/max time and wavelength of a pulse.
+
+        The distance is set to 0 m.
+        """
+        time = sc.concat([time_min, time_max, time_max, time_min], dim='vertex').to(
+            unit='s'
+        )
+        wavelength = sc.concat(
+            [wavelength_min, wavelength_min, wavelength_max, wavelength_max],
+            dim='vertex',
+        ).to(unit='angstrom')
+        frames = [
+            Frame(
+                distance=sc.scalar(0, unit='m'),
+                subframes=[Subframe(time=time, wavelength=wavelength)],
+            )
+        ]
+        return FrameSequence(frames)
+
+    def __len__(self) -> int:
+        """Number of frames."""
+        return len(self.frames)
+
+    def __getitem__(self, item: int | sc.Variable) -> Frame:
+        """Get a frame by index or distance."""
+        if isinstance(item, int):
+            return self.frames[item]
+        distance = item.to(unit='m')
+        frame_before_detector = None
+        for frame in self:
+            if frame.distance > distance:
+                break
+            frame_before_detector = frame
+
+        return frame_before_detector.propagate_to(distance)
+
+    def propagate_to(self, distance: sc.Variable) -> FrameSequence:
+        """
+        Propagate the frame sequence to a distance, adding a new frame.
+
+        Use this, e.g., to propagate to the sample position after applying choppers.
+
+        Parameters
+        ----------
+        distance:
+            Distance to propagate.
+
+        Returns
+        -------
+        :
+            New frame sequence.
+        """
+        return FrameSequence([*self.frames, self.frames[-1].propagate_to(distance)])
+
+    def chop(self, choppers: list[Chopper]) -> FrameSequence:
+        """
+        Chop the frame sequence by a list of choppers.
+
+        The choppers will be sorted by their distance, and applied in order.
+
+        Parameters
+        ----------
+        choppers:
+            List of choppers.
+
+        Returns
+        -------
+        :
+            New frame sequence.
+        """
+        choppers = sorted(choppers, key=lambda x: x.distance)
+        frames = list(self.frames)
+        for chopper in choppers:
+            frames.append(frames[-1].chop(chopper))
+        return FrameSequence(frames)
+
+    def draw(
+        self,
+        linewidth: float = 0,
+        fill: bool = True,
+        alpha: float | None = None,
+        transpose: bool = False,
+        colors: list[str] | None = None,
+        grid: bool = True,
+        title: str = 'Frame propagation through chopper cascade',
+        time_unit: str = 'ms',
+        wavelength_unit: str = 'angstrom',
+    ) -> Any:
+        """
+        Draw frames using matplotlib.
+
+        Parameters
+        ----------
+        linewidth:
+            Line width of frame edges.
+        fill:
+            Fill frame with color.
+        alpha:
+            Transparency of frame.
+        transpose:
+            Transpose axes.
+        colors:
+            List of colors to use for frames. If None, use default matplotlib colors.
+        grid:
+            Show grid.
+        time_unit:
+            Unit for time axis. Default is ms.
+        wavelength_unit:
+            Unit for wavelength axis. Default is angstrom.
+        """
+        import matplotlib.colors as mcolors
+        import matplotlib.patches as patches
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots()
+        x_max = 0
+        y_max = 0
+        for i, frame in enumerate(self.frames):
+            color = colors[i] if colors is not None else f'C{i}'
+            # Add label to legend
+            ax.plot([], [], color=color, label=f'{frame.distance:c}')
+            # All subframes have same color
+            for subframe in frame.subframes:
+                x = subframe.time.to(unit=time_unit, copy=False)
+                y = subframe.wavelength.to(unit=wavelength_unit, copy=False)
+                if transpose:
+                    x, y = y, x
+                x_unit = x.unit
+                y_unit = y.unit
+                x_max = max(x_max, x.max().value)
+                y_max = max(y_max, y.max().value)
+                if alpha:
+                    color = mcolors.to_rgba(color, alpha=alpha)
+                polygon = patches.Polygon(
+                    np.stack((x.values, y.values), axis=1),
+                    closed=True,
+                    fill=fill,
+                    color=color,
+                    linewidth=linewidth,
+                )
+                ax.add_patch(polygon)
+        ax.set_xlabel(x_unit)
+        ax.set_ylabel(y_unit)
+        ax.set_xlim(0, x_max)
+        ax.set_ylim(0, y_max)
+        ax.minorticks_on()
+        if grid:
+            ax.grid(True, linestyle='-', linewidth='0.5', color='gray')
+            ax.grid(which='minor', linestyle=':', linewidth='0.5', color='gray')
+        ax.legend(loc='best')
+        ax.set_title(title)
+        fig.tight_layout()
+        return fig, ax
+
+    def acceptance_diagram(self, transpose: bool = True):
+        """
+        Draw a chopper acceptance diagram.
+
+        See, e.g., J.R.D. Copley, An acceptance diagram analysis of the contaminant
+        pulse removal problem with direct geometry neutron chopper spectrometers,
+        https://doi.org/10.1016/S0168-9002(03)01731-5 for more background.
+
+        Parameters
+        ----------
+        transpose:
+            Transpose axes if ``True``, i.e., show wavelength on horizontal axis and
+            time on vertical axis.
+        """
+        import matplotlib.pyplot as plt
+
+        source_distance = self.frames[0].distance
+        frames = FrameSequence(
+            [frame.propagate_to(source_distance) for frame in self.frames]
+        )
+        # Reset frame distance for plotting labels. This is a bit of a hack.
+        # We should use chopper names if available.
+        for i, frame in enumerate(frames.frames):
+            frame.distance = self.frames[i].distance
+        blue = np.linspace(0.1, 1, len(frames.frames))
+        colors = plt.cm.Blues(blue)
+        fig, ax = frames.draw(
+            fill=True,
+            linewidth=0.5,
+            transpose=transpose,
+            colors=colors,
+            title='Chopper acceptance diagram',
+        )
+        # Put legend outside of plot
+        box = ax.get_position()
+        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
+        ax.legend(
+            loc='center left',
+            bbox_to_anchor=(1, 0.5),
+            title='Distance',
+            frameon=False,
+        )
+        return fig, ax
+
+
+# @dataclass
+# class Chopper:
+#     distance: sc.Variable
+#     time_open: sc.Variable
+#     time_close: sc.Variable
+
+#     def __post_init__(self):
+#         if self.time_open.sizes != self.time_close.sizes:
+#             raise sc.DimensionError(
+#                 f'Inconsistent dims or shape: {self.time_open.sizes} vs '
+#                 f'{self.time_close.sizes}'
+#             )
+
+#     def __getitem__(self, key) -> Chopper:
+#         return Chopper(
+#             distance=self.distance,
+#             time_open=self.time_open[key],
+#             time_close=self.time_close[key],
+#         )
+
+#     @classmethod
+#     def from_disk_chopper(
+#         cls,
+#         disk_chopper: DiskChopper,
+#         pulse_frequency: sc.Variable,
+#         npulses: int,
+#     ) -> Chopper:
+#         """
+#         Create a chopper from a NeXus disk chopper.
+
+#         Parameters
+#         ----------
+#         disk_chopper:
+#             Disk chopper.
+#         pulse_frequency:
+#             The frequency of pulses in the neutron source (this is to determine how
+#             many rotations the chopper performs per pulse).
+#         npulses:
+#             Number of pulses to rotate the chopper for.
+#         """
+#         tpulse = 1.0 / pulse_frequency
+#         topen = disk_chopper.time_offset_open(pulse_frequency=pulse_frequency)
+#         tclose = disk_chopper.time_offset_close(pulse_frequency=pulse_frequency)
+#         offsets = sc.arange('pulse', npulses) * tpulse
+#         return cls(
+#             distance=sc.norm(disk_chopper.axle_position),
+#             time_open=(offsets + topen).flatten(to=topen.dim),
+#             time_close=(offsets + tclose).flatten(to=tclose.dim),
+#         )
+
+
+def _chop(frame: Subframe, time: sc.Variable, close_to_open: bool) -> Subframe | None:
+    inside = frame.time >= time if close_to_open else frame.time <= time
+    output = []
+    for i in range(len(frame.time)):
+        # Note how j wraps around to 0
+        j = (i + 1) % len(frame.time)
+        inside_i = inside[i]
+        inside_j = inside[j]
+        if inside_i:
+            output.append((frame.time[i], frame.wavelength[i]))
+        if inside_i != inside_j:
+            # Intersection
+            t = (time - frame.time[i]) / (frame.time[j] - frame.time[i])
+            v = (1 - t) * frame.wavelength[i] + t * frame.wavelength[j]
+            output.append((time, v))
+    if not output:
+        return None
+    time = sc.concat([t for t, _ in output], dim=frame.time.dim)
+    wavelength = sc.concat([v for _, v in output], dim=frame.wavelength.dim)
+    return Subframe(time=time, wavelength=wavelength)
diff --git a/src/tof/utils.py b/src/tof/utils.py
index 8805fe5..55955e4 100644
--- a/src/tof/utils.py
+++ b/src/tof/utils.py
@@ -42,6 +42,20 @@ def wavelength_to_speed(x: sc.Variable, unit: str = "m/s") -> sc.Variable:
     return (1.0 / (m_over_h * x)).to(unit=unit)
 
 
+def wavelength_to_inverse_speed(x: sc.Variable, unit: str = "s/m") -> sc.Variable:
+    """
+    Convert neutron wavelengths to inverse speeds.
+
+    Parameters
+    ----------
+    x:
+        Input wavelengths.
+    unit:
+        The unit of the output inverse speeds.
+    """
+    return (x * m_over_h).to(unit=unit)
+
+
 def speed_to_energy(x: sc.Variable, unit="meV") -> sc.Variable:
     """
     Convert neutron speeds to energies.

From 3272b2f6763d2999db29e08f76dc7b43fa2856dc Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Sun, 22 Mar 2026 23:12:06 +0100
Subject: [PATCH 02/15] working version of optimize_for

---
 src/tof/optimization.py | 475 ++++++++++++++--------------------------
 src/tof/source.py       | 156 +++++++++++--
 2 files changed, 302 insertions(+), 329 deletions(-)

diff --git a/src/tof/optimization.py b/src/tof/optimization.py
index 3260906..a2651c7 100644
--- a/src/tof/optimization.py
+++ b/src/tof/optimization.py
@@ -1,17 +1,16 @@
 # SPDX-License-Identifier: BSD-3-Clause
-# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
-# @author Simon Heybrock
+# Copyright (c) 2026 Scipp contributors (https://github.com/scipp)
 """
 Compute result of applying a chopper cascade to a neutron pulse at a time-of-flight
 neutron source.
 
-See :py:class:`FrameSequence` for the main entry point.
+This is copied over from scippneutron/tof/chopper_cascade.py, with the unused parts
+stripped out (see https://github.com/scipp/scippneutron).
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
 
 import numpy as np
 import scipp as sc
@@ -19,14 +18,6 @@
 from .chopper import Chopper
 from .utils import wavelength_to_inverse_speed
 
-# from ..chopper import DiskChopper
-
-
-# def wavelength_to_inverse_velocity(wavelength):
-#     h = sc.constants.h
-#     m_n = sc.constants.m_n
-#     return (wavelength * m_n / h).to(unit='s/m')
-
 
 def propagate_times(
     time: sc.Variable, wavelength: sc.Variable, distance: sc.Variable
@@ -66,39 +57,6 @@ def __init__(self, time: sc.Variable, wavelength: sc.Variable):
         self.time = time.to(unit='s', copy=False)
         self.wavelength = wavelength.to(unit='angstrom', copy=False)
 
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Subframe):
-            return NotImplemented
-        # Using sc.identical can lead to flaky behavior (different frames might have
-        # been obtained via a different operation order), so we use allclose instead.
-        same_time = sc.allclose(
-            self.time,
-            other.time,
-            rtol=sc.scalar(1.0e-12),
-            atol=sc.scalar(1.0e-16, unit=self.time.unit),
-        )
-        same_wavelength = sc.allclose(
-            self.wavelength,
-            other.wavelength,
-            rtol=sc.scalar(1.0e-12),
-            atol=sc.scalar(1.0e-16, unit=self.wavelength.unit),
-        )
-        return same_time and same_wavelength
-
-    def is_regular(self) -> bool:
-        """
-        Return True if the subframe is regular, i.e., if the vertex with the minimum
-        wavelength also has the minimum time, and the vertex with the maximum wavelength
-        also has the maximum time.
-        """
-        min_time = self.time == self.time.min()
-        min_wavelength = self.wavelength == self.wavelength.min()
-        max_time = self.time == self.time.max()
-        max_wavelength = self.wavelength == self.wavelength.max()
-        coinciding_min = min_time & min_wavelength
-        coinciding_max = max_time & max_wavelength
-        return coinciding_min.any() and coinciding_max.any()
-
     def propagate_by(self, distance: sc.Variable) -> Subframe:
         """
         Propagate subframe by a distance.
@@ -119,37 +77,6 @@ def propagate_by(self, distance: sc.Variable) -> Subframe:
             wavelength=self.wavelength,
         )
 
-    @property
-    def start_time(self) -> sc.Variable:
-        """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 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 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 for each of the distances in self.time."""
-        return self.wavelength.max()
-
 
 @dataclass
 class Frame:
@@ -161,20 +88,6 @@ class Frame:
     distance: sc.Variable
     subframes: list[Subframe]
 
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Frame):
-            return NotImplemented
-        # Note: we don't use sc.identical here because it can fail on the dtype even
-        # if the values are equal.
-        same_distance = (
-            np.array_equal(self.distance.values, other.distance.values)
-            and self.distance.unit == other.distance.unit
-        )
-        return same_distance and all(
-            self_sub == other_sub
-            for self_sub, other_sub in zip(self.subframes, other.subframes, strict=True)
-        )
-
     def propagate_to(self, distance: sc.Variable) -> Frame:
         """
         Compute new frame by propagating to a distance.
@@ -236,53 +149,6 @@ def chop(self, chopper: Chopper) -> Frame:
                         chopped.subframes.append(tmp)
         return chopped
 
-    def bounds(self) -> sc.DataGroup:
-        """The bounds of the frame, i.e., the global min and max time and wavelength."""
-        start = sc.reduce([sub.start_time for sub in self.subframes]).min()
-        end = sc.reduce([sub.end_time for sub in self.subframes]).max()
-        wav_start = sc.reduce([sub.start_wavelength for sub in self.subframes]).min()
-        wav_end = sc.reduce([sub.end_wavelength for sub in self.subframes]).max()
-        return sc.DataGroup(
-            time=sc.concat([start, end], dim='bound'),
-            wavelength=sc.concat([wav_start, wav_end], dim='bound'),
-        )
-
-    def subbounds(self) -> sc.DataGroup:
-        """
-        The bounds of the individual subframes, stored as a DataGroup.
-        """
-        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.
-        # Exceptions may be subframes that have been created manually.
-        if not all(subframe.is_regular() for subframe in self.subframes):
-            raise NotImplementedError(
-                'Subframes must be regular, i.e., min/max time and wavelength must '
-                'coincide.'
-            )
-        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=time.transpose([*time_dims, 'subframe', 'bound']),
-            wavelength=wavelength.transpose([*wavelength_dims, 'subframe', 'bound']),
-        )
-
 
 @dataclass
 class FrameSequence:
@@ -381,179 +247,6 @@ def chop(self, choppers: list[Chopper]) -> FrameSequence:
             frames.append(frames[-1].chop(chopper))
         return FrameSequence(frames)
 
-    def draw(
-        self,
-        linewidth: float = 0,
-        fill: bool = True,
-        alpha: float | None = None,
-        transpose: bool = False,
-        colors: list[str] | None = None,
-        grid: bool = True,
-        title: str = 'Frame propagation through chopper cascade',
-        time_unit: str = 'ms',
-        wavelength_unit: str = 'angstrom',
-    ) -> Any:
-        """
-        Draw frames using matplotlib.
-
-        Parameters
-        ----------
-        linewidth:
-            Line width of frame edges.
-        fill:
-            Fill frame with color.
-        alpha:
-            Transparency of frame.
-        transpose:
-            Transpose axes.
-        colors:
-            List of colors to use for frames. If None, use default matplotlib colors.
-        grid:
-            Show grid.
-        time_unit:
-            Unit for time axis. Default is ms.
-        wavelength_unit:
-            Unit for wavelength axis. Default is angstrom.
-        """
-        import matplotlib.colors as mcolors
-        import matplotlib.patches as patches
-        import matplotlib.pyplot as plt
-
-        fig, ax = plt.subplots()
-        x_max = 0
-        y_max = 0
-        for i, frame in enumerate(self.frames):
-            color = colors[i] if colors is not None else f'C{i}'
-            # Add label to legend
-            ax.plot([], [], color=color, label=f'{frame.distance:c}')
-            # All subframes have same color
-            for subframe in frame.subframes:
-                x = subframe.time.to(unit=time_unit, copy=False)
-                y = subframe.wavelength.to(unit=wavelength_unit, copy=False)
-                if transpose:
-                    x, y = y, x
-                x_unit = x.unit
-                y_unit = y.unit
-                x_max = max(x_max, x.max().value)
-                y_max = max(y_max, y.max().value)
-                if alpha:
-                    color = mcolors.to_rgba(color, alpha=alpha)
-                polygon = patches.Polygon(
-                    np.stack((x.values, y.values), axis=1),
-                    closed=True,
-                    fill=fill,
-                    color=color,
-                    linewidth=linewidth,
-                )
-                ax.add_patch(polygon)
-        ax.set_xlabel(x_unit)
-        ax.set_ylabel(y_unit)
-        ax.set_xlim(0, x_max)
-        ax.set_ylim(0, y_max)
-        ax.minorticks_on()
-        if grid:
-            ax.grid(True, linestyle='-', linewidth='0.5', color='gray')
-            ax.grid(which='minor', linestyle=':', linewidth='0.5', color='gray')
-        ax.legend(loc='best')
-        ax.set_title(title)
-        fig.tight_layout()
-        return fig, ax
-
-    def acceptance_diagram(self, transpose: bool = True):
-        """
-        Draw a chopper acceptance diagram.
-
-        See, e.g., J.R.D. Copley, An acceptance diagram analysis of the contaminant
-        pulse removal problem with direct geometry neutron chopper spectrometers,
-        https://doi.org/10.1016/S0168-9002(03)01731-5 for more background.
-
-        Parameters
-        ----------
-        transpose:
-            Transpose axes if ``True``, i.e., show wavelength on horizontal axis and
-            time on vertical axis.
-        """
-        import matplotlib.pyplot as plt
-
-        source_distance = self.frames[0].distance
-        frames = FrameSequence(
-            [frame.propagate_to(source_distance) for frame in self.frames]
-        )
-        # Reset frame distance for plotting labels. This is a bit of a hack.
-        # We should use chopper names if available.
-        for i, frame in enumerate(frames.frames):
-            frame.distance = self.frames[i].distance
-        blue = np.linspace(0.1, 1, len(frames.frames))
-        colors = plt.cm.Blues(blue)
-        fig, ax = frames.draw(
-            fill=True,
-            linewidth=0.5,
-            transpose=transpose,
-            colors=colors,
-            title='Chopper acceptance diagram',
-        )
-        # Put legend outside of plot
-        box = ax.get_position()
-        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
-        ax.legend(
-            loc='center left',
-            bbox_to_anchor=(1, 0.5),
-            title='Distance',
-            frameon=False,
-        )
-        return fig, ax
-
-
-# @dataclass
-# class Chopper:
-#     distance: sc.Variable
-#     time_open: sc.Variable
-#     time_close: sc.Variable
-
-#     def __post_init__(self):
-#         if self.time_open.sizes != self.time_close.sizes:
-#             raise sc.DimensionError(
-#                 f'Inconsistent dims or shape: {self.time_open.sizes} vs '
-#                 f'{self.time_close.sizes}'
-#             )
-
-#     def __getitem__(self, key) -> Chopper:
-#         return Chopper(
-#             distance=self.distance,
-#             time_open=self.time_open[key],
-#             time_close=self.time_close[key],
-#         )
-
-#     @classmethod
-#     def from_disk_chopper(
-#         cls,
-#         disk_chopper: DiskChopper,
-#         pulse_frequency: sc.Variable,
-#         npulses: int,
-#     ) -> Chopper:
-#         """
-#         Create a chopper from a NeXus disk chopper.
-
-#         Parameters
-#         ----------
-#         disk_chopper:
-#             Disk chopper.
-#         pulse_frequency:
-#             The frequency of pulses in the neutron source (this is to determine how
-#             many rotations the chopper performs per pulse).
-#         npulses:
-#             Number of pulses to rotate the chopper for.
-#         """
-#         tpulse = 1.0 / pulse_frequency
-#         topen = disk_chopper.time_offset_open(pulse_frequency=pulse_frequency)
-#         tclose = disk_chopper.time_offset_close(pulse_frequency=pulse_frequency)
-#         offsets = sc.arange('pulse', npulses) * tpulse
-#         return cls(
-#             distance=sc.norm(disk_chopper.axle_position),
-#             time_open=(offsets + topen).flatten(to=topen.dim),
-#             time_close=(offsets + tclose).flatten(to=tclose.dim),
-#         )
-
 
 def _chop(frame: Subframe, time: sc.Variable, close_to_open: bool) -> Subframe | None:
     inside = frame.time >= time if close_to_open else frame.time <= time
@@ -575,3 +268,165 @@ def _chop(frame: Subframe, time: sc.Variable, close_to_open: bool) -> Subframe |
     time = sc.concat([t for t, _ in output], dim=frame.time.dim)
     wavelength = sc.concat([v for _, v in output], dim=frame.wavelength.dim)
     return Subframe(time=time, wavelength=wavelength)
+
+
+def polygon_grid_overlap_mask(
+    polygon: np.ndarray, X: np.ndarray, Y: np.ndarray, tol: float | None = None
+) -> np.ndarray:
+    """
+    Compute a mask indicating which grid cells overlap with a polygon.
+    This is used to select relevant regions of the source distribution where the
+    selection is made from the polygons generated by the FrameSequence.
+
+    Parameters
+    ----------
+    polygon : (N,2) array
+        Coordinates of the polygon vertices.
+    X, Y    : (H,W) array
+        Coordinates of the grid corners.
+    tol     : float, optional
+        Tolerance for intersection checks (default = small fraction of grid spacing).
+
+    Notes:
+        Generated by ChatGPT-5.2.
+
+        Original prompt:
+        "Using numpy, can you make a function that would take in a polygon as a set
+        of vertices, and a 2d grid of xy coordinates, and return a mask that would
+        indicate which squares in the grid overlap with the polygon (even in the
+        slightest of ways). The math should be exact, I can't afford to approximate
+        the polygon with a cloud of points.
+        I don't need to know the area of overlap between squares and the polygon,
+        just whether there is overlap or not (True/False).
+
+        Refinement:
+        Can you:
+        1. Integrate bounding-box pruning per edge directly into the function
+        2. Add some tolerance when checking for equalities (maybe something like a
+        fraction of the grid spacing) because it can sometimes go wrong when a
+        polygon edge is perfectly aligned with a grid edge
+
+    Returns:
+        mask (H-1, W-1)
+    """
+
+    H, W = X.shape
+
+    # --- Cell bounds ---
+    cells_xmin = X[:-1, :-1]
+    cells_xmax = X[1:, 1:]
+    cells_ymin = Y[:-1, :-1]
+    cells_ymax = Y[1:, 1:]
+
+    # --- Auto tolerance ---
+    if tol is None:
+        dx = np.min(np.diff(X, axis=1))
+        dy = np.min(np.diff(Y, axis=0))
+        tol = 1e-9 + 1e-6 * min(dx, dy)
+
+    # Expand cell bounds slightly (robustness)
+    cells_xmin -= tol
+    cells_xmax += tol
+    cells_ymin -= tol
+    cells_ymax += tol
+
+    # --- Point in polygon ---
+    def points_in_poly(px, py, poly):
+        x = poly[:, 0]
+        y = poly[:, 1]
+        x2 = np.roll(x, -1)
+        y2 = np.roll(y, -1)
+
+        inside = np.zeros(px.shape, dtype=bool)
+
+        for i in range(len(poly)):
+            cond = ((y[i] > py) != (y2[i] > py)) & (
+                px < (x2[i] - x[i]) * (py - y[i]) / (y2[i] - y[i] + tol) + x[i]
+            )
+            inside ^= cond
+
+        return inside
+
+    # --- 1. Corner inside polygon ---
+    corners_x = np.stack([X[:-1, :-1], X[1:, :-1], X[:-1, 1:], X[1:, 1:]], axis=0)
+    corners_y = np.stack([Y[:-1, :-1], Y[1:, :-1], Y[:-1, 1:], Y[1:, 1:]], axis=0)
+
+    corner_inside = np.any(points_in_poly(corners_x, corners_y, polygon), axis=0)
+
+    # --- 2. Polygon vertex inside cell ---
+    vx = polygon[:, 0][:, None, None]
+    vy = polygon[:, 1][:, None, None]
+
+    vertex_inside = (
+        (vx >= cells_xmin)
+        & (vx <= cells_xmax)
+        & (vy >= cells_ymin)
+        & (vy <= cells_ymax)
+    ).any(axis=0)
+
+    # --- Segment intersection ---
+    def segments_intersect(p1, p2, q1, q2):
+        def orient(a, b, c):
+            return (b[..., 0] - a[..., 0]) * (c[..., 1] - a[..., 1]) - (
+                b[..., 1] - a[..., 1]
+            ) * (c[..., 0] - a[..., 0])
+
+        o1 = orient(p1, p2, q1)
+        o2 = orient(p1, p2, q2)
+        o3 = orient(q1, q2, p1)
+        o4 = orient(q1, q2, p2)
+
+        return (o1 * o2 <= tol) & (o3 * o4 <= tol)
+
+    edges_p1 = polygon
+    edges_p2 = np.roll(polygon, -1, axis=0)
+
+    intersect_mask = np.zeros((H - 1, W - 1), dtype=bool)
+
+    # Cell edges
+    cell_edges = [
+        (
+            np.stack([cells_xmin, cells_ymin], axis=-1),
+            np.stack([cells_xmax, cells_ymin], axis=-1),
+        ),
+        (
+            np.stack([cells_xmin, cells_ymax], axis=-1),
+            np.stack([cells_xmax, cells_ymax], axis=-1),
+        ),
+        (
+            np.stack([cells_xmin, cells_ymin], axis=-1),
+            np.stack([cells_xmin, cells_ymax], axis=-1),
+        ),
+        (
+            np.stack([cells_xmax, cells_ymin], axis=-1),
+            np.stack([cells_xmax, cells_ymax], axis=-1),
+        ),
+    ]
+
+    # --- 3. Edge intersection with pruning ---
+    for p1, p2 in zip(edges_p1, edges_p2, strict=True):
+        # Edge bounding box (+ tolerance)
+        xmin = min(p1[0], p2[0]) - tol
+        xmax = max(p1[0], p2[0]) + tol
+        ymin = min(p1[1], p2[1]) - tol
+        ymax = max(p1[1], p2[1]) + tol
+
+        # Candidate cells only
+        candidate = (
+            (cells_xmax >= xmin)
+            & (cells_xmin <= xmax)
+            & (cells_ymax >= ymin)
+            & (cells_ymin <= ymax)
+        )
+
+        if not np.any(candidate):
+            continue
+
+        p1e = p1[None, None, :]
+        p2e = p2[None, None, :]
+
+        for q1, q2 in cell_edges:
+            hit = segments_intersect(p1e, p2e, q1, q2)
+            intersect_mask[candidate] |= hit[candidate]
+
+    return corner_inside | vertex_inside | intersect_mask
diff --git a/src/tof/source.py b/src/tof/source.py
index fc0489f..b9a76ba 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -9,11 +9,15 @@
 import plopp as pp
 import scipp as sc
 
+from .chopper import Chopper
 from .component import ComponentReading
+from .optimization import FrameSequence, polygon_grid_overlap_mask
 from .utils import wavelength_to_speed
 
 TIME_UNIT = "us"
 WAV_UNIT = "angstrom"
+T_DIM = "birth_time"
+W_DIM = "wavelength"
 
 
 def _default_frequency(frequency: sc.Variable | None, pulses: int) -> sc.Variable:
@@ -38,6 +42,28 @@ def _bin_edges_to_midpoints(
     )
 
 
+def _midpoints_to_edges(x, dim):
+    if x.sizes[dim] < 2:
+        half = sc.scalar(0.5, unit=x.unit)
+        return sc.concat([x[dim, 0:1] - half, x[dim, 0:1] + half], dim)
+    else:
+        center = sc.midpoints(x, dim=dim)
+        # Note: use range of 0:1 to keep dimension dim in the slice to avoid
+        # switching round dimension order in concatenate step.
+        left = center[dim, 0:1] - (x[dim, 1] - x[dim, 0])
+        right = center[dim, -1] + (x[dim, -1] - x[dim, -2])
+        return sc.concat([left, center, right], dim)
+
+
+def _compute_grid_spacing(x: sc.Variable, dim: str) -> sc.Variable:
+    out = sc.empty_like(x)
+    dx = x[dim, 1:] - x[dim, :-1]
+    out[dim, 1:-1] = 0.5 * (dx[dim, :-1] + dx[dim, 1:])
+    out[dim, 0] = dx[dim, 0]
+    out[dim, -1] = dx[dim, -1]
+    return out
+
+
 def _make_pulses(
     neutrons: int,
     frequency: sc.Variable,
@@ -48,6 +74,8 @@ def _make_pulses(
     p_wav: sc.DataArray | None = None,
     wmin: sc.Variable | None = None,
     wmax: sc.Variable | None = None,
+    tmin: sc.Variable | None = None,
+    tmax: sc.Variable | None = None,
 ):
     """
     Create pulses from time a wavelength probability distributions.
@@ -78,9 +106,13 @@ def _make_pulses(
         Minimum neutron wavelength.
     wmax:
         Maximum neutron wavelength.
+    tmin:
+        Minimum neutron birth time.
+    tmax:
+        Maximum neutron birth time.
     """
-    t_dim = "birth_time"
-    w_dim = "wavelength"
+    # t_dim = "birth_time"
+    # w_dim = "wavelength"
 
     if p is None:
         if None in (p_time, p_wav):
@@ -103,21 +135,37 @@ def _make_pulses(
     else:
         p = p.copy(deep=False)
 
-    if p.sizes[t_dim] < 2 or p.sizes[w_dim] < 2:
+    if p.sizes[T_DIM] < 2 or p.sizes[W_DIM] < 2:
         raise ValueError(
             f"Distribution must have at least 2 points in each dimension. "
-            f"Got {p.sizes[t_dim]} in {t_dim}, {p.sizes[w_dim]} in {w_dim}"
+            f"Got {p.sizes[T_DIM]} in {T_DIM}, {p.sizes[W_DIM]} in {W_DIM}"
         )
 
-    p.coords[t_dim] = p.coords[t_dim].to(dtype=float, unit=TIME_UNIT)
-    p.coords[w_dim] = p.coords[w_dim].to(dtype=float, unit=WAV_UNIT)
+    p.coords[T_DIM] = p.coords[T_DIM].to(dtype=float, unit=TIME_UNIT)
+    p.coords[W_DIM] = p.coords[W_DIM].to(dtype=float, unit=WAV_UNIT)
 
-    tmin = p.coords[t_dim].min()
-    tmax = p.coords[t_dim].max()
-    if wmin is None:
-        wmin = p.coords[w_dim].min()
-    if wmax is None:
-        wmax = p.coords[w_dim].max()
+    # Filter parameter space defined by limits
+    ind_tmin, ind_tmax = 0, p.sizes[T_DIM]
+    ind_wmin, ind_wmax = 0, p.sizes[W_DIM]
+    trange = sc.arange(T_DIM, p.sizes[T_DIM])
+    wrange = sc.arange(W_DIM, p.sizes[W_DIM])
+    if tmin is not None:
+        ind_tmin = max(trange[p.coords[T_DIM] >= tmin][0].value - 1, 0)
+    else:
+        tmin = p.coords[T_DIM][0] - 0.5 * (p.coords[T_DIM][1] - p.coords[T_DIM][0])
+    if tmax is not None:
+        ind_tmax = min(trange[p.coords[T_DIM] <= tmax][-1].value + 2, p.sizes[T_DIM])
+    else:
+        tmax = p.coords[T_DIM][-1] + 0.5 * (p.coords[T_DIM][-1] - p.coords[T_DIM][-2])
+    if wmin is not None:
+        ind_wmin = max(wrange[p.coords[W_DIM] >= wmin][0].value - 1, 0)
+    else:
+        wmin = p.coords[W_DIM][0] - 0.5 * (p.coords[W_DIM][1] - p.coords[W_DIM][0])
+    if wmax is not None:
+        ind_wmax = min(wrange[p.coords[W_DIM] <= wmax][-1].value + 2, p.sizes[W_DIM])
+    else:
+        wmax = p.coords[W_DIM][-1] + 0.5 * (p.coords[W_DIM][-1] - p.coords[W_DIM][-2])
+    prob = p[T_DIM, ind_tmin:ind_tmax][W_DIM, ind_wmin:ind_wmax]
 
     # In the following, random.choice only allows to select from the values listed
     # in the coordinate of the probability distribution arrays. This leads to data
@@ -129,10 +177,16 @@ def _make_pulses(
     # prohibitively slow.
     # See https://docs.scipy.org/doc/scipy/tutorial/stats/sampling.html for more
     # information.
-    tsel = (p.coords[t_dim] >= tmin) & (p.coords[t_dim] <= tmax)
-    wsel = (p.coords[w_dim] >= wmin) & (p.coords[w_dim] <= wmax)
-    dt = 0.5 * (tmax - tmin).value / (p[tsel].sizes[t_dim] - 1)
-    dw = 0.5 * (wmax - wmin).value / (p[wsel].sizes[w_dim] - 1)
+
+    t = prob.coords[T_DIM]
+    w = prob.coords[W_DIM]
+    widths = {
+        T_DIM: _compute_grid_spacing(t, T_DIM),
+        W_DIM: _compute_grid_spacing(w, W_DIM),
+    }
+
+    widths[T_DIM] = widths[T_DIM].broadcast(sizes=prob.sizes).flatten(to='x')
+    widths[W_DIM] = widths[W_DIM].broadcast(sizes=prob.sizes).flatten(to='x')
 
     # Because of the added noise, some values end up being outside the specified range
     # for the birth times and wavelengths. Using naive clipping leads to pile-up on the
@@ -143,7 +197,7 @@ def _make_pulses(
     wavs = []
     ntot = pulses * neutrons
     rng = np.random.default_rng(seed)
-    p_flat = p.flatten(to='x')
+    p_flat = prob.flatten(to='x')
 
     p_sum = p_flat.data.sum()
     if p_sum.value <= 0:
@@ -155,8 +209,12 @@ def _make_pulses(
     while n < ntot:
         size = ntot - n
         inds = rng.choice(len(p_flat), size=size, p=p_flat.values)
-        t = p_flat.coords[t_dim].values[inds] + rng.normal(scale=dt, size=size)
-        w = p_flat.coords[w_dim].values[inds] + rng.normal(scale=dw, size=size)
+        t = p_flat.coords[T_DIM].values[inds] + (
+            rng.normal(scale=0.5, size=size) * widths[T_DIM].values[inds]
+        )
+        w = p_flat.coords[W_DIM].values[inds] + (
+            rng.normal(scale=0.5, size=size) * widths[W_DIM].values[inds]
+        )
         mask = (
             (t >= tmin.value)
             & (t <= tmax.value)
@@ -192,6 +250,52 @@ def _make_pulses(
     }
 
 
+def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
+    time_edges = _midpoints_to_edges(p.coords[T_DIM], T_DIM).to(
+        unit=TIME_UNIT, copy=False
+    )
+    wave_edges = _midpoints_to_edges(p.coords[W_DIM], W_DIM).to(
+        unit=WAV_UNIT, copy=False
+    )
+    frames = FrameSequence.from_source_pulse(
+        time_min=time_edges.min(),
+        time_max=time_edges.max(),
+        wavelength_min=wave_edges.min(),
+        wavelength_max=wave_edges.max(),
+    )
+    frames = frames.chop(choppers)
+    # for f in frames:
+    #     print("frame distance", f.distance)
+    # Propagate frames back to source
+    frames = FrameSequence(
+        [frame.propagate_to(p.coords['distance']) for frame in frames]
+    )
+    # for f in frames:
+    #     print("frame distance", f.distance)
+
+    X, Y = np.meshgrid(time_edges.values, wave_edges.values)
+    mask = np.zeros(shape=p.shape, dtype=bool)
+    for subf in frames[-1].subframes:
+        mask |= polygon_grid_overlap_mask(
+            np.column_stack(
+                [
+                    subf.time.to(unit=TIME_UNIT).values,
+                    subf.wavelength.to(unit=WAV_UNIT).values,
+                ]
+            ),
+            X,
+            Y,
+        )
+        # break
+    out = p.copy(deep=True)
+    out.values = np.where(mask, out.values, 0.0)
+
+    # print("Original", p)
+    # print("====================")
+    # print("NEW", out)
+    return out
+
+
 class Source:
     """
     A class that represents a source of neutrons.
@@ -216,6 +320,10 @@ class Source:
         Minimum neutron wavelength.
     wmax:
         Maximum neutron wavelength.
+    tmin:
+        Minimum neutron birth time.
+    tmax:
+        Maximum neutron birth time.
     seed:
         Seed for the random number generator.
     """  # noqa: E501
@@ -227,7 +335,10 @@ def __init__(
         pulses: int = 1,
         wmin: sc.Variable | None = None,
         wmax: sc.Variable | None = None,
+        tmin: sc.Variable | None = None,
+        tmax: sc.Variable | None = None,
         seed: int | None = None,
+        optimize_for: list[Chopper] | None = None,
     ):
         self._facility = facility.lower() if facility is not None else None
         self._neutrons = int(neutrons)
@@ -241,6 +352,11 @@ def __init__(
             file_path = get_source_path(self._facility)
             facility_pulse = sc.io.load_hdf5(file_path)
 
+            if optimize_for is not None:
+                facility_pulse = _optimize_source(
+                    p=facility_pulse, choppers=optimize_for
+                )
+
             self._frequency = facility_pulse.coords["frequency"]
             self._distance = facility_pulse.coords["distance"]
             pulse_params = _make_pulses(
@@ -250,6 +366,8 @@ def __init__(
                 pulses=self._pulses,
                 wmin=wmin,
                 wmax=wmax,
+                tmin=tmin,
+                tmax=tmax,
                 seed=seed,
             )
             self._data = sc.DataArray(

From de91f31f58bd10b5361c7b514e89ea9697b2c710 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Sun, 22 Mar 2026 23:13:45 +0100
Subject: [PATCH 03/15] cleanup

---
 src/tof/optimization.py |  1 -
 src/tof/source.py       | 10 +---------
 2 files changed, 1 insertion(+), 10 deletions(-)

diff --git a/src/tof/optimization.py b/src/tof/optimization.py
index a2651c7..0008894 100644
--- a/src/tof/optimization.py
+++ b/src/tof/optimization.py
@@ -142,7 +142,6 @@ def chop(self, chopper: Chopper) -> Frame:
         chopped = Frame(distance=frame.distance, subframes=[])
         open_times, close_times = (t.to(unit='s') for t in chopper.open_close_times())
         for subframe in frame.subframes:
-            # for open, close in zip(chopper.time_open, chopper.time_close, strict=True):
             for open, close in zip(open_times, close_times, strict=True):
                 if (tmp := _chop(subframe, open, close_to_open=True)) is not None:
                     if (tmp := _chop(tmp, close, close_to_open=False)) is not None:
diff --git a/src/tof/source.py b/src/tof/source.py
index b9a76ba..3b21ee9 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -264,14 +264,10 @@ def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
         wavelength_max=wave_edges.max(),
     )
     frames = frames.chop(choppers)
-    # for f in frames:
-    #     print("frame distance", f.distance)
     # Propagate frames back to source
     frames = FrameSequence(
         [frame.propagate_to(p.coords['distance']) for frame in frames]
     )
-    # for f in frames:
-    #     print("frame distance", f.distance)
 
     X, Y = np.meshgrid(time_edges.values, wave_edges.values)
     mask = np.zeros(shape=p.shape, dtype=bool)
@@ -286,13 +282,9 @@ def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
             X,
             Y,
         )
-        # break
+
     out = p.copy(deep=True)
     out.values = np.where(mask, out.values, 0.0)
-
-    # print("Original", p)
-    # print("====================")
-    # print("NEW", out)
     return out
 
 

From 5582e6248b907186a773b823c5cb1e49451aef0a Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Mon, 23 Mar 2026 16:16:42 +0100
Subject: [PATCH 04/15] filter out more regions from source pulse

---
 src/tof/source.py | 82 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 68 insertions(+), 14 deletions(-)

diff --git a/src/tof/source.py b/src/tof/source.py
index 3b21ee9..660390c 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -64,6 +64,25 @@ def _compute_grid_spacing(x: sc.Variable, dim: str) -> sc.Variable:
     return out
 
 
+def _load_facility_pulse_profile(facility):
+    from .facilities import get_source_path
+
+    file_path = get_source_path(facility)
+    facility_pulse = sc.io.load_hdf5(file_path)
+    facility_pulse.coords.update(
+        {
+            f"{T_DIM}_edges": _midpoints_to_edges(
+                facility_pulse.coords[T_DIM], T_DIM
+            ).to(unit=TIME_UNIT, copy=False),
+            f"{W_DIM}_edges": _midpoints_to_edges(
+                facility_pulse.coords[W_DIM], W_DIM
+            ).to(unit=WAV_UNIT, copy=False),
+        }
+    )
+    # wave_edges = _midpoints_to_edges(w, W_DIM).to(unit=WAV_UNIT, copy=False)
+    return facility_pulse
+
+
 def _make_pulses(
     neutrons: int,
     frequency: sc.Variable,
@@ -188,6 +207,9 @@ def _make_pulses(
     widths[T_DIM] = widths[T_DIM].broadcast(sizes=prob.sizes).flatten(to='x')
     widths[W_DIM] = widths[W_DIM].broadcast(sizes=prob.sizes).flatten(to='x')
 
+    time_edges = p.coords[f"{T_DIM}_edges"]
+    wave_edges = p.coords[f"{W_DIM}_edges"]
+
     # Because of the added noise, some values end up being outside the specified range
     # for the birth times and wavelengths. Using naive clipping leads to pile-up on the
     # edges of the range. To avoid this, we remove the outliers and resample until we
@@ -206,6 +228,11 @@ def _make_pulses(
             f"Sum of probabilities is {p_sum.value}"
         )
     p_flat /= p_sum
+
+    # We want to filter out events that end up in regions where the probability is zero
+    # after adding gaussian spread.
+    zero_mask = prob.data == 0.0
+
     while n < ntot:
         size = ntot - n
         inds = rng.choice(len(p_flat), size=size, p=p_flat.values)
@@ -215,15 +242,41 @@ def _make_pulses(
         w = p_flat.coords[W_DIM].values[inds] + (
             rng.normal(scale=0.5, size=size) * widths[W_DIM].values[inds]
         )
-        mask = (
+        # sel = (
+        #     (t >= tmin.value)
+        #     & (t <= tmax.value)
+        #     & (w >= wmin.value)
+        #     & (w <= wmax.value)
+        # )
+
+        # Additional selection
+        da = sc.DataArray(
+            data=sc.ones(sizes={'event': size}),
+            coords={
+                T_DIM: sc.array(dims=['event'], values=t, unit=TIME_UNIT),
+                W_DIM: sc.array(dims=['event'], values=w, unit=WAV_UNIT),
+            },
+        )
+
+        binned = da.bin({T_DIM: time_edges, W_DIM: wave_edges})
+        filtered = binned.assign_masks(m=zero_mask).bins.concat().value
+
+        t = filtered.coords[T_DIM].values
+        w = filtered.coords[W_DIM].values
+        sel = (
             (t >= tmin.value)
             & (t <= tmax.value)
             & (w >= wmin.value)
             & (w <= wmax.value)
         )
-        times.append(t[mask])
-        wavs.append(w[mask])
-        n += mask.sum()
+
+        # times.append(t[)
+        # wavs.append(filtered.coords[W_DIM].values)
+        # n += filtered.size
+
+        times.append(t[sel])
+        wavs.append(w[sel])
+        n += sel.sum()
 
     dim = "event"
     birth_time = sc.array(
@@ -251,12 +304,14 @@ def _make_pulses(
 
 
 def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
-    time_edges = _midpoints_to_edges(p.coords[T_DIM], T_DIM).to(
-        unit=TIME_UNIT, copy=False
-    )
-    wave_edges = _midpoints_to_edges(p.coords[W_DIM], W_DIM).to(
-        unit=WAV_UNIT, copy=False
-    )
+    # time_edges = _midpoints_to_edges(p.coords[T_DIM], T_DIM).to(
+    #     unit=TIME_UNIT, copy=False
+    # )
+    # wave_edges = _midpoints_to_edges(p.coords[W_DIM], W_DIM).to(
+    #     unit=WAV_UNIT, copy=False
+    # )
+    time_edges = p.coords[f"{T_DIM}_edges"]
+    wave_edges = p.coords[f"{W_DIM}_edges"]
     frames = FrameSequence.from_source_pulse(
         time_min=time_edges.min(),
         time_max=time_edges.max(),
@@ -281,6 +336,7 @@ def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
             ),
             X,
             Y,
+            tol=0,
         )
 
     out = p.copy(deep=True)
@@ -339,15 +395,13 @@ def __init__(
         self.seed = seed
 
         if self._facility is not None:
-            from .facilities import get_source_path
-
-            file_path = get_source_path(self._facility)
-            facility_pulse = sc.io.load_hdf5(file_path)
+            facility_pulse = _load_facility_pulse_profile(self._facility)
 
             if optimize_for is not None:
                 facility_pulse = _optimize_source(
                     p=facility_pulse, choppers=optimize_for
                 )
+                self.probability = facility_pulse
 
             self._frequency = facility_pulse.coords["frequency"]
             self._distance = facility_pulse.coords["distance"]

From 0c3a242201c7ccc6b76c6948002c168bdac352ab Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Mon, 23 Mar 2026 17:55:22 +0100
Subject: [PATCH 05/15] use Path to select only neutrons inside polygons

---
 src/tof/source.py | 245 ++++++++++++++++++++++++++++------------------
 1 file changed, 148 insertions(+), 97 deletions(-)

diff --git a/src/tof/source.py b/src/tof/source.py
index 660390c..c4f68e0 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -8,10 +8,11 @@
 import numpy as np
 import plopp as pp
 import scipp as sc
+from matplotlib.path import Path
 
 from .chopper import Chopper
 from .component import ComponentReading
-from .optimization import FrameSequence, polygon_grid_overlap_mask
+from .optimization import FrameSequence, Subframe, polygon_grid_overlap_mask
 from .utils import wavelength_to_speed
 
 TIME_UNIT = "us"
@@ -79,7 +80,6 @@ def _load_facility_pulse_profile(facility):
             ).to(unit=WAV_UNIT, copy=False),
         }
     )
-    # wave_edges = _midpoints_to_edges(w, W_DIM).to(unit=WAV_UNIT, copy=False)
     return facility_pulse
 
 
@@ -91,10 +91,10 @@ def _make_pulses(
     p: sc.DataArray | None = None,
     p_time: sc.DataArray | None = None,
     p_wav: sc.DataArray | None = None,
-    wmin: sc.Variable | None = None,
-    wmax: sc.Variable | None = None,
-    tmin: sc.Variable | None = None,
-    tmax: sc.Variable | None = None,
+    # wmin: sc.Variable | None = None,
+    # wmax: sc.Variable | None = None,
+    # tmin: sc.Variable | None = None,
+    # tmax: sc.Variable | None = None,
 ):
     """
     Create pulses from time a wavelength probability distributions.
@@ -133,26 +133,29 @@ def _make_pulses(
     # t_dim = "birth_time"
     # w_dim = "wavelength"
 
-    if p is None:
-        if None in (p_time, p_wav):
-            raise ValueError(
-                "Either p (2D) or both p_time (1D) and p_wav (1D) must be supplied."
-            )
-        p_wav_sum = p_wav.data.sum()
-        p_time_sum = p_time.data.sum()
-        if p_wav_sum.value <= 0:
-            raise ValueError(
-                "Wavelength distribution must have at least one positive "
-                f"probability value. Sum of probabilities is {p_wav_sum.value}"
-            )
-        if p_time_sum.value <= 0:
-            raise ValueError(
-                "Time distribution must have at least one positive "
-                f"probability value. Sum of probabilities is {p_time_sum.value}"
-            )
-        p = (p_wav / p_wav_sum) * (p_time / p_time_sum)
-    else:
-        p = p.copy(deep=False)
+    acceptance_polygons = p['acceptance_polygons']
+    p = p['probability'].copy(deep=False)
+
+    # if p is None:
+    #     if None in (p_time, p_wav):
+    #         raise ValueError(
+    #             "Either p (2D) or both p_time (1D) and p_wav (1D) must be supplied."
+    #         )
+    #     p_wav_sum = p_wav.data.sum()
+    #     p_time_sum = p_time.data.sum()
+    #     if p_wav_sum.value <= 0:
+    #         raise ValueError(
+    #             "Wavelength distribution must have at least one positive "
+    #             f"probability value. Sum of probabilities is {p_wav_sum.value}"
+    #         )
+    #     if p_time_sum.value <= 0:
+    #         raise ValueError(
+    #             "Time distribution must have at least one positive "
+    #             f"probability value. Sum of probabilities is {p_time_sum.value}"
+    #         )
+    #     p = (p_wav / p_wav_sum) * (p_time / p_time_sum)
+    # else:
+    #     p = p['probability'].copy(deep=False)
 
     if p.sizes[T_DIM] < 2 or p.sizes[W_DIM] < 2:
         raise ValueError(
@@ -163,28 +166,30 @@ def _make_pulses(
     p.coords[T_DIM] = p.coords[T_DIM].to(dtype=float, unit=TIME_UNIT)
     p.coords[W_DIM] = p.coords[W_DIM].to(dtype=float, unit=WAV_UNIT)
 
-    # Filter parameter space defined by limits
-    ind_tmin, ind_tmax = 0, p.sizes[T_DIM]
-    ind_wmin, ind_wmax = 0, p.sizes[W_DIM]
-    trange = sc.arange(T_DIM, p.sizes[T_DIM])
-    wrange = sc.arange(W_DIM, p.sizes[W_DIM])
-    if tmin is not None:
-        ind_tmin = max(trange[p.coords[T_DIM] >= tmin][0].value - 1, 0)
-    else:
-        tmin = p.coords[T_DIM][0] - 0.5 * (p.coords[T_DIM][1] - p.coords[T_DIM][0])
-    if tmax is not None:
-        ind_tmax = min(trange[p.coords[T_DIM] <= tmax][-1].value + 2, p.sizes[T_DIM])
-    else:
-        tmax = p.coords[T_DIM][-1] + 0.5 * (p.coords[T_DIM][-1] - p.coords[T_DIM][-2])
-    if wmin is not None:
-        ind_wmin = max(wrange[p.coords[W_DIM] >= wmin][0].value - 1, 0)
-    else:
-        wmin = p.coords[W_DIM][0] - 0.5 * (p.coords[W_DIM][1] - p.coords[W_DIM][0])
-    if wmax is not None:
-        ind_wmax = min(wrange[p.coords[W_DIM] <= wmax][-1].value + 2, p.sizes[W_DIM])
-    else:
-        wmax = p.coords[W_DIM][-1] + 0.5 * (p.coords[W_DIM][-1] - p.coords[W_DIM][-2])
-    prob = p[T_DIM, ind_tmin:ind_tmax][W_DIM, ind_wmin:ind_wmax]
+    # # Filter parameter space defined by limits
+    # ind_tmin, ind_tmax = 0, p.sizes[T_DIM]
+    # ind_wmin, ind_wmax = 0, p.sizes[W_DIM]
+    # trange = sc.arange(T_DIM, p.sizes[T_DIM])
+    # wrange = sc.arange(W_DIM, p.sizes[W_DIM])
+    # if tmin is not None:
+    #     ind_tmin = max(trange[p.coords[T_DIM] >= tmin][0].value - 1, 0)
+    # else:
+    #     tmin = p.coords[T_DIM][0] - 0.5 * (p.coords[T_DIM][1] - p.coords[T_DIM][0])
+    # if tmax is not None:
+    #     ind_tmax = min(trange[p.coords[T_DIM] <= tmax][-1].value + 2, p.sizes[T_DIM])
+    # else:
+    #     tmax = p.coords[T_DIM][-1] + 0.5 * (p.coords[T_DIM][-1] - p.coords[T_DIM][-2])
+    # if wmin is not None:
+    #     ind_wmin = max(wrange[p.coords[W_DIM] >= wmin][0].value - 1, 0)
+    # else:
+    #     wmin = p.coords[W_DIM][0] - 0.5 * (p.coords[W_DIM][1] - p.coords[W_DIM][0])
+    # if wmax is not None:
+    #     ind_wmax = min(wrange[p.coords[W_DIM] <= wmax][-1].value + 2, p.sizes[W_DIM])
+    # else:
+    #     wmax = p.coords[W_DIM][-1] + 0.5 * (p.coords[W_DIM][-1] - p.coords[W_DIM][-2])
+    # prob = p[T_DIM, ind_tmin:ind_tmax][W_DIM, ind_wmin:ind_wmax]
+
+    prob = p
 
     # In the following, random.choice only allows to select from the values listed
     # in the coordinate of the probability distribution arrays. This leads to data
@@ -231,7 +236,7 @@ def _make_pulses(
 
     # We want to filter out events that end up in regions where the probability is zero
     # after adding gaussian spread.
-    zero_mask = prob.data == 0.0
+    # zero_mask = prob.data == 0.0
 
     while n < ntot:
         size = ntot - n
@@ -242,6 +247,22 @@ def _make_pulses(
         w = p_flat.coords[W_DIM].values[inds] + (
             rng.normal(scale=0.5, size=size) * widths[W_DIM].values[inds]
         )
+
+        sel = np.zeros(shape=t.shape, dtype=bool)
+        for poly in acceptance_polygons:
+            # time = subf.time.to(unit='us').values
+            # wav = subf.wavelength.values
+
+            verts = np.column_stack(
+                [
+                    poly.time.to(unit=TIME_UNIT).values,
+                    poly.wavelength.to(unit=WAV_UNIT).values,
+                ]
+            )
+            path = Path(verts)
+            points = np.array([t, w]).T
+            sel |= path.contains_points(points)
+
         # sel = (
         #     (t >= tmin.value)
         #     & (t <= tmax.value)
@@ -249,26 +270,26 @@ def _make_pulses(
         #     & (w <= wmax.value)
         # )
 
-        # Additional selection
-        da = sc.DataArray(
-            data=sc.ones(sizes={'event': size}),
-            coords={
-                T_DIM: sc.array(dims=['event'], values=t, unit=TIME_UNIT),
-                W_DIM: sc.array(dims=['event'], values=w, unit=WAV_UNIT),
-            },
-        )
+        # # Additional selection
+        # da = sc.DataArray(
+        #     data=sc.ones(sizes={'event': size}),
+        #     coords={
+        #         T_DIM: sc.array(dims=['event'], values=t, unit=TIME_UNIT),
+        #         W_DIM: sc.array(dims=['event'], values=w, unit=WAV_UNIT),
+        #     },
+        # )
 
-        binned = da.bin({T_DIM: time_edges, W_DIM: wave_edges})
-        filtered = binned.assign_masks(m=zero_mask).bins.concat().value
+        # binned = da.bin({T_DIM: time_edges, W_DIM: wave_edges})
+        # filtered = binned.assign_masks(m=zero_mask).bins.concat().value
 
-        t = filtered.coords[T_DIM].values
-        w = filtered.coords[W_DIM].values
-        sel = (
-            (t >= tmin.value)
-            & (t <= tmax.value)
-            & (w >= wmin.value)
-            & (w <= wmax.value)
-        )
+        # t = filtered.coords[T_DIM].values
+        # w = filtered.coords[W_DIM].values
+        # sel = (
+        #     (t >= tmin.value)
+        #     & (t <= tmax.value)
+        #     & (w >= wmin.value)
+        #     & (w <= wmax.value)
+        # )
 
         # times.append(t[)
         # wavs.append(filtered.coords[W_DIM].values)
@@ -303,7 +324,14 @@ def _make_pulses(
     }
 
 
-def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
+def _optimize_source(
+    p,
+    wmin: sc.Variable | None = None,
+    wmax: sc.Variable | None = None,
+    tmin: sc.Variable | None = None,
+    tmax: sc.Variable | None = None,
+    choppers: list[Chopper] | None = None,
+) -> sc.DataArray:
     # time_edges = _midpoints_to_edges(p.coords[T_DIM], T_DIM).to(
     #     unit=TIME_UNIT, copy=False
     # )
@@ -312,36 +340,54 @@ def _optimize_source(p, choppers: list[Chopper]) -> sc.DataArray:
     # )
     time_edges = p.coords[f"{T_DIM}_edges"]
     wave_edges = p.coords[f"{W_DIM}_edges"]
-    frames = FrameSequence.from_source_pulse(
-        time_min=time_edges.min(),
-        time_max=time_edges.max(),
-        wavelength_min=wave_edges.min(),
-        wavelength_max=wave_edges.max(),
-    )
-    frames = frames.chop(choppers)
-    # Propagate frames back to source
-    frames = FrameSequence(
-        [frame.propagate_to(p.coords['distance']) for frame in frames]
-    )
+
+    if choppers is not None:
+        frames = FrameSequence.from_source_pulse(
+            time_min=time_edges.min(),
+            time_max=time_edges.max(),
+            wavelength_min=wave_edges.min(),
+            wavelength_max=wave_edges.max(),
+        )
+        frames = frames.chop(choppers)
+        # Propagate frames back to source
+        frames = FrameSequence(
+            [frame.propagate_to(p.coords['distance']) for frame in frames]
+        )
+        polygons = frames[-1].subframes
+    else:
+        if tmin is None:
+            tmin = time_edges[0]
+        if tmax is None:
+            tmax = time_edges[-1]
+        if wmin is None:
+            wmin = wave_edges[0]
+        if wmax is None:
+            wmax = wave_edges[-1]
+        polygons = [
+            Subframe(
+                time=sc.concat([tmin, tmax, tmax, tmin, tmin], dim='vertex'),
+                wavelength=sc.concat([wmin, wmin, wmax, wmax, wmin], dim='vertex'),
+            )
+        ]
 
     X, Y = np.meshgrid(time_edges.values, wave_edges.values)
     mask = np.zeros(shape=p.shape, dtype=bool)
-    for subf in frames[-1].subframes:
+    for poly in polygons:
         mask |= polygon_grid_overlap_mask(
             np.column_stack(
                 [
-                    subf.time.to(unit=TIME_UNIT).values,
-                    subf.wavelength.to(unit=WAV_UNIT).values,
+                    poly.time.to(unit=TIME_UNIT).values,
+                    poly.wavelength.to(unit=WAV_UNIT).values,
                 ]
             ),
             X,
             Y,
-            tol=0,
         )
 
-    out = p.copy(deep=True)
-    out.values = np.where(mask, out.values, 0.0)
-    return out
+    da = p.copy(deep=True)
+    da.values = np.where(mask, da.values, 0.0)
+    return sc.DataGroup({"probability": da, "acceptance_polygons": polygons})
+    # return out
 
 
 class Source:
@@ -397,23 +443,28 @@ def __init__(
         if self._facility is not None:
             facility_pulse = _load_facility_pulse_profile(self._facility)
 
-            if optimize_for is not None:
-                facility_pulse = _optimize_source(
-                    p=facility_pulse, choppers=optimize_for
-                )
-                self.probability = facility_pulse
+            # if optimize_for is not None:
+            facility_pulse = _optimize_source(
+                p=facility_pulse,
+                wmin=wmin,
+                wmax=wmax,
+                tmin=tmin,
+                tmax=tmax,
+                choppers=optimize_for,
+            )
+            self.probability = facility_pulse["probability"]
 
-            self._frequency = facility_pulse.coords["frequency"]
-            self._distance = facility_pulse.coords["distance"]
+            self._frequency = facility_pulse["probability"].coords["frequency"]
+            self._distance = facility_pulse["probability"].coords["distance"]
             pulse_params = _make_pulses(
                 neutrons=self._neutrons,
                 p=facility_pulse,
                 frequency=self._frequency,
                 pulses=self._pulses,
-                wmin=wmin,
-                wmax=wmax,
-                tmin=tmin,
-                tmax=tmax,
+                # wmin=wmin,
+                # wmax=wmax,
+                # tmin=tmin,
+                # tmax=tmax,
                 seed=seed,
             )
             self._data = sc.DataArray(

From 4538f54bc105d62449a4f23624ac202353ea1bde Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Mon, 23 Mar 2026 22:32:56 +0100
Subject: [PATCH 06/15] refactor and cleanup

---
 src/tof/source.py | 463 +++++++++++++++++++++++-----------------------
 1 file changed, 230 insertions(+), 233 deletions(-)

diff --git a/src/tof/source.py b/src/tof/source.py
index c4f68e0..dde3ef2 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -21,6 +21,37 @@
 W_DIM = "wavelength"
 
 
+@dataclass(frozen=True)
+class SourceReading(ComponentReading):
+    """
+    Read-only container for the parameters of a source.
+    """
+
+    data: sc.DataArray
+    facility: str | None
+    neutrons: int
+    frequency: sc.Variable
+    pulses: int
+    distance: sc.Variable
+
+    @property
+    def kind(self) -> str:
+        return "source"
+
+    def plot_on_time_distance_diagram(self, ax, pulse) -> None:
+        birth_time = self.data.coords["birth_time"]["pulse", pulse]
+        tmin = birth_time.min().value
+        dist = self.distance.value
+        ax.plot([tmin, birth_time.max().value], [dist] * 2, color="gray", lw=3)
+        ax.text(tmin, dist, "Pulse", ha="left", va="top", color="gray")
+
+
+@dataclass
+class SourceDistribution:
+    probability: sc.DataArray
+    acceptance_polygons: list[Subframe]
+
+
 def _default_frequency(frequency: sc.Variable | None, pulses: int) -> sc.Variable:
     if frequency is None:
         if pulses > 1:
@@ -89,21 +120,11 @@ def _make_pulses(
     pulses: int,
     seed: int | None,
     p: sc.DataArray | None = None,
-    p_time: sc.DataArray | None = None,
-    p_wav: sc.DataArray | None = None,
-    # wmin: sc.Variable | None = None,
-    # wmax: sc.Variable | None = None,
-    # tmin: sc.Variable | None = None,
-    # tmax: sc.Variable | None = None,
 ):
     """
-    Create pulses from time a wavelength probability distributions.
-    The distributions should be supplied as DataArrays where the coordinates
+    Create pulses from a 2D time and wavelength probability distribution.
+    The distribution ``p`` should be supplied as a DataArray where the coordinates
     are the values of the distribution, and the values are the probability.
-    Note that the time and wavelength distributions are independent. A neutron with
-    a randomly selected birth time from ``p_time`` can adopt any wavelength in
-    ``p_wav`` (in other words, the two distributions are simply broadcast into a
-    square 2D parameter space).
 
     Parameters
     ----------
@@ -117,45 +138,10 @@ def _make_pulses(
         Seed for the random number generator.
     p:
         2D probability distribution for a single pulse.
-    p_time:
-        Time probability distribution for a single pulse.
-    p_wav:
-        Wavelength probability distribution for a single pulse.
-    wmin:
-        Minimum neutron wavelength.
-    wmax:
-        Maximum neutron wavelength.
-    tmin:
-        Minimum neutron birth time.
-    tmax:
-        Maximum neutron birth time.
     """
-    # t_dim = "birth_time"
-    # w_dim = "wavelength"
-
-    acceptance_polygons = p['acceptance_polygons']
-    p = p['probability'].copy(deep=False)
-
-    # if p is None:
-    #     if None in (p_time, p_wav):
-    #         raise ValueError(
-    #             "Either p (2D) or both p_time (1D) and p_wav (1D) must be supplied."
-    #         )
-    #     p_wav_sum = p_wav.data.sum()
-    #     p_time_sum = p_time.data.sum()
-    #     if p_wav_sum.value <= 0:
-    #         raise ValueError(
-    #             "Wavelength distribution must have at least one positive "
-    #             f"probability value. Sum of probabilities is {p_wav_sum.value}"
-    #         )
-    #     if p_time_sum.value <= 0:
-    #         raise ValueError(
-    #             "Time distribution must have at least one positive "
-    #             f"probability value. Sum of probabilities is {p_time_sum.value}"
-    #         )
-    #     p = (p_wav / p_wav_sum) * (p_time / p_time_sum)
-    # else:
-    #     p = p['probability'].copy(deep=False)
+
+    acceptance_polygons = p.acceptance_polygons
+    p = p.probability.copy(deep=False)
 
     if p.sizes[T_DIM] < 2 or p.sizes[W_DIM] < 2:
         raise ValueError(
@@ -166,31 +152,6 @@ def _make_pulses(
     p.coords[T_DIM] = p.coords[T_DIM].to(dtype=float, unit=TIME_UNIT)
     p.coords[W_DIM] = p.coords[W_DIM].to(dtype=float, unit=WAV_UNIT)
 
-    # # Filter parameter space defined by limits
-    # ind_tmin, ind_tmax = 0, p.sizes[T_DIM]
-    # ind_wmin, ind_wmax = 0, p.sizes[W_DIM]
-    # trange = sc.arange(T_DIM, p.sizes[T_DIM])
-    # wrange = sc.arange(W_DIM, p.sizes[W_DIM])
-    # if tmin is not None:
-    #     ind_tmin = max(trange[p.coords[T_DIM] >= tmin][0].value - 1, 0)
-    # else:
-    #     tmin = p.coords[T_DIM][0] - 0.5 * (p.coords[T_DIM][1] - p.coords[T_DIM][0])
-    # if tmax is not None:
-    #     ind_tmax = min(trange[p.coords[T_DIM] <= tmax][-1].value + 2, p.sizes[T_DIM])
-    # else:
-    #     tmax = p.coords[T_DIM][-1] + 0.5 * (p.coords[T_DIM][-1] - p.coords[T_DIM][-2])
-    # if wmin is not None:
-    #     ind_wmin = max(wrange[p.coords[W_DIM] >= wmin][0].value - 1, 0)
-    # else:
-    #     wmin = p.coords[W_DIM][0] - 0.5 * (p.coords[W_DIM][1] - p.coords[W_DIM][0])
-    # if wmax is not None:
-    #     ind_wmax = min(wrange[p.coords[W_DIM] <= wmax][-1].value + 2, p.sizes[W_DIM])
-    # else:
-    #     wmax = p.coords[W_DIM][-1] + 0.5 * (p.coords[W_DIM][-1] - p.coords[W_DIM][-2])
-    # prob = p[T_DIM, ind_tmin:ind_tmax][W_DIM, ind_wmin:ind_wmax]
-
-    prob = p
-
     # In the following, random.choice only allows to select from the values listed
     # in the coordinate of the probability distribution arrays. This leads to data
     # grouped into spikes and empty in between because the sampling resolution used
@@ -202,19 +163,13 @@ def _make_pulses(
     # See https://docs.scipy.org/doc/scipy/tutorial/stats/sampling.html for more
     # information.
 
-    t = prob.coords[T_DIM]
-    w = prob.coords[W_DIM]
     widths = {
-        T_DIM: _compute_grid_spacing(t, T_DIM),
-        W_DIM: _compute_grid_spacing(w, W_DIM),
+        dim: _compute_grid_spacing(p.coords[dim], dim)
+        .broadcast(sizes=p.sizes)
+        .flatten(to='x')
+        for dim in (T_DIM, W_DIM)
     }
 
-    widths[T_DIM] = widths[T_DIM].broadcast(sizes=prob.sizes).flatten(to='x')
-    widths[W_DIM] = widths[W_DIM].broadcast(sizes=prob.sizes).flatten(to='x')
-
-    time_edges = p.coords[f"{T_DIM}_edges"]
-    wave_edges = p.coords[f"{W_DIM}_edges"]
-
     # Because of the added noise, some values end up being outside the specified range
     # for the birth times and wavelengths. Using naive clipping leads to pile-up on the
     # edges of the range. To avoid this, we remove the outliers and resample until we
@@ -224,7 +179,7 @@ def _make_pulses(
     wavs = []
     ntot = pulses * neutrons
     rng = np.random.default_rng(seed)
-    p_flat = prob.flatten(to='x')
+    p_flat = p.flatten(to='x')
 
     p_sum = p_flat.data.sum()
     if p_sum.value <= 0:
@@ -234,12 +189,15 @@ def _make_pulses(
         )
     p_flat /= p_sum
 
-    # We want to filter out events that end up in regions where the probability is zero
-    # after adding gaussian spread.
-    # zero_mask = prob.data == 0.0
-
     while n < ntot:
-        size = ntot - n
+        # The first iteration, we sample all neutrons. Because some get discarded, we
+        # will have to iterate at least twice. Instead of just sampling the missing
+        # number of neutrons, subsequent iterations over-sample (here we choose 50% of
+        # the total number requested) and then trim the excess.
+        # This leads to less iterations overall, as we can have many iterations at the
+        # end when we are missing just one or two neutrons and we keep sampling them
+        # outside of the accepted regions.
+        size = max(ntot - n, ntot // 2)
         inds = rng.choice(len(p_flat), size=size, p=p_flat.values)
         t = p_flat.coords[T_DIM].values[inds] + (
             rng.normal(scale=0.5, size=size) * widths[T_DIM].values[inds]
@@ -248,56 +206,29 @@ def _make_pulses(
             rng.normal(scale=0.5, size=size) * widths[W_DIM].values[inds]
         )
 
+        # We filter events that are outside the accepted regions. Accepted regions can
+        # be defined by either a cropping of the probability distribution (using tmin,
+        # tmax, wmin, wmax) or generated by a chopper acceptance diagram for sources
+        # that are optimized for a specific chopper cascade.
+        points = np.array([t, w]).T
         sel = np.zeros(shape=t.shape, dtype=bool)
         for poly in acceptance_polygons:
-            # time = subf.time.to(unit='us').values
-            # wav = subf.wavelength.values
-
             verts = np.column_stack(
                 [
                     poly.time.to(unit=TIME_UNIT).values,
                     poly.wavelength.to(unit=WAV_UNIT).values,
                 ]
             )
+            # Use Matplotlib's efficient Path.contains_points
             path = Path(verts)
-            points = np.array([t, w]).T
             sel |= path.contains_points(points)
 
-        # sel = (
-        #     (t >= tmin.value)
-        #     & (t <= tmax.value)
-        #     & (w >= wmin.value)
-        #     & (w <= wmax.value)
-        # )
-
-        # # Additional selection
-        # da = sc.DataArray(
-        #     data=sc.ones(sizes={'event': size}),
-        #     coords={
-        #         T_DIM: sc.array(dims=['event'], values=t, unit=TIME_UNIT),
-        #         W_DIM: sc.array(dims=['event'], values=w, unit=WAV_UNIT),
-        #     },
-        # )
-
-        # binned = da.bin({T_DIM: time_edges, W_DIM: wave_edges})
-        # filtered = binned.assign_masks(m=zero_mask).bins.concat().value
-
-        # t = filtered.coords[T_DIM].values
-        # w = filtered.coords[W_DIM].values
-        # sel = (
-        #     (t >= tmin.value)
-        #     & (t <= tmax.value)
-        #     & (w >= wmin.value)
-        #     & (w <= wmax.value)
-        # )
-
-        # times.append(t[)
-        # wavs.append(filtered.coords[W_DIM].values)
-        # n += filtered.size
-
-        times.append(t[sel])
-        wavs.append(w[sel])
-        n += sel.sum()
+        # We could end up with more neutrons than needed, so only select the number
+        # we need.
+        inds = np.arange(size)[sel][: ntot - n]
+        times.append(t[inds])
+        wavs.append(w[inds])
+        n += len(inds)
 
     dim = "event"
     birth_time = sc.array(
@@ -331,13 +262,16 @@ def _optimize_source(
     tmin: sc.Variable | None = None,
     tmax: sc.Variable | None = None,
     choppers: list[Chopper] | None = None,
-) -> sc.DataArray:
-    # time_edges = _midpoints_to_edges(p.coords[T_DIM], T_DIM).to(
-    #     unit=TIME_UNIT, copy=False
-    # )
-    # wave_edges = _midpoints_to_edges(p.coords[W_DIM], W_DIM).to(
-    #     unit=WAV_UNIT, copy=False
-    # )
+) -> SourceDistribution:
+    """
+    Optimize the source distribution by applying acceptance criteria.
+    If choppers are provided, they will be used to define the acceptance diagram
+    computed analytically.
+    Otherwise, a single acceptance polygon will be created based on the provided
+    min/max values for time and wavelength.
+
+    Probabilities outside the acceptance regions will be set to zero.
+    """
     time_edges = p.coords[f"{T_DIM}_edges"]
     wave_edges = p.coords[f"{W_DIM}_edges"]
 
@@ -384,10 +318,97 @@ def _optimize_source(
             Y,
         )
 
-    da = p.copy(deep=True)
-    da.values = np.where(mask, da.values, 0.0)
-    return sc.DataGroup({"probability": da, "acceptance_polygons": polygons})
-    # return out
+    prob = p.copy(deep=True)
+    prob.values = np.where(mask, prob.values, 0.0)
+    return SourceDistribution(probability=prob, acceptance_polygons=polygons)
+
+
+def _sample_source_data_from_distribution(
+    neutrons: int,
+    frequency: sc.Variable,
+    p: sc.Variable | None,
+    p_time: sc.Variable | None,
+    p_wav: sc.Variable | None,
+    distance: sc.Variable | None,
+    wmin: sc.Variable | None,
+    wmax: sc.Variable | None,
+    tmin: sc.Variable | None,
+    tmax: sc.Variable | None,
+    seed: int | None,
+    optimize_for: list[Chopper] | None,
+    pulses: int,
+) -> sc.DataArray:
+    """
+    Sample neutrons from the source distribution.
+    """
+    if p is not None:
+        if (p_time is not None) or (p_wav is not None):
+            raise ValueError(
+                "If p (2D) is supplied, both p_time (1D) and p_wav (1D) must be None."
+            )
+        p = _bin_edges_to_midpoints(p, dims=["birth_time", "wavelength"])
+    else:
+        if None in (p_time, p_wav):
+            raise ValueError(
+                "Either p (2D) or both p_time (1D) and p_wav (1D) must be supplied."
+            )
+        p_time = _bin_edges_to_midpoints(p_time, dims=["birth_time"])
+        p_wav = _bin_edges_to_midpoints(p_wav, dims=["wavelength"])
+
+        p_wav_sum = p_wav.data.sum()
+        p_time_sum = p_time.data.sum()
+        if p_wav_sum.value <= 0:
+            raise ValueError(
+                "Wavelength distribution must have at least one positive "
+                f"probability value. Sum of probabilities is {p_wav_sum.value}"
+            )
+        if p_time_sum.value <= 0:
+            raise ValueError(
+                "Time distribution must have at least one positive "
+                f"probability value. Sum of probabilities is {p_time_sum.value}"
+            )
+        p = (p_wav / p_wav_sum) * (p_time / p_time_sum)
+
+    p.coords.update(
+        {
+            f"{T_DIM}_edges": _midpoints_to_edges(p.coords[T_DIM], T_DIM).to(
+                unit=TIME_UNIT, copy=False
+            ),
+            f"{W_DIM}_edges": _midpoints_to_edges(p.coords[W_DIM], W_DIM).to(
+                unit=WAV_UNIT, copy=False
+            ),
+        }
+    )
+    if 'distance' not in p.coords:
+        p.coords['distance'] = distance
+
+    distribution = _optimize_source(
+        p=p,
+        wmin=wmin,
+        wmax=wmax,
+        tmin=tmin,
+        tmax=tmax,
+        choppers=optimize_for,
+    )
+
+    pulse_params = _make_pulses(
+        neutrons=neutrons,
+        p=distribution,
+        frequency=frequency,
+        pulses=pulses,
+        seed=seed,
+    )
+    return sc.DataArray(
+        data=sc.ones(sizes=pulse_params["birth_time"].sizes, unit="counts"),
+        coords={
+            "birth_time": pulse_params["birth_time"],
+            "wavelength": pulse_params["wavelength"],
+            "speed": pulse_params["speed"],
+            "id": sc.arange("event", pulse_params["birth_time"].size, unit=None).fold(
+                "event", sizes=pulse_params["birth_time"].sizes
+            ),
+        },
+    )
 
 
 class Source:
@@ -420,6 +441,12 @@ class Source:
         Maximum neutron birth time.
     seed:
         Seed for the random number generator.
+    optimize_for:
+        List of choppers to optimize the source for. A chopper acceptance diagram will
+        be overlaid on the source distribution, and samples will be taken only from
+        the regions where the chopper cascade is accepting neutrons.
+
+        .. versionadded:: 26.4.0
     """  # noqa: E501
 
     def __init__(
@@ -440,44 +467,28 @@ def __init__(
         self._data = None
         self.seed = seed
 
-        if self._facility is not None:
-            facility_pulse = _load_facility_pulse_profile(self._facility)
-
-            # if optimize_for is not None:
-            facility_pulse = _optimize_source(
-                p=facility_pulse,
-                wmin=wmin,
-                wmax=wmax,
-                tmin=tmin,
-                tmax=tmax,
-                choppers=optimize_for,
-            )
-            self.probability = facility_pulse["probability"]
-
-            self._frequency = facility_pulse["probability"].coords["frequency"]
-            self._distance = facility_pulse["probability"].coords["distance"]
-            pulse_params = _make_pulses(
-                neutrons=self._neutrons,
-                p=facility_pulse,
-                frequency=self._frequency,
-                pulses=self._pulses,
-                # wmin=wmin,
-                # wmax=wmax,
-                # tmin=tmin,
-                # tmax=tmax,
-                seed=seed,
-            )
-            self._data = sc.DataArray(
-                data=sc.ones(sizes=pulse_params["birth_time"].sizes, unit="counts"),
-                coords={
-                    "birth_time": pulse_params["birth_time"],
-                    "wavelength": pulse_params["wavelength"],
-                    "speed": pulse_params["speed"],
-                    "id": sc.arange(
-                        "event", pulse_params["birth_time"].size, unit=None
-                    ).fold("event", sizes=pulse_params["birth_time"].sizes),
-                },
-            )
+        if self._facility is None:
+            return
+
+        facility_pulse = _load_facility_pulse_profile(self._facility)
+
+        self._distance = facility_pulse.coords['distance']
+        self._frequency = facility_pulse.coords['frequency']
+        self._data = _sample_source_data_from_distribution(
+            p=facility_pulse,
+            neutrons=neutrons,
+            pulses=pulses,
+            frequency=self._frequency,
+            seed=seed,
+            distance=self._distance,
+            wmin=wmin,
+            wmax=wmax,
+            tmin=tmin,
+            tmax=tmax,
+            optimize_for=optimize_for,
+            p_time=None,
+            p_wav=None,
+        )
 
     @property
     def facility(self) -> str | None:
@@ -599,15 +610,22 @@ def from_distribution(
         frequency: sc.Variable | None = None,
         seed: int | None = None,
         distance: sc.Variable | None = None,
+        wmin: sc.Variable | None = None,
+        wmax: sc.Variable | None = None,
+        tmin: sc.Variable | None = None,
+        tmax: sc.Variable | None = None,
+        optimize_for: list[Chopper] | None = None,
     ):
         """
-        Create source pulses from time a wavelength probability distributions.
+        Create source pulses from time and wavelength probability distributions.
         The distributions should be supplied as DataArrays where the coordinates
         are the values of the distribution, and the values are the probability.
-        Note that the time and wavelength distributions are independent. A neutron with
-        a randomly selected birth time from ``p_time`` can adopt any wavelength in
-        ``p_wav`` (in other words, the two distributions are simply broadcast into a
-        square 2D parameter space).
+        This can either be a single 2D distribution, or separate 1D distributions
+        for time and wavelength. Note that in this case, the time and wavelength
+        distributions are considered independent. A neutron with a randomly selected
+        birth time from ``p_time`` can adopt any wavelength in ``p_wav`` (in other
+        words, the two distributions are simply broadcast into a square 2D parameter
+        space).
 
         Parameters
         ----------
@@ -629,40 +647,44 @@ def from_distribution(
             Seed for the random number generator.
         distance:
             Position of the source along the beamline.
+        wmin:
+            Minimum wavelength.
+        wmax:
+            Maximum wavelength.
+        tmin:
+            Minimum birth time.
+        tmax:
+            Maximum birth time.
+        optimize_for:
+            List of choppers to optimize the source for. A chopper acceptance diagram
+            will be overlaid on the source distribution, and samples will be taken only
+            from the regions where the chopper cascade is accepting neutrons.
+
+            .. versionadded:: 26.4.0
         """
 
         source = cls(facility=None, neutrons=neutrons, pulses=pulses)
-        source._distance = (
-            distance if distance is not None else sc.scalar(0.0, unit="m")
-        )
-        source._frequency = _default_frequency(frequency, pulses)
 
-        if p is not None:
-            p = _bin_edges_to_midpoints(p, dims=["birth_time", "wavelength"])
-        if p_time is not None:
-            p_time = _bin_edges_to_midpoints(p_time, dims=["birth_time"])
-        if p_wav is not None:
-            p_wav = _bin_edges_to_midpoints(p_wav, dims=["wavelength"])
+        if distance is None:
+            distance = sc.scalar(0.0, unit="m")
+        frequency = _default_frequency(frequency, pulses)
 
-        pulse_params = _make_pulses(
-            neutrons=neutrons,
+        source._distance = distance
+        source._frequency = frequency
+        source._data = _sample_source_data_from_distribution(
             p=p,
             p_time=p_time,
             p_wav=p_wav,
-            frequency=source._frequency,
+            neutrons=neutrons,
             pulses=pulses,
+            frequency=frequency,
             seed=seed,
-        )
-        source._data = sc.DataArray(
-            data=sc.ones(sizes=pulse_params["birth_time"].sizes, unit="counts"),
-            coords={
-                "birth_time": pulse_params["birth_time"],
-                "wavelength": pulse_params["wavelength"],
-                "speed": pulse_params["speed"],
-                "id": sc.arange(
-                    "event", pulse_params["birth_time"].size, unit=None
-                ).fold("event", sizes=pulse_params["birth_time"].sizes),
-            },
+            distance=distance,
+            wmin=wmin,
+            wmax=wmax,
+            tmin=tmin,
+            tmax=tmax,
+            optimize_for=optimize_for,
         )
         return source
 
@@ -746,28 +768,3 @@ def as_json(self) -> dict:
             'seed': self.seed,
             'type': 'source',
         }
-
-
-@dataclass(frozen=True)
-class SourceReading(ComponentReading):
-    """
-    Read-only container for the parameters of a source.
-    """
-
-    data: sc.DataArray
-    facility: str | None
-    neutrons: int
-    frequency: sc.Variable
-    pulses: int
-    distance: sc.Variable
-
-    @property
-    def kind(self) -> str:
-        return "source"
-
-    def plot_on_time_distance_diagram(self, ax, pulse) -> None:
-        birth_time = self.data.coords["birth_time"]["pulse", pulse]
-        tmin = birth_time.min().value
-        dist = self.distance.value
-        ax.plot([tmin, birth_time.max().value], [dist] * 2, color="gray", lw=3)
-        ax.text(tmin, dist, "Pulse", ha="left", va="top", color="gray")

From 28bb5e6823f3b4a253a66e5a0bc34c16cd2c195e Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Mon, 23 Mar 2026 22:55:29 +0100
Subject: [PATCH 07/15] fix args passing

---
 src/tof/source.py | 38 ++++++++++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/tof/source.py b/src/tof/source.py
index dde3ef2..f91d609 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -454,6 +454,7 @@ def __init__(
         facility: str | None,
         neutrons: int = 1_000_000,
         pulses: int = 1,
+        frequency: sc.Variable | None = None,
         wmin: sc.Variable | None = None,
         wmax: sc.Variable | None = None,
         tmin: sc.Variable | None = None,
@@ -465,21 +466,24 @@ def __init__(
         self._neutrons = int(neutrons)
         self._pulses = int(pulses)
         self._data = None
-        self.seed = seed
+        self._seed = seed
 
         if self._facility is None:
             return
 
         facility_pulse = _load_facility_pulse_profile(self._facility)
 
+        if frequency is not None:
+            facility_pulse = facility_pulse.assign_coords(frequency=frequency)
+
         self._distance = facility_pulse.coords['distance']
         self._frequency = facility_pulse.coords['frequency']
         self._data = _sample_source_data_from_distribution(
             p=facility_pulse,
-            neutrons=neutrons,
-            pulses=pulses,
+            neutrons=self._neutrons,
+            pulses=self._pulses,
             frequency=self._frequency,
-            seed=seed,
+            seed=self._seed,
             distance=self._distance,
             wmin=wmin,
             wmax=wmax,
@@ -511,6 +515,13 @@ def frequency(self) -> sc.Variable:
         """
         return self._frequency
 
+    @property
+    def period(self) -> sc.Variable:
+        """
+        The period of the pulse.
+        """
+        return 1.0 / self._frequency
+
     @property
     def distance(self) -> sc.Variable:
         """
@@ -525,6 +536,13 @@ def pulses(self) -> int:
         """
         return self._pulses
 
+    @property
+    def seed(self) -> int | None:
+        """
+        The seed for the random number generator.
+        """
+        return self._seed
+
     @property
     def data(self) -> sc.DataArray:
         """
@@ -663,7 +681,7 @@ def from_distribution(
             .. versionadded:: 26.4.0
         """
 
-        source = cls(facility=None, neutrons=neutrons, pulses=pulses)
+        source = cls(facility=None, neutrons=neutrons, pulses=pulses, seed=seed)
 
         if distance is None:
             distance = sc.scalar(0.0, unit="m")
@@ -675,11 +693,11 @@ def from_distribution(
             p=p,
             p_time=p_time,
             p_wav=p_wav,
-            neutrons=neutrons,
-            pulses=pulses,
-            frequency=frequency,
-            seed=seed,
-            distance=distance,
+            neutrons=source._neutrons,
+            pulses=source._pulses,
+            frequency=source._frequency,
+            seed=source._seed,
+            distance=source._distance,
             wmin=wmin,
             wmax=wmax,
             tmin=tmin,

From 4375a8d0da8fba47b922192023da8ddb2dcf1bbb Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Mon, 23 Mar 2026 23:07:51 +0100
Subject: [PATCH 08/15] fix broken test

---
 tests/source_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/source_test.py b/tests/source_test.py
index 10b3b70..5d5db47 100644
--- a/tests/source_test.py
+++ b/tests/source_test.py
@@ -148,7 +148,7 @@ def test_multiple_pulses_from_distribution():
     p_time = sc.DataArray(
         data=sc.array(dims=['birth_time'], values=v),
         coords={
-            'birth_time': sc.linspace('birth_time', 0.0, 8000.0, len(v), unit='us')
+            'birth_time': sc.linspace('birth_time', 1.0e3, 6.0e3, len(v), unit='us')
         },
     )
     p_wav = sc.DataArray(

From 2317fe572faa79ffb88a091e201259042dbf0cd9 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Tue, 24 Mar 2026 14:41:35 +0100
Subject: [PATCH 09/15] use patches instead to mask

---
 src/tof/source.py | 51 ++++++++++++++++++++++++++++++++---------------
 1 file changed, 35 insertions(+), 16 deletions(-)

diff --git a/src/tof/source.py b/src/tof/source.py
index f91d609..d62a60b 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 import numpy as np
@@ -189,7 +190,10 @@ def _make_pulses(
         )
     p_flat /= p_sum
 
+    it = 0
+
     while n < ntot:
+        it += 1
         # The first iteration, we sample all neutrons. Because some get discarded, we
         # will have to iterate at least twice. Instead of just sampling the missing
         # number of neutrons, subsequent iterations over-sample (here we choose 50% of
@@ -197,7 +201,8 @@ def _make_pulses(
         # This leads to less iterations overall, as we can have many iterations at the
         # end when we are missing just one or two neutrons and we keep sampling them
         # outside of the accepted regions.
-        size = max(ntot - n, ntot // 2)
+        # size = max(ntot - n, ntot // 2)
+        size = ntot
         inds = rng.choice(len(p_flat), size=size, p=p_flat.values)
         t = p_flat.coords[T_DIM].values[inds] + (
             rng.normal(scale=0.5, size=size) * widths[T_DIM].values[inds]
@@ -229,6 +234,9 @@ def _make_pulses(
         times.append(t[inds])
         wavs.append(w[inds])
         n += len(inds)
+        print("Selected", len(inds), "neutrons", ntot - n, "remaining")
+
+    print("Sampling took", it, "iterations")
 
     dim = "event"
     birth_time = sc.array(
@@ -261,7 +269,7 @@ def _optimize_source(
     wmax: sc.Variable | None = None,
     tmin: sc.Variable | None = None,
     tmax: sc.Variable | None = None,
-    choppers: list[Chopper] | None = None,
+    choppers: list[Chopper] | Mapping[str, Chopper] | None = None,
 ) -> SourceDistribution:
     """
     Optimize the source distribution by applying acceptance criteria.
@@ -282,7 +290,9 @@ def _optimize_source(
             wavelength_min=wave_edges.min(),
             wavelength_max=wave_edges.max(),
         )
-        frames = frames.chop(choppers)
+        frames = frames.chop(
+            choppers.values() if hasattr(choppers, "items") else choppers
+        )
         # Propagate frames back to source
         frames = FrameSequence(
             [frame.propagate_to(p.coords['distance']) for frame in frames]
@@ -304,22 +314,30 @@ def _optimize_source(
             )
         ]
 
-    X, Y = np.meshgrid(time_edges.values, wave_edges.values)
-    mask = np.zeros(shape=p.shape, dtype=bool)
+    # X, Y = np.meshgrid(time_edges.values, wave_edges.values)
+    mask = sc.zeros(sizes=p.sizes, dtype=bool)
     for poly in polygons:
-        mask |= polygon_grid_overlap_mask(
-            np.column_stack(
-                [
-                    poly.time.to(unit=TIME_UNIT).values,
-                    poly.wavelength.to(unit=WAV_UNIT).values,
-                ]
-            ),
-            X,
-            Y,
+        # Mask with bounding box of each polygon
+        mask |= (
+            (time_edges[1:] >= poly.time.min().to(unit=TIME_UNIT))
+            & (time_edges[:-1] <= poly.time.max().to(unit=TIME_UNIT))
+            & (wave_edges[1:] >= poly.wavelength.min().to(unit=WAV_UNIT))
+            & (wave_edges[:-1] <= poly.wavelength.max().to(unit=WAV_UNIT))
         )
 
-    prob = p.copy(deep=True)
-    prob.values = np.where(mask, prob.values, 0.0)
+    #     mask |= polygon_grid_overlap_mask(
+    #         np.column_stack(
+    #             [
+    #                 poly.time.to(unit=TIME_UNIT).values,
+    #                 poly.wavelength.to(unit=WAV_UNIT).values,
+    #             ]
+    #         ),
+    #         X,
+    #         Y,
+    #     )
+
+    # prob = p.copy(deep=True)
+    prob = sc.where(mask, p, sc.scalar(0.0, unit=p.unit))
     return SourceDistribution(probability=prob, acceptance_polygons=polygons)
 
 
@@ -554,6 +572,7 @@ def data(self) -> sc.DataArray:
                 "eto": self._data.coords["birth_time"]
                 % (1.0 / self._frequency).to(unit=TIME_UNIT, copy=False),
                 "toa": self._data.coords["birth_time"],
+                "birth_wavelength": self._data.coords["wavelength"],
             }
         )
 

From 0b737fd83dc5aa9184ebefff98b7c13f6e31243c Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Tue, 24 Mar 2026 17:26:34 +0100
Subject: [PATCH 10/15] remove masking P using polygons

---
 src/tof/optimization.py | 162 ----------------------------------------
 src/tof/source.py       |  21 ++----
 2 files changed, 5 insertions(+), 178 deletions(-)

diff --git a/src/tof/optimization.py b/src/tof/optimization.py
index 0008894..d0742db 100644
--- a/src/tof/optimization.py
+++ b/src/tof/optimization.py
@@ -267,165 +267,3 @@ def _chop(frame: Subframe, time: sc.Variable, close_to_open: bool) -> Subframe |
     time = sc.concat([t for t, _ in output], dim=frame.time.dim)
     wavelength = sc.concat([v for _, v in output], dim=frame.wavelength.dim)
     return Subframe(time=time, wavelength=wavelength)
-
-
-def polygon_grid_overlap_mask(
-    polygon: np.ndarray, X: np.ndarray, Y: np.ndarray, tol: float | None = None
-) -> np.ndarray:
-    """
-    Compute a mask indicating which grid cells overlap with a polygon.
-    This is used to select relevant regions of the source distribution where the
-    selection is made from the polygons generated by the FrameSequence.
-
-    Parameters
-    ----------
-    polygon : (N,2) array
-        Coordinates of the polygon vertices.
-    X, Y    : (H,W) array
-        Coordinates of the grid corners.
-    tol     : float, optional
-        Tolerance for intersection checks (default = small fraction of grid spacing).
-
-    Notes:
-        Generated by ChatGPT-5.2.
-
-        Original prompt:
-        "Using numpy, can you make a function that would take in a polygon as a set
-        of vertices, and a 2d grid of xy coordinates, and return a mask that would
-        indicate which squares in the grid overlap with the polygon (even in the
-        slightest of ways). The math should be exact, I can't afford to approximate
-        the polygon with a cloud of points.
-        I don't need to know the area of overlap between squares and the polygon,
-        just whether there is overlap or not (True/False).
-
-        Refinement:
-        Can you:
-        1. Integrate bounding-box pruning per edge directly into the function
-        2. Add some tolerance when checking for equalities (maybe something like a
-        fraction of the grid spacing) because it can sometimes go wrong when a
-        polygon edge is perfectly aligned with a grid edge
-
-    Returns:
-        mask (H-1, W-1)
-    """
-
-    H, W = X.shape
-
-    # --- Cell bounds ---
-    cells_xmin = X[:-1, :-1]
-    cells_xmax = X[1:, 1:]
-    cells_ymin = Y[:-1, :-1]
-    cells_ymax = Y[1:, 1:]
-
-    # --- Auto tolerance ---
-    if tol is None:
-        dx = np.min(np.diff(X, axis=1))
-        dy = np.min(np.diff(Y, axis=0))
-        tol = 1e-9 + 1e-6 * min(dx, dy)
-
-    # Expand cell bounds slightly (robustness)
-    cells_xmin -= tol
-    cells_xmax += tol
-    cells_ymin -= tol
-    cells_ymax += tol
-
-    # --- Point in polygon ---
-    def points_in_poly(px, py, poly):
-        x = poly[:, 0]
-        y = poly[:, 1]
-        x2 = np.roll(x, -1)
-        y2 = np.roll(y, -1)
-
-        inside = np.zeros(px.shape, dtype=bool)
-
-        for i in range(len(poly)):
-            cond = ((y[i] > py) != (y2[i] > py)) & (
-                px < (x2[i] - x[i]) * (py - y[i]) / (y2[i] - y[i] + tol) + x[i]
-            )
-            inside ^= cond
-
-        return inside
-
-    # --- 1. Corner inside polygon ---
-    corners_x = np.stack([X[:-1, :-1], X[1:, :-1], X[:-1, 1:], X[1:, 1:]], axis=0)
-    corners_y = np.stack([Y[:-1, :-1], Y[1:, :-1], Y[:-1, 1:], Y[1:, 1:]], axis=0)
-
-    corner_inside = np.any(points_in_poly(corners_x, corners_y, polygon), axis=0)
-
-    # --- 2. Polygon vertex inside cell ---
-    vx = polygon[:, 0][:, None, None]
-    vy = polygon[:, 1][:, None, None]
-
-    vertex_inside = (
-        (vx >= cells_xmin)
-        & (vx <= cells_xmax)
-        & (vy >= cells_ymin)
-        & (vy <= cells_ymax)
-    ).any(axis=0)
-
-    # --- Segment intersection ---
-    def segments_intersect(p1, p2, q1, q2):
-        def orient(a, b, c):
-            return (b[..., 0] - a[..., 0]) * (c[..., 1] - a[..., 1]) - (
-                b[..., 1] - a[..., 1]
-            ) * (c[..., 0] - a[..., 0])
-
-        o1 = orient(p1, p2, q1)
-        o2 = orient(p1, p2, q2)
-        o3 = orient(q1, q2, p1)
-        o4 = orient(q1, q2, p2)
-
-        return (o1 * o2 <= tol) & (o3 * o4 <= tol)
-
-    edges_p1 = polygon
-    edges_p2 = np.roll(polygon, -1, axis=0)
-
-    intersect_mask = np.zeros((H - 1, W - 1), dtype=bool)
-
-    # Cell edges
-    cell_edges = [
-        (
-            np.stack([cells_xmin, cells_ymin], axis=-1),
-            np.stack([cells_xmax, cells_ymin], axis=-1),
-        ),
-        (
-            np.stack([cells_xmin, cells_ymax], axis=-1),
-            np.stack([cells_xmax, cells_ymax], axis=-1),
-        ),
-        (
-            np.stack([cells_xmin, cells_ymin], axis=-1),
-            np.stack([cells_xmin, cells_ymax], axis=-1),
-        ),
-        (
-            np.stack([cells_xmax, cells_ymin], axis=-1),
-            np.stack([cells_xmax, cells_ymax], axis=-1),
-        ),
-    ]
-
-    # --- 3. Edge intersection with pruning ---
-    for p1, p2 in zip(edges_p1, edges_p2, strict=True):
-        # Edge bounding box (+ tolerance)
-        xmin = min(p1[0], p2[0]) - tol
-        xmax = max(p1[0], p2[0]) + tol
-        ymin = min(p1[1], p2[1]) - tol
-        ymax = max(p1[1], p2[1]) + tol
-
-        # Candidate cells only
-        candidate = (
-            (cells_xmax >= xmin)
-            & (cells_xmin <= xmax)
-            & (cells_ymax >= ymin)
-            & (cells_ymin <= ymax)
-        )
-
-        if not np.any(candidate):
-            continue
-
-        p1e = p1[None, None, :]
-        p2e = p2[None, None, :]
-
-        for q1, q2 in cell_edges:
-            hit = segments_intersect(p1e, p2e, q1, q2)
-            intersect_mask[candidate] |= hit[candidate]
-
-    return corner_inside | vertex_inside | intersect_mask
diff --git a/src/tof/source.py b/src/tof/source.py
index d62a60b..13ab0df 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -13,7 +13,7 @@
 
 from .chopper import Chopper
 from .component import ComponentReading
-from .optimization import FrameSequence, Subframe, polygon_grid_overlap_mask
+from .optimization import FrameSequence, Subframe
 from .utils import wavelength_to_speed
 
 TIME_UNIT = "us"
@@ -201,8 +201,10 @@ def _make_pulses(
         # This leads to less iterations overall, as we can have many iterations at the
         # end when we are missing just one or two neutrons and we keep sampling them
         # outside of the accepted regions.
-        # size = max(ntot - n, ntot // 2)
-        size = ntot
+        size = max(ntot - n, ntot // 3)
+        # size = max(ntot - n, ntot // 3)
+        # size = ntot
+        # size = ntot - n
         inds = rng.choice(len(p_flat), size=size, p=p_flat.values)
         t = p_flat.coords[T_DIM].values[inds] + (
             rng.normal(scale=0.5, size=size) * widths[T_DIM].values[inds]
@@ -314,7 +316,6 @@ def _optimize_source(
             )
         ]
 
-    # X, Y = np.meshgrid(time_edges.values, wave_edges.values)
     mask = sc.zeros(sizes=p.sizes, dtype=bool)
     for poly in polygons:
         # Mask with bounding box of each polygon
@@ -325,18 +326,6 @@ def _optimize_source(
             & (wave_edges[:-1] <= poly.wavelength.max().to(unit=WAV_UNIT))
         )
 
-    #     mask |= polygon_grid_overlap_mask(
-    #         np.column_stack(
-    #             [
-    #                 poly.time.to(unit=TIME_UNIT).values,
-    #                 poly.wavelength.to(unit=WAV_UNIT).values,
-    #             ]
-    #         ),
-    #         X,
-    #         Y,
-    #     )
-
-    # prob = p.copy(deep=True)
     prob = sc.where(mask, p, sc.scalar(0.0, unit=p.unit))
     return SourceDistribution(probability=prob, acceptance_polygons=polygons)
 

From eb459b9ddf49b86ef0693fa1588583b6920f2541 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Mon, 30 Mar 2026 17:01:03 +0200
Subject: [PATCH 11/15] use uniform noise and put tight masks back in to
 optimization

---
 src/tof/optimization.py | 162 ++++++++++++++++++++++++++++++++++++++++
 src/tof/source.py       |  36 ++++-----
 2 files changed, 181 insertions(+), 17 deletions(-)

diff --git a/src/tof/optimization.py b/src/tof/optimization.py
index d0742db..0008894 100644
--- a/src/tof/optimization.py
+++ b/src/tof/optimization.py
@@ -267,3 +267,165 @@ def _chop(frame: Subframe, time: sc.Variable, close_to_open: bool) -> Subframe |
     time = sc.concat([t for t, _ in output], dim=frame.time.dim)
     wavelength = sc.concat([v for _, v in output], dim=frame.wavelength.dim)
     return Subframe(time=time, wavelength=wavelength)
+
+
+def polygon_grid_overlap_mask(
+    polygon: np.ndarray, X: np.ndarray, Y: np.ndarray, tol: float | None = None
+) -> np.ndarray:
+    """
+    Compute a mask indicating which grid cells overlap with a polygon.
+    This is used to select relevant regions of the source distribution where the
+    selection is made from the polygons generated by the FrameSequence.
+
+    Parameters
+    ----------
+    polygon : (N,2) array
+        Coordinates of the polygon vertices.
+    X, Y    : (H,W) array
+        Coordinates of the grid corners.
+    tol     : float, optional
+        Tolerance for intersection checks (default = small fraction of grid spacing).
+
+    Notes:
+        Generated by ChatGPT-5.2.
+
+        Original prompt:
+        "Using numpy, can you make a function that would take in a polygon as a set
+        of vertices, and a 2d grid of xy coordinates, and return a mask that would
+        indicate which squares in the grid overlap with the polygon (even in the
+        slightest of ways). The math should be exact, I can't afford to approximate
+        the polygon with a cloud of points.
+        I don't need to know the area of overlap between squares and the polygon,
+        just whether there is overlap or not (True/False).
+
+        Refinement:
+        Can you:
+        1. Integrate bounding-box pruning per edge directly into the function
+        2. Add some tolerance when checking for equalities (maybe something like a
+        fraction of the grid spacing) because it can sometimes go wrong when a
+        polygon edge is perfectly aligned with a grid edge
+
+    Returns:
+        mask (H-1, W-1)
+    """
+
+    H, W = X.shape
+
+    # --- Cell bounds ---
+    cells_xmin = X[:-1, :-1]
+    cells_xmax = X[1:, 1:]
+    cells_ymin = Y[:-1, :-1]
+    cells_ymax = Y[1:, 1:]
+
+    # --- Auto tolerance ---
+    if tol is None:
+        dx = np.min(np.diff(X, axis=1))
+        dy = np.min(np.diff(Y, axis=0))
+        tol = 1e-9 + 1e-6 * min(dx, dy)
+
+    # Expand cell bounds slightly (robustness)
+    cells_xmin -= tol
+    cells_xmax += tol
+    cells_ymin -= tol
+    cells_ymax += tol
+
+    # --- Point in polygon ---
+    def points_in_poly(px, py, poly):
+        x = poly[:, 0]
+        y = poly[:, 1]
+        x2 = np.roll(x, -1)
+        y2 = np.roll(y, -1)
+
+        inside = np.zeros(px.shape, dtype=bool)
+
+        for i in range(len(poly)):
+            cond = ((y[i] > py) != (y2[i] > py)) & (
+                px < (x2[i] - x[i]) * (py - y[i]) / (y2[i] - y[i] + tol) + x[i]
+            )
+            inside ^= cond
+
+        return inside
+
+    # --- 1. Corner inside polygon ---
+    corners_x = np.stack([X[:-1, :-1], X[1:, :-1], X[:-1, 1:], X[1:, 1:]], axis=0)
+    corners_y = np.stack([Y[:-1, :-1], Y[1:, :-1], Y[:-1, 1:], Y[1:, 1:]], axis=0)
+
+    corner_inside = np.any(points_in_poly(corners_x, corners_y, polygon), axis=0)
+
+    # --- 2. Polygon vertex inside cell ---
+    vx = polygon[:, 0][:, None, None]
+    vy = polygon[:, 1][:, None, None]
+
+    vertex_inside = (
+        (vx >= cells_xmin)
+        & (vx <= cells_xmax)
+        & (vy >= cells_ymin)
+        & (vy <= cells_ymax)
+    ).any(axis=0)
+
+    # --- Segment intersection ---
+    def segments_intersect(p1, p2, q1, q2):
+        def orient(a, b, c):
+            return (b[..., 0] - a[..., 0]) * (c[..., 1] - a[..., 1]) - (
+                b[..., 1] - a[..., 1]
+            ) * (c[..., 0] - a[..., 0])
+
+        o1 = orient(p1, p2, q1)
+        o2 = orient(p1, p2, q2)
+        o3 = orient(q1, q2, p1)
+        o4 = orient(q1, q2, p2)
+
+        return (o1 * o2 <= tol) & (o3 * o4 <= tol)
+
+    edges_p1 = polygon
+    edges_p2 = np.roll(polygon, -1, axis=0)
+
+    intersect_mask = np.zeros((H - 1, W - 1), dtype=bool)
+
+    # Cell edges
+    cell_edges = [
+        (
+            np.stack([cells_xmin, cells_ymin], axis=-1),
+            np.stack([cells_xmax, cells_ymin], axis=-1),
+        ),
+        (
+            np.stack([cells_xmin, cells_ymax], axis=-1),
+            np.stack([cells_xmax, cells_ymax], axis=-1),
+        ),
+        (
+            np.stack([cells_xmin, cells_ymin], axis=-1),
+            np.stack([cells_xmin, cells_ymax], axis=-1),
+        ),
+        (
+            np.stack([cells_xmax, cells_ymin], axis=-1),
+            np.stack([cells_xmax, cells_ymax], axis=-1),
+        ),
+    ]
+
+    # --- 3. Edge intersection with pruning ---
+    for p1, p2 in zip(edges_p1, edges_p2, strict=True):
+        # Edge bounding box (+ tolerance)
+        xmin = min(p1[0], p2[0]) - tol
+        xmax = max(p1[0], p2[0]) + tol
+        ymin = min(p1[1], p2[1]) - tol
+        ymax = max(p1[1], p2[1]) + tol
+
+        # Candidate cells only
+        candidate = (
+            (cells_xmax >= xmin)
+            & (cells_xmin <= xmax)
+            & (cells_ymax >= ymin)
+            & (cells_ymin <= ymax)
+        )
+
+        if not np.any(candidate):
+            continue
+
+        p1e = p1[None, None, :]
+        p2e = p2[None, None, :]
+
+        for q1, q2 in cell_edges:
+            hit = segments_intersect(p1e, p2e, q1, q2)
+            intersect_mask[candidate] |= hit[candidate]
+
+    return corner_inside | vertex_inside | intersect_mask
diff --git a/src/tof/source.py b/src/tof/source.py
index 13ab0df..cb795ef 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -13,7 +13,7 @@
 
 from .chopper import Chopper
 from .component import ComponentReading
-from .optimization import FrameSequence, Subframe
+from .optimization import FrameSequence, Subframe, polygon_grid_overlap_mask
 from .utils import wavelength_to_speed
 
 TIME_UNIT = "us"
@@ -171,9 +171,9 @@ def _make_pulses(
         for dim in (T_DIM, W_DIM)
     }
 
-    # Because of the added noise, some values end up being outside the specified range
-    # for the birth times and wavelengths. Using naive clipping leads to pile-up on the
-    # edges of the range. To avoid this, we remove the outliers and resample until we
+    # When using an acceptance region (either through time/wavelength limits, or
+    # acceptance polygon), some values end up being outside the specified range
+    # for the birth times and wavelengths. We remove the outliers and resample until we
     # have the desired number of neutrons.
     n = 0
     times = []
@@ -202,15 +202,12 @@ def _make_pulses(
         # end when we are missing just one or two neutrons and we keep sampling them
         # outside of the accepted regions.
         size = max(ntot - n, ntot // 3)
-        # size = max(ntot - n, ntot // 3)
-        # size = ntot
-        # size = ntot - n
         inds = rng.choice(len(p_flat), size=size, p=p_flat.values)
         t = p_flat.coords[T_DIM].values[inds] + (
-            rng.normal(scale=0.5, size=size) * widths[T_DIM].values[inds]
+            rng.uniform(-0.5, 0.5, size=size) * widths[T_DIM].values[inds]
         )
         w = p_flat.coords[W_DIM].values[inds] + (
-            rng.normal(scale=0.5, size=size) * widths[W_DIM].values[inds]
+            rng.uniform(-0.5, 0.5, size=size) * widths[W_DIM].values[inds]
         )
 
         # We filter events that are outside the accepted regions. Accepted regions can
@@ -316,17 +313,22 @@ def _optimize_source(
             )
         ]
 
-    mask = sc.zeros(sizes=p.sizes, dtype=bool)
+    X, Y = np.meshgrid(time_edges.values, wave_edges.values)
+    mask = np.zeros(shape=p.shape, dtype=bool)
     for poly in polygons:
-        # Mask with bounding box of each polygon
-        mask |= (
-            (time_edges[1:] >= poly.time.min().to(unit=TIME_UNIT))
-            & (time_edges[:-1] <= poly.time.max().to(unit=TIME_UNIT))
-            & (wave_edges[1:] >= poly.wavelength.min().to(unit=WAV_UNIT))
-            & (wave_edges[:-1] <= poly.wavelength.max().to(unit=WAV_UNIT))
+        mask |= polygon_grid_overlap_mask(
+            np.column_stack(
+                [
+                    poly.time.to(unit=TIME_UNIT).values,
+                    poly.wavelength.to(unit=WAV_UNIT).values,
+                ]
+            ),
+            X,
+            Y,
         )
 
-    prob = sc.where(mask, p, sc.scalar(0.0, unit=p.unit))
+    prob = p.copy(deep=True)
+    prob.values = np.where(mask, p.values, 0.0)
     return SourceDistribution(probability=prob, acceptance_polygons=polygons)
 
 

From 681f47d9aa2a0cc4a3568761aa22735f2575a0eb Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Tue, 7 Apr 2026 20:47:36 +0200
Subject: [PATCH 12/15] improve performance of masking prob distribution

---
 src/tof/optimization.py | 182 +++++++++++++++++++++-------------------
 src/tof/source.py       |  22 ++---
 2 files changed, 105 insertions(+), 99 deletions(-)

diff --git a/src/tof/optimization.py b/src/tof/optimization.py
index 0008894..88e5e35 100644
--- a/src/tof/optimization.py
+++ b/src/tof/optimization.py
@@ -270,7 +270,7 @@ def _chop(frame: Subframe, time: sc.Variable, close_to_open: bool) -> Subframe |
 
 
 def polygon_grid_overlap_mask(
-    polygon: np.ndarray, X: np.ndarray, Y: np.ndarray, tol: float | None = None
+    polygons: list[np.ndarray], X: np.ndarray, Y: np.ndarray, tol: float | None = None
 ) -> np.ndarray:
     """
     Compute a mask indicating which grid cells overlap with a polygon.
@@ -279,8 +279,8 @@ def polygon_grid_overlap_mask(
 
     Parameters
     ----------
-    polygon : (N,2) array
-        Coordinates of the polygon vertices.
+    polygons : list of (N,2) arrays
+        List of polygons, each defined by coordinates of its vertices.
     X, Y    : (H,W) array
         Coordinates of the grid corners.
     tol     : float, optional
@@ -304,6 +304,8 @@ def polygon_grid_overlap_mask(
         2. Add some tolerance when checking for equalities (maybe something like a
         fraction of the grid spacing) because it can sometimes go wrong when a
         polygon edge is perfectly aligned with a grid edge
+        3. It is a bit slow, can you improve the performance replacing boolean masks
+        with index slicing pruning?
 
     Returns:
         mask (H-1, W-1)
@@ -311,25 +313,25 @@ def polygon_grid_overlap_mask(
 
     H, W = X.shape
 
-    # --- Cell bounds ---
-    cells_xmin = X[:-1, :-1]
-    cells_xmax = X[1:, 1:]
-    cells_ymin = Y[:-1, :-1]
-    cells_ymax = Y[1:, 1:]
+    # Assume structured grid
+    x = X[0, :]
+    y = Y[:, 0]
+
+    dx = np.min(np.diff(x))
+    dy = np.min(np.diff(y))
 
-    # --- Auto tolerance ---
     if tol is None:
-        dx = np.min(np.diff(X, axis=1))
-        dy = np.min(np.diff(Y, axis=0))
         tol = 1e-9 + 1e-6 * min(dx, dy)
 
-    # Expand cell bounds slightly (robustness)
-    cells_xmin -= tol
-    cells_xmax += tol
-    cells_ymin -= tol
-    cells_ymax += tol
+    # Cell bounds
+    cells_xmin = X[:-1, :-1] - tol
+    cells_xmax = X[1:, 1:] + tol
+    cells_ymin = Y[:-1, :-1] - tol
+    cells_ymax = Y[1:, 1:] + tol
+
+    union_mask = np.zeros((H - 1, W - 1), dtype=bool)
 
-    # --- Point in polygon ---
+    # --- Helpers ---
     def points_in_poly(px, py, poly):
         x = poly[:, 0]
         y = poly[:, 1]
@@ -337,33 +339,13 @@ def points_in_poly(px, py, poly):
         y2 = np.roll(y, -1)
 
         inside = np.zeros(px.shape, dtype=bool)
-
         for i in range(len(poly)):
             cond = ((y[i] > py) != (y2[i] > py)) & (
                 px < (x2[i] - x[i]) * (py - y[i]) / (y2[i] - y[i] + tol) + x[i]
             )
             inside ^= cond
-
         return inside
 
-    # --- 1. Corner inside polygon ---
-    corners_x = np.stack([X[:-1, :-1], X[1:, :-1], X[:-1, 1:], X[1:, 1:]], axis=0)
-    corners_y = np.stack([Y[:-1, :-1], Y[1:, :-1], Y[:-1, 1:], Y[1:, 1:]], axis=0)
-
-    corner_inside = np.any(points_in_poly(corners_x, corners_y, polygon), axis=0)
-
-    # --- 2. Polygon vertex inside cell ---
-    vx = polygon[:, 0][:, None, None]
-    vy = polygon[:, 1][:, None, None]
-
-    vertex_inside = (
-        (vx >= cells_xmin)
-        & (vx <= cells_xmax)
-        & (vy >= cells_ymin)
-        & (vy <= cells_ymax)
-    ).any(axis=0)
-
-    # --- Segment intersection ---
     def segments_intersect(p1, p2, q1, q2):
         def orient(a, b, c):
             return (b[..., 0] - a[..., 0]) * (c[..., 1] - a[..., 1]) - (
@@ -377,55 +359,85 @@ def orient(a, b, c):
 
         return (o1 * o2 <= tol) & (o3 * o4 <= tol)
 
-    edges_p1 = polygon
-    edges_p2 = np.roll(polygon, -1, axis=0)
-
-    intersect_mask = np.zeros((H - 1, W - 1), dtype=bool)
-
-    # Cell edges
-    cell_edges = [
-        (
-            np.stack([cells_xmin, cells_ymin], axis=-1),
-            np.stack([cells_xmax, cells_ymin], axis=-1),
-        ),
-        (
-            np.stack([cells_xmin, cells_ymax], axis=-1),
-            np.stack([cells_xmax, cells_ymax], axis=-1),
-        ),
-        (
-            np.stack([cells_xmin, cells_ymin], axis=-1),
-            np.stack([cells_xmin, cells_ymax], axis=-1),
-        ),
-        (
-            np.stack([cells_xmax, cells_ymin], axis=-1),
-            np.stack([cells_xmax, cells_ymax], axis=-1),
-        ),
-    ]
-
-    # --- 3. Edge intersection with pruning ---
-    for p1, p2 in zip(edges_p1, edges_p2, strict=True):
-        # Edge bounding box (+ tolerance)
-        xmin = min(p1[0], p2[0]) - tol
-        xmax = max(p1[0], p2[0]) + tol
-        ymin = min(p1[1], p2[1]) - tol
-        ymax = max(p1[1], p2[1]) + tol
-
-        # Candidate cells only
-        candidate = (
-            (cells_xmax >= xmin)
-            & (cells_xmin <= xmax)
-            & (cells_ymax >= ymin)
-            & (cells_ymin <= ymax)
-        )
+    # Precompute corners
+    corners_x = np.stack([X[:-1, :-1], X[1:, :-1], X[:-1, 1:], X[1:, 1:]], axis=0)
+    corners_y = np.stack([Y[:-1, :-1], Y[1:, :-1], Y[:-1, 1:], Y[1:, 1:]], axis=0)
+
+    for polygon in polygons:
+        if union_mask.all():
+            break
+
+        # --- 1. Corner test (global, cheap enough) ---
+        corner_inside = np.any(points_in_poly(corners_x, corners_y, polygon), axis=0)
+
+        # --- 2. Vertex test (localized per vertex) ---
+        vertex_inside = np.zeros_like(union_mask)
+
+        for vx, vy in polygon:
+            ix = np.searchsorted(x, vx)
+            iy = np.searchsorted(y, vy)
+
+            ix0 = max(ix - 1, 0)
+            ix1 = min(ix + 1, W - 1)
+            iy0 = max(iy - 1, 0)
+            iy1 = min(iy + 1, H - 1)
+
+            vertex_inside[iy0:iy1, ix0:ix1] |= (
+                (vx >= cells_xmin[iy0:iy1, ix0:ix1])
+                & (vx <= cells_xmax[iy0:iy1, ix0:ix1])
+                & (vy >= cells_ymin[iy0:iy1, ix0:ix1])
+                & (vy <= cells_ymax[iy0:iy1, ix0:ix1])
+            )
+
+        # --- 3. Edge intersections (INDEX PRUNED) ---
+        intersect_mask = np.zeros_like(union_mask)
+
+        edges_p1 = polygon
+        edges_p2 = np.roll(polygon, -1, axis=0)
+
+        for p1, p2 in zip(edges_p1, edges_p2, strict=True):
+            xmin = min(p1[0], p2[0]) - tol
+            xmax = max(p1[0], p2[0]) + tol
+            ymin = min(p1[1], p2[1]) - tol
+            ymax = max(p1[1], p2[1]) + tol
+
+            # Convert to index ranges
+            ix0 = max(np.searchsorted(x, xmin) - 1, 0)
+            ix1 = min(np.searchsorted(x, xmax), W - 1)
+            iy0 = max(np.searchsorted(y, ymin) - 1, 0)
+            iy1 = min(np.searchsorted(y, ymax), H - 1)
+
+            if ix0 >= ix1 or iy0 >= iy1:
+                continue
+
+            # Slice relevant region
+            sl = np.s_[iy0:iy1, ix0:ix1]
+
+            # Skip already-filled cells
+            remaining = ~union_mask[sl]
+            if not np.any(remaining):
+                continue
+
+            # Build cell edges only for slice
+            cxmin = cells_xmin[sl]
+            cxmax = cells_xmax[sl]
+            cymin = cells_ymin[sl]
+            cymax = cells_ymax[sl]
+
+            cell_edges = [
+                (np.stack([cxmin, cymin], axis=-1), np.stack([cxmax, cymin], axis=-1)),
+                (np.stack([cxmin, cymax], axis=-1), np.stack([cxmax, cymax], axis=-1)),
+                (np.stack([cxmin, cymin], axis=-1), np.stack([cxmin, cymax], axis=-1)),
+                (np.stack([cxmax, cymin], axis=-1), np.stack([cxmax, cymax], axis=-1)),
+            ]
 
-        if not np.any(candidate):
-            continue
+            p1e = p1[None, None, :]
+            p2e = p2[None, None, :]
 
-        p1e = p1[None, None, :]
-        p2e = p2[None, None, :]
+            for q1, q2 in cell_edges:
+                hit = segments_intersect(p1e, p2e, q1, q2)
+                intersect_mask[sl] |= hit & remaining
 
-        for q1, q2 in cell_edges:
-            hit = segments_intersect(p1e, p2e, q1, q2)
-            intersect_mask[candidate] |= hit[candidate]
+        union_mask |= corner_inside | vertex_inside | intersect_mask
 
-    return corner_inside | vertex_inside | intersect_mask
+    return union_mask
diff --git a/src/tof/source.py b/src/tof/source.py
index cb795ef..97c48a7 100644
--- a/src/tof/source.py
+++ b/src/tof/source.py
@@ -140,7 +140,6 @@ def _make_pulses(
     p:
         2D probability distribution for a single pulse.
     """
-
     acceptance_polygons = p.acceptance_polygons
     p = p.probability.copy(deep=False)
 
@@ -190,10 +189,7 @@ def _make_pulses(
         )
     p_flat /= p_sum
 
-    it = 0
-
     while n < ntot:
-        it += 1
         # The first iteration, we sample all neutrons. Because some get discarded, we
         # will have to iterate at least twice. Instead of just sampling the missing
         # number of neutrons, subsequent iterations over-sample (here we choose 50% of
@@ -233,9 +229,6 @@ def _make_pulses(
         times.append(t[inds])
         wavs.append(w[inds])
         n += len(inds)
-        print("Selected", len(inds), "neutrons", ntot - n, "remaining")
-
-    print("Sampling took", it, "iterations")
 
     dim = "event"
     birth_time = sc.array(
@@ -314,18 +307,19 @@ def _optimize_source(
         ]
 
     X, Y = np.meshgrid(time_edges.values, wave_edges.values)
-    mask = np.zeros(shape=p.shape, dtype=bool)
-    for poly in polygons:
-        mask |= polygon_grid_overlap_mask(
+    mask = polygon_grid_overlap_mask(
+        [
             np.column_stack(
                 [
                     poly.time.to(unit=TIME_UNIT).values,
                     poly.wavelength.to(unit=WAV_UNIT).values,
                 ]
-            ),
-            X,
-            Y,
-        )
+            )
+            for poly in polygons
+        ],
+        X,
+        Y,
+    )
 
     prob = p.copy(deep=True)
     prob.values = np.where(mask, p.values, 0.0)

From b1a71337f18f8b5511da0a2b4b5143f5e5b044b4 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Tue, 7 Apr 2026 21:02:49 +0200
Subject: [PATCH 13/15] add some tests

---
 tests/optimization_test.py | 51 ++++++++++++++++++++++++++++++++++++++
 tests/source_test.py       | 16 ++++++++++++
 2 files changed, 67 insertions(+)
 create mode 100644 tests/optimization_test.py

diff --git a/tests/optimization_test.py b/tests/optimization_test.py
new file mode 100644
index 0000000..37d337f
--- /dev/null
+++ b/tests/optimization_test.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2026 Scipp contributors (https://github.com/scipp)
+
+
+import tof
+
+
+def test_optimized_source_has_no_blocked_neutrons():
+    beamline = tof.facilities.ess.odin(pulse_skipping=True)
+    N = 100_000
+    s1 = tof.Source(facility='ess', neutrons=N)
+    m1 = tof.Model(source=s1, **beamline)
+    r1 = m1.run()
+    sum1 = r1['detector'].data.sum().value
+    assert sum1 > 0
+    assert sum1 < N
+
+    choppers = {
+        comp.name: comp
+        for comp in beamline['components']
+        if isinstance(comp, tof.Chopper)
+    }
+    s2 = tof.Source(facility='ess', neutrons=N, optimize_for=choppers)
+    m2 = tof.Model(source=s2, **beamline)
+    r2 = m2.run()
+    sum2 = r2['detector'].data.sum().value
+    assert sum2 == N
+
+
+def test_optimize_for_an_early_chopper_still_has_some_blocked_neutrons():
+    beamline = tof.facilities.ess.odin(pulse_skipping=True)
+    N = 100_000
+    s1 = tof.Source(facility='ess', neutrons=N)
+    m1 = tof.Model(source=s1, **beamline)
+    r1 = m1.run()
+    sum1 = r1['detector'].data.sum().value
+    assert sum1 > 0
+    assert sum1 < N
+
+    choppers = {
+        comp.name: comp
+        for comp in beamline['components']
+        if comp.name in ("WFM1", "WFM2")
+    }
+    s2 = tof.Source(facility='ess', neutrons=N, optimize_for=choppers)
+    m2 = tof.Model(source=s2, **beamline)
+    r2 = m2.run()
+    sum2 = r2['detector'].data.sum().value
+    assert sum2 > 0
+    assert sum2 < N
+    assert sum2 > sum1
diff --git a/tests/source_test.py b/tests/source_test.py
index 56054f1..dd4402e 100644
--- a/tests/source_test.py
+++ b/tests/source_test.py
@@ -31,6 +31,22 @@ def test_ess_pulse_uppercase(facility):
     assert source.facility == 'ess'
 
 
+def test_ess_pulse_tmin_tmax():
+    tmin = sc.scalar(1000.0, unit='us')
+    tmax = sc.scalar(2500.0, unit='us')
+    source = tof.Source(facility='ess', neutrons=100_000, tmin=tmin, tmax=tmax)
+    assert sc.all(source.data['pulse', 0].coords['birth_time'] >= tmin)
+    assert sc.all(source.data['pulse', 0].coords['birth_time'] <= tmax)
+
+
+def test_ess_pulse_wmin_wmax():
+    wmin = sc.scalar(2.0, unit='angstrom')
+    wmax = sc.scalar(6.0, unit='angstrom')
+    source = tof.Source(facility='ess', neutrons=100_000, wmin=wmin, wmax=wmax)
+    assert sc.all(source.data['pulse', 0].coords['wavelength'] >= wmin)
+    assert sc.all(source.data['pulse', 0].coords['wavelength'] <= wmax)
+
+
 def test_creation_from_neutrons():
     birth_times = sc.array(dims=['event'], values=[1000.0, 1500.0, 2000.0], unit='us')
     wavelengths = sc.array(dims=['event'], values=[1.0, 5.0, 10.0], unit='angstrom')

From 2f0b78b53e98062cf2bbe64eae1efa665e073940 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Wed, 8 Apr 2026 00:28:47 +0200
Subject: [PATCH 14/15] add notebook on optimizations

---
 docs/index.md            |   3 +-
 docs/optimizations.ipynb | 233 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 235 insertions(+), 1 deletion(-)
 create mode 100644 docs/optimizations.ipynb

diff --git a/docs/index.md b/docs/index.md
index c254beb..4382a53 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -63,8 +63,9 @@ sources
 components
 multiple-pulses
 wfm
-dashboard
+optimizations
 ess-instruments
+dashboard
 api-reference/index
 developer/index
 about/index
diff --git a/docs/optimizations.ipynb b/docs/optimizations.ipynb
new file mode 100644
index 0000000..719d6a2
--- /dev/null
+++ b/docs/optimizations.ipynb
@@ -0,0 +1,233 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# Optimizations\n",
+    "\n",
+    "## Optimize for a given chopper cascade\n",
+    "\n",
+    "Most times when we run a `tof` model,\n",
+    "the vast majority of neutrons in a pulse get blocked by the first few choppers in the beam path.\n",
+    "\n",
+    "For example, using the chopper settings for the Odin (ESS) instrument:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import scipp as sc\n",
+    "import matplotlib.pyplot as plt\n",
+    "import tof"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "s1 = tof.Source(facility=\"ess\", neutrons=1_000_000)\n",
+    "beamline = tof.facilities.ess.odin(pulse_skipping=True)\n",
+    "m1 = tof.Model(source=s1, **beamline)\n",
+    "r1 = m1.run()\n",
+    "r1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "We can see that out of 1M neutrons that left the source, over 900K were blocked by the first chopper.\n",
+    "In the end, only ~57K make it to the detector.\n",
+    "\n",
+    "This is incredibly wasteful in both memory and compute.\n",
+    "\n",
+    "We can however use the information of the opening and closing times of the choppers to predict which parts of the pulse (in the `birth_time`/`wavelength` space) will make it through and which regions will be blocked.\n",
+    "This is otherwise known as a 'chopper acceptance diagram'.\n",
+    "\n",
+    "This can be visualized by looking at the birth times and wavelengths of the neutrons that made it to the detector,\n",
+    "and compare that to the original distribution of neutrons in the source."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig1 = s1.data.hist(wavelength=300, birth_time=300).plot(norm='log', title=\"Sampled from source\")\n",
+    "fig2 = r1['detector'].data.hist(wavelength=300, birth_time=300).plot(norm='log', title=\"Neutrons that make it to the detector\")\n",
+    "\n",
+    "fig1 + fig2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "We can however use the information of the opening and closing times of the choppers to predict which parts of the pulse (in the `birth_time`/`wavelength` space) will make it through and which regions will be blocked.\n",
+    "This is otherwise known as a 'chopper acceptance diagram'.\n",
+    "\n",
+    "The source has an in-built method to apply the chopper acceptance criteria during the sampling of neutrons;\n",
+    "it is activated via the `optimize_for` argument.\n",
+    "It expects a list of choppers, and only neutrons that would make it through all choppers end up in the source."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Filter out choppers from list of Odin components\n",
+    "choppers = {comp.name: comp for comp in beamline['components'] if isinstance(comp, tof.Chopper)}\n",
+    "\n",
+    "# Create optimized source\n",
+    "s2 = tof.Source(facility=\"ess\", neutrons=1_000_000, optimize_for=choppers)\n",
+    "\n",
+    "m2 = tof.Model(source=s2, **beamline)\n",
+    "r2 = m2.run()\n",
+    "r2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "We can now see that **all 1M neutrons make to the detector**,\n",
+    "and plotting the birth time/wavelength distribution illustrates the optimization:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig1 = s2.data.hist(wavelength=300, birth_time=300).plot(norm='log', title=\"Sampled from source\")\n",
+    "fig2 = r2['detector'].data.hist(wavelength=300, birth_time=300).plot(norm='log', title=\"Neutrons that make it to the detector\")\n",
+    "\n",
+    "fig1 + fig2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9",
+   "metadata": {},
+   "source": [
+    "It is also clear that the signal recorded at detector is much less noisy due to the improved statistics.\n",
+    "It is also important to note that the overall shape of the data (relative intensities) was not changed by the optimization."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 2, figsize=(12, 4))\n",
+    "\n",
+    "r1['detector'].toa.plot(ax=ax[0])\n",
+    "r2['detector'].toa.plot(color='C1', ax=ax[0].twinx())\n",
+    "\n",
+    "r1['detector'].wavelength.plot(ax=ax[1])\n",
+    "r2['detector'].wavelength.plot(color='C1', ax=ax[1].twinx())\n",
+    "\n",
+    "fig.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "11",
+   "metadata": {},
+   "source": [
+    "## Pulse skipping\n",
+    "\n",
+    "Some instruments (such as Odin) use a pulse skipping chopper to 'skip' every other pulse, thus allowing to record a wider wavelength range at the detector without having issues where neutrons from successive pulses mix (also known as pulse-overlap).\n",
+    "\n",
+    "In such a setup, when running multiple pulses, all neutrons from every other pulse are rendered useless.\n",
+    "Yet, `tof` naively treats them as normal neutrons and tries to follow them all the way to the detector."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "s1 = tof.Source(facility=\"ess\", neutrons=1_000_000, pulses=3)\n",
+    "m1 = tof.Model(source=s1, **beamline)\n",
+    "r1 = m1.run()\n",
+    "r1.plot(blocked_rays=5000)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13",
+   "metadata": {},
+   "source": [
+    "To avoid wasting all the neutrons in the second pulse, a simple trick is to override the frequency of the source.\n",
+    "Here we set it to 7 Hz (half of the original 14 Hz), meaning that the second pulse above will not exist at all."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "s2 = tof.Source(facility=\"ess\", neutrons=1_000_000, pulses=2, frequency=sc.scalar(7, unit=\"Hz\"))\n",
+    "m2 = tof.Model(source=s2, **beamline)\n",
+    "r2 = m2.run()\n",
+    "r2.plot(blocked_rays=5000)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15",
+   "metadata": {},
+   "source": [
+    "We have now used 1/3 less neutrons to achieve the same results.\n",
+    "In combination with the `optimize_for` option introduced above, these optimizations can lead to significant speedups."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

From 0330f833be497b298d48fa48eb2e3fe1ebe45869 Mon Sep 17 00:00:00 2001
From: Neil Vaytet <neil.vaytet@ess.eu>
Date: Wed, 8 Apr 2026 10:52:09 +0200
Subject: [PATCH 15/15] fix chopper names in test

---
 tests/optimization_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/optimization_test.py b/tests/optimization_test.py
index 37d337f..50563e8 100644
--- a/tests/optimization_test.py
+++ b/tests/optimization_test.py
@@ -40,7 +40,7 @@ def test_optimize_for_an_early_chopper_still_has_some_blocked_neutrons():
     choppers = {
         comp.name: comp
         for comp in beamline['components']
-        if comp.name in ("WFM1", "WFM2")
+        if comp.name in ("WFMC_1", "WFMC_2")
     }
     s2 = tof.Source(facility='ess', neutrons=N, optimize_for=choppers)
     m2 = tof.Model(source=s2, **beamline)