essreduce 24.12.0__py3-none-any.whl → 25.1.0__py3-none-any.whl

ess/reduce/time_of_flight/toa_to_tof.py

@@ -0,0 +1,541 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2025 Scipp contributors (https://github.com/scipp)
+"""
+Time-of-flight workflow for unwrapping the time of arrival of the neutron at the
+detector.
+This workflow is used to convert raw detector data with event_time_zero and
+event_time_offset coordinates to data with a time-of-flight coordinate.
+"""
+
+from collections.abc import Callable
+
+import numpy as np
+import scipp as sc
+from scipp._scipp.core import _bins_no_validate
+from scippneutron._utils import elem_unit
+
+from .to_events import to_events
+from .types import (
+    DistanceResolution,
+    FastestNeutron,
+    FrameFoldedTimeOfArrival,
+    FramePeriod,
+    LookupTableRelativeErrorThreshold,
+    Ltotal,
+    LtotalRange,
+    MaskedTimeOfFlightLookupTable,
+    PivotTimeAtDetector,
+    PulsePeriod,
+    PulseStride,
+    PulseStrideOffset,
+    RawData,
+    ResampledTofData,
+    SimulationResults,
+    TimeOfArrivalMinusPivotTimeModuloPeriod,
+    TimeOfArrivalResolution,
+    TimeOfFlightLookupTable,
+    TofData,
+    UnwrappedTimeOfArrival,
+    UnwrappedTimeOfArrivalMinusPivotTime,
+)
+
+
+def frame_period(pulse_period: PulsePeriod, pulse_stride: PulseStride) -> FramePeriod:
+    """
+    Return the period of a frame, which is defined by the pulse period times the pulse
+    stride.
+
+    Parameters
+    ----------
+    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.
+    """
+    return FramePeriod(pulse_period * pulse_stride)
+
+
+def extract_ltotal(da: RawData) -> Ltotal:
+    """
+    Extract the total length of the flight path from the source to the detector from the
+    detector data.
+
+    Parameters
+    ----------
+    da:
+        Raw detector data loaded from a NeXus file, e.g., NXdetector containing
+        NXevent_data.
+    """
+    return Ltotal(da.coords["Ltotal"])
+
+
+def compute_tof_lookup_table(
+    simulation: SimulationResults,
+    ltotal_range: LtotalRange,
+    distance_resolution: DistanceResolution,
+    toa_resolution: TimeOfArrivalResolution,
+) -> 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.
+    toa_resolution:
+        Resolution of the time-of-arrival axis in the lookup table.
+    """
+    distance_unit = "m"
+    res = distance_resolution.to(unit=distance_unit)
+    simulation_distance = simulation.distance.to(unit=distance_unit)
+
+    min_dist, max_dist = (
+        x.to(unit=distance_unit) - simulation_distance 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
+    dist_edges = sc.array(
+        dims=["distance"],
+        values=np.arange((min_dist - pad).value, (max_dist + pad).value, res.value),
+        unit=distance_unit,
+    )
+    distances = sc.midpoints(dist_edges)
+
+    time_unit = simulation.time_of_arrival.unit
+    toas = simulation.time_of_arrival + (distances / simulation.speed).to(
+        unit=time_unit, copy=False
+    )
+
+    # Compute time-of-flight for all neutrons
+    wavs = sc.broadcast(simulation.wavelength.to(unit="m"), sizes=toas.sizes).flatten(
+        to="event"
+    )
+    dist = sc.broadcast(distances + simulation_distance, sizes=toas.sizes).flatten(
+        to="event"
+    )
+    tofs = dist * sc.constants.m_n
+    tofs *= wavs
+    tofs /= sc.constants.h
+
+    data = sc.DataArray(
+        data=sc.broadcast(simulation.weight, sizes=toas.sizes).flatten(to="event"),
+        coords={
+            "toa": toas.flatten(to="event"),
+            "tof": tofs.to(unit=time_unit, copy=False),
+            "distance": dist,
+        },
+    )
+
+    binned = data.bin(distance=dist_edges + simulation_distance, toa=toa_resolution)
+    # 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
+
+    # Convert coordinates to midpoints
+    mean_tof.coords["toa"] = sc.midpoints(mean_tof.coords["toa"])
+    mean_tof.coords["distance"] = sc.midpoints(mean_tof.coords["distance"])
+
+    return TimeOfFlightLookupTable(mean_tof)
+
+
+def masked_tof_lookup_table(
+    tof_lookup: TimeOfFlightLookupTable,
+    error_threshold: LookupTableRelativeErrorThreshold,
+) -> MaskedTimeOfFlightLookupTable:
+    """
+    Mask regions of the lookup table where the variance of the projected time-of-flight
+    is larger than a given threshold.
+
+    Parameters
+    ----------
+    tof_lookup:
+        Lookup table giving time-of-flight as a function of distance and
+        time-of-arrival.
+    variance_threshold:
+        Threshold for the variance of the projected time-of-flight above which regions
+        are masked.
+    """
+    relative_error = sc.stddevs(tof_lookup.data) / sc.values(tof_lookup.data)
+    mask = relative_error > sc.scalar(error_threshold)
+    out = tof_lookup.copy()
+    # Use numpy for indexing as table is 2D
+    out.values[mask.values] = np.nan
+    return MaskedTimeOfFlightLookupTable(out)
+
+
+def find_fastest_neutron(simulation: SimulationResults) -> FastestNeutron:
+    """
+    Find the fastest neutron in the simulation results.
+    """
+    ind = np.argmax(simulation.speed.values)
+    return FastestNeutron(
+        time_of_arrival=simulation.time_of_arrival[ind],
+        speed=simulation.speed[ind],
+        distance=simulation.distance,
+    )
+
+
+def pivot_time_at_detector(
+    fastest_neutron: FastestNeutron, ltotal: Ltotal
+) -> PivotTimeAtDetector:
+    """
+    Compute the pivot time at the detector, i.e., the time of the start of the frame at
+    the detector.
+    The assumption here is that the fastest neutron in the simulation results is the one
+    that arrives at the detector first.
+    One could have an edge case where a slightly slower neutron which is born earlier
+    could arrive at the detector first, but this edge case is most probably uncommon,
+    and the difference in arrival times is likely to be small.
+
+    Parameters
+    ----------
+    fastest_neutron:
+        Properties of the fastest neutron in the simulation results.
+    ltotal:
+        Total length of the flight path from the source to the detector.
+    """
+    dist = ltotal - fastest_neutron.distance.to(unit=ltotal.unit)
+    toa = fastest_neutron.time_of_arrival + (dist / fastest_neutron.speed).to(
+        unit=fastest_neutron.time_of_arrival.unit, copy=False
+    )
+    return PivotTimeAtDetector(toa)
+
+
+def unwrapped_time_of_arrival(
+    da: RawData, offset: PulseStrideOffset, pulse_period: PulsePeriod
+) -> UnwrappedTimeOfArrival:
+    """
+    Compute the unwrapped time of arrival of the neutron at the detector.
+    For event data, this is essentially ``event_time_offset + event_time_zero``.
+
+    Parameters
+    ----------
+    da:
+        Raw detector data loaded from a NeXus file, e.g., NXdetector containing
+        NXevent_data.
+    offset:
+        Integer offset of the first pulse in the stride (typically zero unless we are
+        using pulse-skipping and the events do not begin with the first pulse in the
+        stride).
+    pulse_period:
+        Period of the source pulses, i.e., time between consecutive pulse starts.
+    """
+    if da.bins is None:
+        # 'time_of_flight' is the canonical name in NXmonitor, but in some files, it
+        # may be called 'tof'.
+        key = next(iter(set(da.coords.keys()) & {"time_of_flight", "tof"}))
+        toa = da.coords[key]
+    else:
+        # To unwrap the time of arrival, we want to add the event_time_zero to the
+        # event_time_offset. However, we do not really care about the exact datetimes,
+        # we just want to know the offsets with respect to the start of the run.
+        # Hence we use the smallest event_time_zero as the time origin.
+        time_zero = da.coords["event_time_zero"] - da.coords["event_time_zero"].min()
+        coord = da.bins.coords["event_time_offset"]
+        unit = elem_unit(coord)
+        toa = (
+            coord
+            + time_zero.to(dtype=float, unit=unit, copy=False)
+            - (offset * pulse_period).to(unit=unit, copy=False)
+        )
+    return UnwrappedTimeOfArrival(toa)
+
+
+def unwrapped_time_of_arrival_minus_frame_pivot_time(
+    toa: UnwrappedTimeOfArrival, pivot_time: PivotTimeAtDetector
+) -> UnwrappedTimeOfArrivalMinusPivotTime:
+    """
+    Compute the time of arrival of the neutron at the detector, unwrapped at the pulse
+    period, minus the start time of the frame.
+    We subtract the start time of the frame so that we can use a modulo operation to
+    wrap the time of arrival at the frame period in the case of pulse-skipping.
+
+    Parameters
+    ----------
+    toa:
+        Time of arrival of the neutron at the detector, unwrapped at the pulse period.
+    pivot_time:
+        Pivot time at the detector, i.e., the time of the start of the frame at the
+        detector.
+    """
+    # Order of operation to preserve dimension order
+    return UnwrappedTimeOfArrivalMinusPivotTime(
+        -pivot_time.to(unit=elem_unit(toa), copy=False) + toa
+    )
+
+
+def time_of_arrival_minus_pivot_time_modulo_period(
+    toa_minus_pivot_time: UnwrappedTimeOfArrivalMinusPivotTime,
+    frame_period: FramePeriod,
+) -> TimeOfArrivalMinusPivotTimeModuloPeriod:
+    """
+    Compute the time of arrival of the neutron at the detector, unwrapped at the pulse
+    period, minus the start time of the frame, modulo the frame period.
+
+    Parameters
+    ----------
+    toa_minus_pivot_time:
+        Time of arrival of the neutron at the detector, unwrapped at the pulse period,
+        minus the start time of the frame.
+    frame_period:
+        Period of the frame, i.e., time between the start of two consecutive frames.
+    """
+    return TimeOfArrivalMinusPivotTimeModuloPeriod(
+        toa_minus_pivot_time
+        % frame_period.to(unit=elem_unit(toa_minus_pivot_time), copy=False)
+    )
+
+
+def time_of_arrival_folded_by_frame(
+    toa: TimeOfArrivalMinusPivotTimeModuloPeriod,
+    pivot_time: PivotTimeAtDetector,
+) -> FrameFoldedTimeOfArrival:
+    """
+    The time of arrival of the neutron at the detector, folded by the frame period.
+
+    Parameters
+    ----------
+    toa:
+        Time of arrival of the neutron at the detector, unwrapped at the pulse period,
+        minus the start time of the frame, modulo the frame period.
+    pivot_time:
+        Pivot time at the detector, i.e., the time of the start of the frame at the
+        detector.
+    """
+    return FrameFoldedTimeOfArrival(
+        toa + pivot_time.to(unit=elem_unit(toa), copy=False)
+    )
+
+
+def time_of_flight_data(
+    da: RawData,
+    lookup: MaskedTimeOfFlightLookupTable,
+    ltotal: Ltotal,
+    toas: FrameFoldedTimeOfArrival,
+) -> TofData:
+    """
+    Convert the time-of-arrival data to time-of-flight data using a lookup table.
+    The output data will have a time-of-flight coordinate.
+
+    Parameters
+    ----------
+    da:
+        Raw detector data loaded from a NeXus file, e.g., NXdetector containing
+        NXevent_data.
+    lookup:
+        Lookup table giving time-of-flight as a function of distance and time of
+        arrival.
+    ltotal:
+        Total length of the flight path from the source to the detector.
+    toas:
+        Time of arrival of the neutron at the detector, folded by the frame period.
+    """
+    from scipy.interpolate import RegularGridInterpolator
+
+    # TODO: to make use of multi-threading, we could write our own interpolator.
+    # This should be simple enough as we are making the bins linspace, so computing
+    # bin indices is fast.
+    f = RegularGridInterpolator(
+        (
+            lookup.coords["toa"].to(unit=elem_unit(toas), copy=False).values,
+            lookup.coords["distance"].to(unit=ltotal.unit, copy=False).values,
+        ),
+        lookup.data.to(unit=elem_unit(toas), copy=False).values.T,
+        method="linear",
+        bounds_error=False,
+    )
+
+    if da.bins is not None:
+        ltotal = sc.bins_like(toas, ltotal).bins.constituents["data"]
+        toas = toas.bins.constituents["data"]
+
+    tofs = sc.array(
+        dims=toas.dims, values=f((toas.values, ltotal.values)), unit=elem_unit(toas)
+    )
+
+    if da.bins is not None:
+        parts = da.bins.constituents
+        parts["data"] = tofs
+        out = da.bins.assign_coords(tof=_bins_no_validate(**parts))
+    else:
+        out = da.assign_coords(tof=tofs)
+
+    return TofData(out)
+
+
+def resample_tof_data(da: TofData) -> ResampledTofData:
+    """
+    Histogrammed data that has been converted to `tof` will typically have
+    unsorted bin edges (due to either wrapping of `time_of_flight` or wavelength
+    overlap between subframes).
+    This function re-histograms the data to ensure that the bin edges are sorted.
+    It makes use of the ``to_events`` helper which generates a number of events in each
+    bin with a uniform distribution. The new events are then histogrammed using a set of
+    sorted bin edges.
+
+    WARNING:
+    This function is highly experimental, has limitations and should be used with
+    caution. It is a workaround to the issue that rebinning data with unsorted bin
+    edges is not supported in scipp.
+    As such, this function is not part of the default set of providers, and needs to be
+    inserted manually into the workflow.
+
+    Parameters
+    ----------
+    da:
+        Histogrammed data with the time-of-flight coordinate.
+    """
+    dim = next(iter(set(da.dims) & {"time_of_flight", "tof"}))
+    events = to_events(da.rename_dims({dim: "tof"}), "event")
+
+    # Define a new bin width, close to the original bin width.
+    # TODO: this could be a workflow parameter
+    coord = da.coords["tof"]
+    bin_width = (coord[dim, 1:] - coord[dim, :-1]).nanmedian()
+    rehist = events.hist(tof=bin_width)
+    for key, var in da.coords.items():
+        if dim not in var.dims:
+            rehist.coords[key] = var
+    return ResampledTofData(rehist)
+
+
+def default_parameters() -> dict:
+    """
+    Default parameters of the time-of-flight workflow.
+    """
+    return {
+        PulsePeriod: 1.0 / sc.scalar(14.0, unit="Hz"),
+        PulseStride: 1,
+        PulseStrideOffset: 0,
+        DistanceResolution: sc.scalar(0.1, unit="m"),
+        TimeOfArrivalResolution: 500,
+        LookupTableRelativeErrorThreshold: 0.1,
+    }
+
+
+def providers() -> tuple[Callable]:
+    """
+    Providers of the time-of-flight workflow.
+    """
+    return (
+        compute_tof_lookup_table,
+        extract_ltotal,
+        find_fastest_neutron,
+        frame_period,
+        masked_tof_lookup_table,
+        pivot_time_at_detector,
+        time_of_arrival_folded_by_frame,
+        time_of_arrival_minus_pivot_time_modulo_period,
+        time_of_flight_data,
+        unwrapped_time_of_arrival,
+        unwrapped_time_of_arrival_minus_frame_pivot_time,
+    )
+
+
+class TofWorkflow:
+    """
+    Helper class to build a time-of-flight workflow and cache the expensive part of the
+    computation: running the simulation and building the lookup table.
+
+    Parameters
+    ----------
+    simulated_neutrons:
+        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.
+        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.
+    pulse_stride:
+        Stride of used pulses. Usually 1, but may be a small integer when
+        pulse-skipping.
+    pulse_stride_offset:
+        Integer offset of the first pulse in the stride (typically zero unless we
+        are using pulse-skipping and the events do not begin with the first pulse in
+        the stride).
+    distance_resolution:
+        Resolution 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.
+    toa_resolution:
+        Resolution of the time of arrival axis in the lookup table.
+        Can be an integer (number of bins) or a sc.Variable (bin width).
+    error_threshold:
+        Threshold for the variance of the projected time-of-flight above which
+        regions are masked.
+    """
+
+    def __init__(
+        self,
+        simulated_neutrons: SimulationResults,
+        ltotal_range: LtotalRange,
+        pulse_stride: PulseStride | None = None,
+        pulse_stride_offset: PulseStrideOffset | None = None,
+        distance_resolution: DistanceResolution | None = None,
+        toa_resolution: TimeOfArrivalResolution | None = None,
+        error_threshold: LookupTableRelativeErrorThreshold | None = None,
+    ):
+        import sciline as sl
+
+        self.pipeline = sl.Pipeline(providers())
+        self.pipeline[SimulationResults] = simulated_neutrons
+        self.pipeline[LtotalRange] = ltotal_range
+
+        params = default_parameters()
+        self.pipeline[PulsePeriod] = params[PulsePeriod]
+        self.pipeline[PulseStride] = pulse_stride or params[PulseStride]
+        self.pipeline[PulseStrideOffset] = (
+            pulse_stride_offset or params[PulseStrideOffset]
+        )
+        self.pipeline[DistanceResolution] = (
+            distance_resolution or params[DistanceResolution]
+        )
+        self.pipeline[TimeOfArrivalResolution] = (
+            toa_resolution or params[TimeOfArrivalResolution]
+        )
+        self.pipeline[LookupTableRelativeErrorThreshold] = (
+            error_threshold or params[LookupTableRelativeErrorThreshold]
+        )
+
+    def cache_results(
+        self,
+        results=(SimulationResults, MaskedTimeOfFlightLookupTable, FastestNeutron),
+    ) -> None:
+        """
+        Cache a list of (usually expensive to compute) intermediate results of the
+        time-of-flight workflow.
+
+        Parameters
+        ----------
+        results:
+            List of results to cache.
+        """
+        for t in results:
+            self.pipeline[t] = self.pipeline.compute(t)