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 +} diff --git a/src/tof/optimization.py b/src/tof/optimization.py new file mode 100644 index 0000000..88e5e35 --- /dev/null +++ b/src/tof/optimization.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: BSD-3-Clause +# 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. + +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 + +import numpy as np +import scipp as sc + +from .chopper import Chopper +from .utils import wavelength_to_inverse_speed + + +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 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, + ) + + +@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 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(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 + + +@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 _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) + + +def polygon_grid_overlap_mask( + 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. + This is used to select relevant regions of the source distribution where the + selection is made from the polygons generated by the FrameSequence. + + Parameters + ---------- + 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 + 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 + 3. It is a bit slow, can you improve the performance replacing boolean masks + with index slicing pruning? + + Returns: + mask (H-1, W-1) + """ + + H, W = X.shape + + # Assume structured grid + x = X[0, :] + y = Y[:, 0] + + dx = np.min(np.diff(x)) + dy = np.min(np.diff(y)) + + if tol is None: + tol = 1e-9 + 1e-6 * min(dx, dy) + + # 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) + + # --- Helpers --- + 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 + + 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) + + # 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)), + ] + + 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 + + union_mask |= corner_inside | vertex_inside | intersect_mask + + return union_mask diff --git a/src/tof/source.py b/src/tof/source.py index fc0489f..97c48a7 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -3,17 +3,54 @@ from __future__ import annotations import warnings +from collections.abc import Mapping from dataclasses import dataclass 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, Subframe, polygon_grid_overlap_mask from .utils import wavelength_to_speed TIME_UNIT = "us" WAV_UNIT = "angstrom" +T_DIM = "birth_time" +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: @@ -38,25 +75,57 @@ 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 _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), + } + ) + return facility_pulse + + def _make_pulses( neutrons: int, frequency: sc.Variable, 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, ): """ - 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 ---------- @@ -70,54 +139,18 @@ 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. """ - 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.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) - - 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() + 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) # 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,14 +162,17 @@ 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) - - # 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 + + widths = { + dim: _compute_grid_spacing(p.coords[dim], dim) + .broadcast(sizes=p.sizes) + .flatten(to='x') + for dim in (T_DIM, W_DIM) + } + + # 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 = [] @@ -152,20 +188,47 @@ def _make_pulses( f"Sum of probabilities is {p_sum.value}" ) p_flat /= p_sum + 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 // 3) 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) - mask = ( - (t >= tmin.value) - & (t <= tmax.value) - & (w >= wmin.value) - & (w <= wmax.value) + t = p_flat.coords[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.uniform(-0.5, 0.5, size=size) * widths[W_DIM].values[inds] ) - times.append(t[mask]) - wavs.append(w[mask]) - n += mask.sum() + + # 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: + 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) + sel |= path.contains_points(points) + + # 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( @@ -192,6 +255,165 @@ def _make_pulses( } +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] | Mapping[str, Chopper] | None = None, +) -> 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"] + + 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.values() if hasattr(choppers, "items") else 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 = polygon_grid_overlap_mask( + [ + np.column_stack( + [ + poly.time.to(unit=TIME_UNIT).values, + poly.wavelength.to(unit=WAV_UNIT).values, + ] + ) + for poly in polygons + ], + X, + Y, + ) + + prob = p.copy(deep=True) + prob.values = np.where(mask, p.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: """ A class that represents a source of neutrons. @@ -216,8 +438,18 @@ 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. + 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__( @@ -225,44 +457,45 @@ 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, + 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) self._pulses = int(pulses) self._data = None - 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) - - self._frequency = facility_pulse.coords["frequency"] - self._distance = facility_pulse.coords["distance"] - pulse_params = _make_pulses( - neutrons=self._neutrons, - p=facility_pulse, - frequency=self._frequency, - pulses=self._pulses, - wmin=wmin, - wmax=wmax, - 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), - }, - ) + 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=self._neutrons, + pulses=self._pulses, + frequency=self._frequency, + seed=self._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: @@ -285,6 +518,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: """ @@ -299,6 +539,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: """ @@ -310,6 +557,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"], } ) @@ -384,15 +632,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 ---------- @@ -414,40 +669,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) + source = cls(facility=None, neutrons=neutrons, pulses=pulses, seed=seed) - 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, + neutrons=source._neutrons, + pulses=source._pulses, frequency=source._frequency, - pulses=pulses, - 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), - }, + seed=source._seed, + distance=source._distance, + wmin=wmin, + wmax=wmax, + tmin=tmin, + tmax=tmax, + optimize_for=optimize_for, ) return source @@ -531,28 +790,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") 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. diff --git a/tests/optimization_test.py b/tests/optimization_test.py new file mode 100644 index 0000000..50563e8 --- /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 ("WFMC_1", "WFMC_2") + } + 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 b173400..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') @@ -148,7 +164,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(