essreduce 25.5.2__py3-none-any.whl → 25.7.0__py3-none-any.whl

ess/reduce/time_of_flight/lut.py

@@ -0,0 +1,478 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2025 Scipp contributors (https://github.com/scipp)
+"""
+Utilities for computing time-of-flight lookup tables from neutron simulations.
+"""
+
+from dataclasses import dataclass
+from typing import NewType
+
+import numpy as np
+import sciline as sl
+import scipp as sc
+
+from ..nexus.types import DiskChoppers
+from .types import TimeOfFlightLookupTable
+
+
+@dataclass
+class SimulationResults:
+    """
+    Results of a time-of-flight simulation used to create a lookup table.
+
+    The results (apart from ``distance``) should be flat lists (1d arrays) of length N
+    where N is the number of neutrons, containing the properties of the neutrons in the
+    simulation.
+
+    Parameters
+    ----------
+    time_of_arrival:
+        Time of arrival of the neutrons at the position where the events were recorded
+        (1d array of size N).
+    speed:
+        Speed of the neutrons, typically derived from the wavelength of the neutrons
+        (1d array of size N).
+    wavelength:
+        Wavelength of the neutrons (1d array of size N).
+    weight:
+        Weight/probability of the neutrons (1d array of size N).
+    distance:
+        Distance from the source to the position where the events were recorded
+        (single value; we assume all neutrons were recorded at the same position).
+        For a ``tof`` simulation, this is just the position of the detector where the
+        events are recorded. For a ``McStas`` simulation, this is the distance between
+        the source and the event monitor.
+    """
+
+    time_of_arrival: sc.Variable
+    speed: sc.Variable
+    wavelength: sc.Variable
+    weight: sc.Variable
+    distance: sc.Variable
+
+
+NumberOfSimulatedNeutrons = NewType("NumberOfSimulatedNeutrons", int)
+"""
+Number of neutrons simulated in the simulation that is used to create the lookup table.
+This is typically a large number, e.g., 1e6 or 1e7.
+"""
+
+LtotalRange = NewType("LtotalRange", tuple[sc.Variable, sc.Variable])
+"""
+Range (min, max) of the total length of the flight path from the source to the detector.
+This is used to create the lookup table to compute the neutron time-of-flight.
+Note that the resulting table will extend slightly beyond this range, as the supplied
+range is not necessarily a multiple of the distance resolution.
+
+Note also that the range of total flight paths is supplied manually to the workflow
+instead of being read from the input data, as it allows us to compute the expensive part
+of the workflow in advance (the lookup table) and does not need to be repeated for each
+run, or for new data coming in in the case of live data collection.
+"""
+
+DistanceResolution = NewType("DistanceResolution", sc.Variable)
+"""
+Step size of the distance axis in the lookup table.
+Should be a single scalar value with a unit of length.
+This is typically of the order of 1-10 cm.
+"""
+
+TimeResolution = NewType("TimeResolution", sc.Variable)
+"""
+Step size of the event_time_offset axis in the lookup table.
+This is basically the 'time-of-flight' resolution of the detector.
+Should be a single scalar value with a unit of time.
+This is typically of the order of 0.1-0.5 ms.
+
+Since the event_time_offset range needs to span exactly one pulse period, the final
+resolution in the lookup table will be at least the supplied value here, but may be
+smaller if the pulse period is not an integer multiple of the time resolution.
+"""
+
+
+LookupTableRelativeErrorThreshold = NewType("LookupTableRelativeErrorThreshold", float)
+"""
+Threshold for the relative standard deviation (coefficient of variation) of the
+projected time-of-flight above which values are masked.
+"""
+
+PulsePeriod = NewType("PulsePeriod", sc.Variable)
+"""
+Period of the source pulses, i.e., time between consecutive pulse starts.
+"""
+
+PulseStride = NewType("PulseStride", int)
+"""
+Stride of used pulses. Usually 1, but may be a small integer when pulse-skipping.
+"""
+
+SourcePosition = NewType("SourcePosition", sc.Variable)
+"""
+Position of the neutron source in the coordinate system of the choppers.
+"""
+
+SimulationSeed = NewType("SimulationSeed", int | None)
+"""Seed for the random number generator used in the simulation.
+"""
+
+SimulationFacility = NewType("SimulationFacility", str)
+"""
+Facility where the experiment is performed, e.g., 'ess'.
+"""
+
+
+def _mask_large_uncertainty(table: sc.DataArray, error_threshold: float):
+    """
+    Mask regions with large uncertainty with NaNs.
+    The values are modified in place in the input table.
+
+    Parameters
+    ----------
+    table:
+        Lookup table with time-of-flight as a function of distance and time-of-arrival.
+    error_threshold:
+        Threshold for the relative standard deviation (coefficient of variation) of the
+        projected time-of-flight above which values are masked.
+    """
+    # Finally, mask regions with large uncertainty with NaNs.
+    relative_error = sc.stddevs(table.data) / sc.values(table.data)
+    mask = relative_error > sc.scalar(error_threshold)
+    # Use numpy for indexing as table is 2D
+    table.values[mask.values] = np.nan
+
+
+def _compute_mean_tof_in_distance_range(
+    simulation: SimulationResults,
+    distance_bins: sc.Variable,
+    time_bins: sc.Variable,
+    distance_unit: str,
+    time_unit: str,
+    frame_period: sc.Variable,
+    time_bins_half_width: sc.Variable,
+) -> sc.DataArray:
+    """
+    Compute the mean time-of-flight inside event_time_offset bins for a given range of
+    distances.
+
+    Parameters
+    ----------
+    simulation:
+        Results of a time-of-flight simulation used to create a lookup table.
+    distance_bins:
+        Bin edges for the distance axis in the lookup table.
+    time_bins:
+        Bin edges for the event_time_offset axis in the lookup table.
+    distance_unit:
+        Unit of the distance axis.
+    time_unit:
+        Unit of the event_time_offset axis.
+    frame_period:
+        Period of the source pulses, i.e., time between consecutive pulse starts.
+    time_bins_half_width:
+        Half width of the time bins in the event_time_offset axis.
+    """
+    simulation_distance = simulation.distance.to(unit=distance_unit)
+    distances = sc.midpoints(distance_bins)
+    # Compute arrival and flight times for all neutrons
+    toas = simulation.time_of_arrival + (distances / simulation.speed).to(
+        unit=time_unit, copy=False
+    )
+    dist = distances + simulation_distance
+    tofs = dist * (sc.constants.m_n / sc.constants.h) * simulation.wavelength
+
+    data = sc.DataArray(
+        data=sc.broadcast(simulation.weight, sizes=toas.sizes),
+        coords={
+            "toa": toas,
+            "tof": tofs.to(unit=time_unit, copy=False),
+            "distance": dist,
+        },
+    ).flatten(to="event")
+
+    # Add the event_time_offset coordinate, wrapped to the frame_period
+    data.coords['event_time_offset'] = data.coords['toa'] % frame_period
+
+    # Because we staggered the mesh by half a bin width, we want the values above
+    # the last bin edge to wrap around to the first bin.
+    # Technically, those values should end up between -0.5*bin_width and 0, but
+    # a simple modulo also works here because even if they end up between 0 and
+    # 0.5*bin_width, we are (below) computing the mean between -0.5*bin_width and
+    # 0.5*bin_width and it yields the same result.
+    # data.coords['event_time_offset'] %= pulse_period - time_bins_half_width
+    data.coords['event_time_offset'] %= frame_period - time_bins_half_width
+
+    binned = data.bin(
+        distance=distance_bins + simulation_distance, event_time_offset=time_bins
+    )
+
+    # Weighted mean of tof inside each bin
+    mean_tof = (
+        binned.bins.data * binned.bins.coords["tof"]
+    ).bins.sum() / binned.bins.sum()
+    # Compute the variance of the tofs to track regions with large uncertainty
+    variance = (
+        binned.bins.data * (binned.bins.coords["tof"] - mean_tof) ** 2
+    ).bins.sum() / binned.bins.sum()
+
+    mean_tof.variances = variance.values
+    return mean_tof
+
+
+def make_tof_lookup_table(
+    simulation: SimulationResults,
+    ltotal_range: LtotalRange,
+    distance_resolution: DistanceResolution,
+    time_resolution: TimeResolution,
+    pulse_period: PulsePeriod,
+    pulse_stride: PulseStride,
+    error_threshold: LookupTableRelativeErrorThreshold,
+) -> TimeOfFlightLookupTable:
+    """
+    Compute a lookup table for time-of-flight as a function of distance and
+    time-of-arrival.
+
+    Parameters
+    ----------
+    simulation:
+        Results of a time-of-flight simulation used to create a lookup table.
+        The results should be a flat table with columns for time-of-arrival, speed,
+        wavelength, and weight.
+    ltotal_range:
+        Range of total flight path lengths from the source to the detector.
+    distance_resolution:
+        Resolution of the distance axis in the lookup table.
+    time_resolution:
+        Resolution of the time-of-arrival axis in the lookup table. Must be an integer.
+    pulse_period:
+        Period of the source pulses, i.e., time between consecutive pulse starts.
+    pulse_stride:
+        Stride of used pulses. Usually 1, but may be a small integer when
+        pulse-skipping.
+    error_threshold:
+        Threshold for the relative standard deviation (coefficient of variation) of the
+        projected time-of-flight above which values are masked.
+
+    Notes
+    -----
+
+    Below are some details about the binning and wrapping around frame period in the
+    time dimension.
+
+    We have some simulated ``toa`` (events) from a Tof/McStas simulation.
+    Those are absolute ``toa``, unwrapped.
+    First we compute the usual ``event_time_offset = toa % frame_period``.
+
+    Now, we want to ensure periodic boundaries. If we make a bin centered around 0,
+    and a bin centered around 71ms: the first bin will use events between 0 and
+    ``0.5 * dt`` (where ``dt`` is the bin width).
+    The last bin will use events between ``frame_period - 0.5*dt`` and
+    ``frame_period + 0.5 * dt``. So when we compute the mean inside those two bins,
+    they will not yield the same results.
+    It is as if the first bin is missing the events it should have between
+    ``-0.5 * dt`` and 0 (because of the modulo we computed above).
+
+    To fix this, we do not make a last bin around 71ms (the bins stop at
+    ``frame_period - 0.5*dt``). Instead, we compute modulo a second time,
+    but this time using ``event_time_offset %= (frame_period - 0.5*dt)``.
+    (we cannot directly do ``event_time_offset = toa % (frame_period - 0.5*dt)`` in a
+    single step because it would introduce a gradual shift,
+    as the pulse number increases).
+
+    This second modulo effectively takes all the events that would have gone in the
+    last bin (between ``frame_period - 0.5*dt`` and ``frame_period``) and puts them in
+    the first bin. Instead of placing them between ``-0.5*dt`` and 0,
+    it places them between 0 and ``0.5*dt``, but this does not really matter,
+    because we then take the mean inside the first bin.
+    Whether the events are on the left or right side of zero does not matter.
+
+    Finally, we make a copy of the left edge, and append it to the right of the table,
+    thus ensuring that the values on the right edge are strictly the same as on the
+    left edge.
+    """
+    distance_unit = "m"
+    time_unit = simulation.time_of_arrival.unit
+    res = distance_resolution.to(unit=distance_unit)
+    pulse_period = pulse_period.to(unit=time_unit)
+    frame_period = pulse_period * pulse_stride
+
+    min_dist, max_dist = (
+        x.to(unit=distance_unit) - simulation.distance.to(unit=distance_unit)
+        for x in ltotal_range
+    )
+    # We need to bin the data below, to compute the weighted mean of the wavelength.
+    # This results in data with bin edges.
+    # However, the 2d interpolator expects bin centers.
+    # We want to give the 2d interpolator a table that covers the requested range,
+    # hence we need to extend the range by at least half a resolution in each direction.
+    # Then, we make the choice that the resolution in distance is the quantity that
+    # should be preserved. Because the difference between min and max distance is
+    # not necessarily an integer multiple of the resolution, we need to add a pad to
+    # ensure that the last bin is not cut off. We want the upper edge to be higher than
+    # the maximum distance, hence we pad with an additional 1.5 x resolution.
+    pad = 2.0 * res
+    distance_bins = sc.arange('distance', min_dist - pad, max_dist + pad, res)
+
+    # Create some time bins for event_time_offset.
+    # We want our final table to strictly cover the range [0, frame_period].
+    # However, binning the data associates mean values inside the bins to the bin
+    # centers. Instead, we stagger the mesh by half a bin width so we are computing
+    # values for the final mesh edges (the bilinear interpolation needs values on the
+    # edges/corners).
+    nbins = int(frame_period / time_resolution.to(unit=time_unit)) + 1
+    time_bins = sc.linspace(
+        'event_time_offset', 0.0, frame_period.value, nbins + 1, unit=pulse_period.unit
+    )
+    time_bins_half_width = 0.5 * (time_bins[1] - time_bins[0])
+    time_bins -= time_bins_half_width
+
+    # To avoid a too large RAM usage, we compute the table in chunks, and piece them
+    # together at the end.
+    ndist = len(distance_bins) - 1
+    max_size = 2e7
+    total_size = ndist * len(simulation.time_of_arrival)
+    nchunks = total_size / max_size
+    chunk_size = int(ndist / nchunks) + 1
+    pieces = []
+    for i in range(int(nchunks) + 1):
+        dist_edges = distance_bins[i * chunk_size : (i + 1) * chunk_size + 1]
+
+        pieces.append(
+            _compute_mean_tof_in_distance_range(
+                simulation=simulation,
+                distance_bins=dist_edges,
+                time_bins=time_bins,
+                distance_unit=distance_unit,
+                time_unit=time_unit,
+                frame_period=frame_period,
+                time_bins_half_width=time_bins_half_width,
+            )
+        )
+
+    table = sc.concat(pieces, 'distance')
+    table.coords["distance"] = sc.midpoints(table.coords["distance"])
+    table.coords["event_time_offset"] = sc.midpoints(table.coords["event_time_offset"])
+
+    # Copy the left edge to the right to create periodic boundary conditions
+    table = sc.DataArray(
+        data=sc.concat(
+            [table.data, table.data['event_time_offset', 0]], dim='event_time_offset'
+        ),
+        coords={
+            "distance": table.coords["distance"],
+            "event_time_offset": sc.concat(
+                [table.coords["event_time_offset"], frame_period],
+                dim='event_time_offset',
+            ),
+            "pulse_period": pulse_period,
+            "pulse_stride": sc.scalar(pulse_stride, unit=None),
+            "distance_resolution": table.coords["distance"][1]
+            - table.coords["distance"][0],
+            "time_resolution": table.coords["event_time_offset"][1]
+            - table.coords["event_time_offset"][0],
+            "error_threshold": sc.scalar(error_threshold),
+        },
+    )
+
+    # In-place masking for better performance
+    _mask_large_uncertainty(table, error_threshold)
+
+    return TimeOfFlightLookupTable(table)
+
+
+def simulate_chopper_cascade_using_tof(
+    choppers: DiskChoppers,
+    source_position: SourcePosition,
+    neutrons: NumberOfSimulatedNeutrons,
+    pulse_stride: PulseStride,
+    seed: SimulationSeed,
+    facility: SimulationFacility,
+) -> SimulationResults:
+    """
+    Simulate a pulse of neutrons propagating through a chopper cascade using the
+    ``tof`` package (https://tof.readthedocs.io).
+
+    Parameters
+    ----------
+    choppers:
+        A dict of DiskChopper objects representing the choppers in the beamline. See
+        https://scipp.github.io/scippneutron/user-guide/chopper/processing-nexus-choppers.html#Build-DiskChopper
+        for more information.
+    source_position:
+        A scalar variable with ``dtype=vector3`` that defines the source position.
+        Must be in the same coordinate system as the choppers' axle positions.
+    neutrons:
+        Number of neutrons to simulate.
+    pulse_stride:
+        The pulse strinde; we need to simulate at least enough pulses to cover the
+        requested stride.
+    seed:
+        Seed for the random number generator used in the simulation.
+    facility:
+        Facility where the experiment is performed.
+    """
+    import tof
+
+    tof_choppers = [
+        tof.Chopper(
+            frequency=abs(ch.frequency),
+            direction=tof.AntiClockwise
+            if (ch.frequency.value > 0.0)
+            else tof.Clockwise,
+            open=ch.slit_begin,
+            close=ch.slit_end,
+            phase=abs(ch.phase),
+            distance=sc.norm(
+                ch.axle_position - source_position.to(unit=ch.axle_position.unit)
+            ),
+            name=name,
+        )
+        for name, ch in choppers.items()
+    ]
+    source = tof.Source(
+        facility=facility, neutrons=neutrons, pulses=pulse_stride, seed=seed
+    )
+    if not tof_choppers:
+        events = source.data.squeeze().flatten(to='event')
+        return SimulationResults(
+            time_of_arrival=events.coords["birth_time"],
+            speed=events.coords["speed"],
+            wavelength=events.coords["wavelength"],
+            weight=events.data,
+            distance=0.0 * sc.units.m,
+        )
+    model = tof.Model(source=source, choppers=tof_choppers)
+    results = model.run()
+    # Find name of the furthest chopper in tof_choppers
+    furthest_chopper = max(tof_choppers, key=lambda c: c.distance)
+    events = results[furthest_chopper.name].data.squeeze().flatten(to='event')
+    events = events[
+        ~(events.masks["blocked_by_others"] | events.masks["blocked_by_me"])
+    ]
+    return SimulationResults(
+        time_of_arrival=events.coords["toa"],
+        speed=events.coords["speed"],
+        wavelength=events.coords["wavelength"],
+        weight=events.data,
+        distance=furthest_chopper.distance,
+    )
+
+
+def TofLookupTableWorkflow():
+    """
+    Create a workflow for computing a time-of-flight lookup table from a
+    simulation of neutrons propagating through a chopper cascade.
+    """
+    wf = sl.Pipeline(
+        (make_tof_lookup_table, simulate_chopper_cascade_using_tof),
+        params={
+            PulsePeriod: 1.0 / sc.scalar(14.0, unit="Hz"),
+            PulseStride: 1,
+            DistanceResolution: sc.scalar(0.1, unit="m"),
+            TimeResolution: sc.scalar(250.0, unit='us'),
+            LookupTableRelativeErrorThreshold: 0.1,
+            NumberOfSimulatedNeutrons: 1_000_000,
+            SimulationSeed: None,
+            SimulationFacility: 'ess',
+        },
+    )
+    return wf

ess/reduce/time_of_flight/resample.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2025 Scipp contributors (https://github.com/scipp)
+
+
+import numpy as np
+import scipp as sc
+
+
+def find_strictly_increasing_sections(var: sc.Variable) -> list[slice]:
+    """
+    Find strictly increasing sections in a coordinate dimension (minimum length 2).
+
+    Parameters
+    ----------
+    var:
+        The variable to analyze, which should be one-dimensional.
+
+    Returns
+    -------
+    sections:
+        Slice objects that can be used extract strictly increasing sections.
+    """
+    values = var.values
+    finite = np.isfinite(values)
+    increasing = (np.sign(np.diff(values)) > 0) & finite[:-1] & finite[1:]
+    # 1 marks the start of an increasing section, -1 marks the end
+    transitions = np.diff(np.concatenate(([False], increasing, [False])).astype(int))
+    section_starts = np.where(transitions == 1)[0]
+    section_ends = np.where(transitions == -1)[0] + np.array(1)
+    return [
+        slice(start, end)
+        for start, end in zip(section_starts, section_ends, strict=True)
+        if end - start >= 2  # Ensure section has at least 2 points
+    ]
+
+
+def get_min_max(
+    var: sc.Variable, *, dim: str, slices: list[slice]
+) -> tuple[sc.Variable, sc.Variable]:
+    if not slices:
+        raise ValueError("No strictly increasing sections found.")
+    combined = sc.concat([var[dim, slice] for slice in slices], dim)
+    return combined.min(), combined.max()
+
+
+def make_regular_grid(
+    var: sc.Variable, *, dim: str, slices: list[slice]
+) -> sc.Variable:
+    """
+    Create a regular grid variable based on the min and max of the slices.
+
+    The grid is constructed such that it includes the minimum and maximum values
+    of the strictly increasing sections, with a step size equal to the difference
+    between the first two values of the section with the minimum start value (which is
+    not necessarily the first section).
+    """
+    min_val, max_val = get_min_max(var, dim=dim, slices=slices)
+    first: sc.Variable | None = None
+    for s in slices:
+        first = var[dim, s]
+        if sc.identical(first[0], min_val):
+            break
+    if first is None:
+        # This should not happen if slices are correctly identified and passed from
+        # find_strictly_increasing_sections.
+        raise ValueError("Section is not strictly increasing.")
+    step = first[1] - first[0]
+    return sc.arange(
+        dim=dim,
+        start=min_val.value,
+        stop=max_val.value + step.value,  # Ensure the last bin edge is included
+        step=step.value,
+        unit=step.unit,
+        dtype=step.dtype,
+    )
+
+
+def rebin_strictly_increasing(da: sc.DataArray, dim: str) -> sc.DataArray:
+    """
+    Find strictly monotonic sections in a coordinate dimension and rebin the data array
+    into a regular grid based on these sections.
+    """
+    # Ensure the dimension is named like the coordinate.
+    da = da.rename_dims({da.coords[dim].dim: dim})
+    slices = find_strictly_increasing_sections(da.coords[dim])
+    if len(slices) == 1:
+        return da[dim, slices[0]]
+    if not slices:
+        raise ValueError("No strictly increasing sections found.")
+    if da.coords[dim].dtype not in (sc.DType.float64, sc.DType.float32):
+        # rebin does not like integer coords.
+        da = da.assign_coords({dim: da.coords[dim].to(dtype='float64')})
+    # Slices refer to the indices in the coord, which are bin edges. For slicing data
+    # we need to stop at the last index minus one.
+    sections = [da[dim, section.start : section.stop - 1] for section in slices]
+    edges = make_regular_grid(da.coords[dim], dim=dim, slices=slices)
+    return sc.reduce([sc.rebin(section, {dim: edges}) for section in sections]).sum()