From 9f527aac8ff380a7061e2a52665c1aa3a60725b1 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Fri, 20 Feb 2026 21:24:12 +0100 Subject: [PATCH 01/22] start modifying the way model runs --- src/tof/model.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/tof/model.py b/src/tof/model.py index 36c241f..7234df7 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -274,9 +274,10 @@ def run(self): "itself. Please check the distances of the components." ) - birth_time = self.source.data.coords['birth_time'] - speed = self.source.data.coords['speed'] - initial_mask = sc.ones(sizes=birth_time.sizes, unit=None, dtype=bool) + # birth_time = self.source.data.coords['birth_time'] + # speed = self.source.data.coords['speed'] + neutrons = self.source.data + initial_mask = sc.ones(sizes=neutrons.sizes, unit=None, dtype=bool) result_choppers = {} result_detectors = {} From 6f49224e206f1dd2086ceee64565a0f8be600efa Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Fri, 20 Feb 2026 23:39:39 +0100 Subject: [PATCH 02/22] update model to work with components --- src/tof/chopper.py | 111 +++++++++++---- src/tof/component.py | 26 ++++ src/tof/detector.py | 40 +++++- src/tof/model.py | 317 +++++++++++++++++++++++++------------------ 4 files changed, 324 insertions(+), 170 deletions(-) create mode 100644 src/tof/component.py diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 939a4c4..50a255c 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -27,6 +27,16 @@ class Direction(Enum): AntiClockwise = Direction.ANTICLOCKWISE +def _array_or_none(container: dict, key: str) -> sc.Variable | None: + return ( + sc.array( + dims=["cutout"], values=container[key]["value"], unit=container[key]["unit"] + ) + if key in container + else None + ) + + class Chopper: """ A chopper is a rotating device with cutouts that blocks the beam at certain times. @@ -178,6 +188,18 @@ def __repr__(self) -> str: f"direction={self.direction.name}, cutouts={len(self.open)})" ) + def __eq__(self, other: object) -> bool: + if not isinstance(other, Chopper): + return NotImplemented + if self.name != other.name: + return False + if self.direction != other.direction: + return False + return all( + sc.identical(getattr(self, field), getattr(other, field)) + for field in ('frequency', 'distance', 'phase', 'open', 'close') + ) + def as_dict(self) -> dict: """ Return the chopper as a dictionary. @@ -192,6 +214,31 @@ def as_dict(self) -> dict: 'direction': self.direction, } + @classmethod + def from_json(cls, name: str, params: dict) -> Chopper: + direction = params["direction"].lower() + if direction == "clockwise": + _dir = Clockwise + elif any(x in direction for x in ("anti", "counter")): + _dir = AntiClockwise + else: + raise ValueError( + f"Chopper direction must be 'clockwise' or 'anti-clockwise', got " + f"'{params['direction']}' for component {name}." + ) + return cls( + frequency=params["frequency"]["value"] + * sc.Unit(params["frequency"]["unit"]), + direction=_dir, + open=_array_or_none(params, "open"), + close=_array_or_none(params, "close"), + centers=_array_or_none(params, "centers"), + widths=_array_or_none(params, "widths"), + phase=params["phase"]["value"] * sc.Unit(params["phase"]["unit"]), + distance=params["distance"]["value"] * sc.Unit(params["distance"]["unit"]), + name=name, + ) + def as_json(self) -> dict: """ Return the chopper as a JSON-serializable dictionary. @@ -212,18 +259,6 @@ def as_json(self) -> dict: ) return out - def __eq__(self, other: object) -> bool: - if not isinstance(other, Chopper): - return NotImplemented - if self.name != other.name: - return False - if self.direction != other.direction: - return False - return all( - sc.identical(getattr(self, field), getattr(other, field)) - for field in ('frequency', 'distance', 'phase', 'open', 'close') - ) - @classmethod def from_diskchopper( cls, disk_chopper: DiskChopper, name: str | None = None @@ -273,6 +308,27 @@ def from_diskchopper( name=name, ) + def to_diskchopper(self) -> DiskChopper: + """ + Export the chopper as a scippneutron DiskChopper. + """ + from scippneutron.chopper import DiskChopper + + frequency = ( + self.frequency if self.direction == AntiClockwise else -self.frequency + ) + phase = self.phase if self.direction == AntiClockwise else -self.phase + return DiskChopper( + frequency=frequency, + beam_position=sc.scalar(0.0, unit='deg'), + slit_begin=self.open, + slit_end=self.close, + phase=phase, + axle_position=sc.vector( + value=[0.0, 0.0, self.distance.value], unit=self.distance.unit + ), + ) + @classmethod def from_nexus(cls, nexus_chopper, name: str | None = None) -> Chopper: """ @@ -321,26 +377,21 @@ def from_nexus(cls, nexus_chopper, name: str | None = None) -> Chopper: name=name, ) - def to_diskchopper(self) -> DiskChopper: + def apply(self, neutrons: sc.DataArray) -> sc.DataArray: """ - Export the chopper as a scippneutron DiskChopper. + Apply the chopper's properties to the given data array. """ - from scippneutron.chopper import DiskChopper - - frequency = ( - self.frequency if self.direction == AntiClockwise else -self.frequency - ) - phase = self.phase if self.direction == AntiClockwise else -self.phase - return DiskChopper( - frequency=frequency, - beam_position=sc.scalar(0.0, unit='deg'), - slit_begin=self.open, - slit_end=self.close, - phase=phase, - axle_position=sc.vector( - value=[0.0, 0.0, self.distance.value], unit=self.distance.unit - ), - ) + # Apply the chopper's open/close times to the data + m = sc.zeros(sizes=t.sizes, unit=None, dtype=bool) + to, tc = c.open_close_times(time_limit=time_limit) + results[c.name].update({'open_times': to, 'close_times': tc}) + for i in range(len(to)): + m |= (t > to[i]) & (t < tc[i]) + combined = initial_mask & m + data_at_comp.masks['blocked_by_others'] = ~initial_mask + data_at_comp.masks['blocked_by_me'] = ~m & initial_mask + + return data @dataclass(frozen=True) diff --git a/src/tof/component.py b/src/tof/component.py new file mode 100644 index 0000000..8c7cdcf --- /dev/null +++ b/src/tof/component.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2026 Scipp contributors (https://github.com/scipp) + +from abc import abstractmethod + +import scipp as sc + + +class Component: + kind: str + + @abstractmethod + def apply(self, neutrons: sc.DataArray) -> sc.DataArray: + """ + Apply the component to the given neutrons. + + Parameters + ---------- + neutrons: + The neutrons to which the component will be applied. + + Returns + ------- + The modified neutrons. + """ + raise NotImplementedError diff --git a/src/tof/detector.py b/src/tof/detector.py index 9dc8075..c2d24a4 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -6,11 +6,12 @@ import scipp as sc +from .component import Component from .reading import ComponentReading from .utils import var_to_dict -class Detector: +class Detector(Component): """ A detector component does not block any neutrons, it sees all neutrons passing through it. @@ -26,16 +27,34 @@ class Detector: def __init__(self, distance: sc.Variable, name: str): self.distance = distance.to(dtype=float, copy=False) self.name = name + self.kind = "detector" def __repr__(self) -> str: return f"Detector(name={self.name}, distance={self.distance:c})" + def __eq__(self, other: object) -> bool: + if not isinstance(other, Detector): + return NotImplemented + return self.name == other.name and sc.identical(self.distance, other.distance) + def as_dict(self) -> dict: """ Return the detector as a dictionary. """ return {'distance': self.distance, 'name': self.name} + @classmethod + def from_json(cls, name: str, params: dict) -> Detector: + """ + Create a detector from a JSON-serializable dictionary. + """ + return cls( + distance=sc.scalar( + params["distance"]["value"], unit=params["distance"]["unit"] + ), + name=name, + ) + def as_json(self) -> dict: """ Return the detector as a JSON-serializable dictionary. @@ -48,10 +67,21 @@ def as_json(self) -> dict: 'name': self.name, } - def __eq__(self, other: object) -> bool: - if not isinstance(other, Detector): - return NotImplemented - return self.name == other.name and sc.identical(self.distance, other.distance) + def apply(self, neutrons: sc.DataArray) -> sc.DataArray: + """ + Apply the detector to the given neutrons. + + Parameters + ---------- + neutrons: + The neutrons to which the detector will be applied. + + Returns + ------- + The modified neutrons. + """ + # A detector does not modify the neutrons, it simply records them. + return neutrons @dataclass(frozen=True) diff --git a/src/tof/model.py b/src/tof/model.py index 7234df7..30fd285 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -4,11 +4,14 @@ from __future__ import annotations import warnings +from collections import defaultdict from itertools import chain import scipp as sc +# from tof.component import Component from .chopper import AntiClockwise, Chopper, Clockwise +from .component import Component from .detector import Detector from .result import Result from .source import Source @@ -16,14 +19,14 @@ ComponentType = Chopper | Detector -def _array_or_none(container: dict, key: str) -> sc.Variable | None: - return ( - sc.array( - dims=["cutout"], values=container[key]["value"], unit=container[key]["unit"] - ) - if key in container - else None - ) +# def _array_or_none(results: dict, key: str) -> sc.Variable | None: +# return ( +# sc.array( +# dims=["cutout"], values=results[key]["value"], unit=results[key]["unit"] +# ) +# if key in results +# else None +# ) def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]]: @@ -42,85 +45,118 @@ def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]] type, see the documentation of the :class:`Chopper` and :class:`Detector` classes for details. """ - choppers = [] - detectors = [] + components = [] + # detectors = [] + mapping = {"chopper": Chopper, "detector": Detector} for name, comp in instrument.items(): - if comp["type"] == "chopper": - direction = comp["direction"].lower() - if direction == "clockwise": - _dir = Clockwise - elif any(x in direction for x in ("anti", "counter")): - _dir = AntiClockwise - else: - raise ValueError( - f"Chopper direction must be 'clockwise' or 'anti-clockwise', got " - f"'{comp['direction']}' for component {name}." - ) - choppers.append( - Chopper( - frequency=comp["frequency"]["value"] - * sc.Unit(comp["frequency"]["unit"]), - direction=_dir, - open=_array_or_none(comp, "open"), - close=_array_or_none(comp, "close"), - centers=_array_or_none(comp, "centers"), - widths=_array_or_none(comp, "widths"), - phase=comp["phase"]["value"] * sc.Unit(comp["phase"]["unit"]), - distance=comp["distance"]["value"] - * sc.Unit(comp["distance"]["unit"]), - name=name, - ) - ) - elif comp["type"] == "detector": - detectors.append( - Detector( - distance=comp["distance"]["value"] - * sc.Unit(comp["distance"]["unit"]), - name=name, - ) - ) - elif comp["type"] == "source": - continue - else: + if comp["type"] not in mapping: raise ValueError( f"Unknown component type: {comp['type']} for component {name}. " - "Supported types are 'chopper', 'detector', and 'source'." ) - return {"choppers": choppers, "detectors": detectors} + components.append(mapping[comp["type"]].from_json(name=name, params=comp)) + # component_class = comp_mapping[comp["type"]] + # if comp["type"] == "chopper": + # direction = comp["direction"].lower() + # if direction == "clockwise": + # _dir = Clockwise + # elif any(x in direction for x in ("anti", "counter")): + # _dir = AntiClockwise + # else: + # raise ValueError( + # f"Chopper direction must be 'clockwise' or 'anti-clockwise', got " + # f"'{comp['direction']}' for component {name}." + # ) + # choppers.append( + # Chopper( + # frequency=comp["frequency"]["value"] + # * sc.Unit(comp["frequency"]["unit"]), + # direction=_dir, + # open=_array_or_none(comp, "open"), + # close=_array_or_none(comp, "close"), + # centers=_array_or_none(comp, "centers"), + # widths=_array_or_none(comp, "widths"), + # phase=comp["phase"]["value"] * sc.Unit(comp["phase"]["unit"]), + # distance=comp["distance"]["value"] + # * sc.Unit(comp["distance"]["unit"]), + # name=name, + # ) + # ) + # elif comp["type"] == "detector": + # detectors.append( + # Detector( + # distance=comp["distance"]["value"] + # * sc.Unit(comp["distance"]["unit"]), + # name=name, + # ) + # ) + # elif comp["type"] == "source": + # continue + # else: + # raise ValueError( + # f"Unknown component type: {comp['type']} for component {name}. " + # "Supported types are 'chopper', 'detector', and 'source'." + # ) + return components class Model: """ A class that represents a neutron instrument. - It is defined by a list of choppers, a list of detectors, and a source. + It is defined by a source and a list of components (choppers, detectors, etc.). Parameters ---------- - choppers: - A list of choppers. - detectors: - A list of detectors. source: A source of neutrons. + components: + A list of components. + choppers: + A list of choppers. This is kept for backwards-compatibility; new code + should use the `components` parameter instead. + detectors: + A list of detectors. This is kept for backwards-compatibility; new code + should use the `components` parameter instead. """ def __init__( self, source: Source | None = None, + components: list[Component] | tuple[Component, ...] | None = None, choppers: list[Chopper] | tuple[Chopper, ...] | None = None, detectors: list[Detector] | tuple[Detector, ...] | None = None, ): - self.choppers = {} - self.detectors = {} + # self.choppers = {} + # self.detectors = {} self.source = source - for components, kind in ((choppers, Chopper), (detectors, Detector)): - for c in components or (): - if not isinstance(c, kind): - raise TypeError( - f"Beamline components: expected {kind.__name__} instance, " - f"got {type(c)}." - ) - self.add(c) + # for components, kind in ((choppers, Chopper), (detectors, Detector)): + for comp in chain((choppers or ()), (detectors or ()), (components or ())): + self.add(comp) + + # @property + # def choppers(self) -> dict[str, Chopper]: + # """ + # Return a dictionary of all choppers in the model. + # This is meant for retro-compatibility with older versions of the code, and will + # be removed in a future version. + # """ + # return { + # name: comp + # for name, comp in self.components.items() + # if comp.kind == "chopper" + # } + + # @property + # def detectors(self) -> dict[str, Detector]: + # """ + # Return a dictionary of all detectors in the model. + # This is meant for retro-compatibility with older versions of the code, and will + # be removed in a future version. + # """ + # return { + # name: comp + # for name, comp in self.components.items() + # if comp.kind == "detector" + # } @classmethod def from_json(cls, filename: str) -> Model: @@ -173,13 +209,11 @@ def as_json(self) -> dict: ) else: instrument_dict['source'] = self.source.as_json() - for ch in self.choppers.values(): - instrument_dict[ch.name] = ch.as_json() - for det in self.detectors.values(): - instrument_dict[det.name] = det.as_json() + for comp in self.components.values(): + instrument_dict[comp.name] = comp.as_json() return instrument_dict - def to_json(self, filename: str): + def to_json(self, filename: str) -> None: """ Save the model to a JSON file. If the source is not from a facility, it is not included in the output. @@ -196,7 +230,7 @@ def to_json(self, filename: str): with open(filename, 'w') as f: json.dump(self.as_json(), f, indent=2) - def add(self, component: Chopper | Detector): + def add(self, component: Component) -> None: """ Add a component to the instrument. Component names must be unique across choppers and detectors. @@ -208,21 +242,20 @@ def add(self, component: Chopper | Detector): component: A chopper or detector. """ - if not isinstance(component, (Chopper | Detector)): - raise TypeError( - f"Cannot add component of type {type(component)} to the model. " - "Only Chopper and Detector instances are allowed." - ) + # if not isinstance(component, (Chopper | Detector)): + # raise TypeError( + # f"Cannot add component of type {type(component)} to the model. " + # "Only Chopper and Detector instances are allowed." + # ) # Note that the name "source" is reserved for the source. - if component.name in chain(self.choppers, self.detectors, ("source",)): + if component.name in (*self.components, "source"): raise KeyError( f"Component with name {component.name} already exists. " "If you wish to replace/update an existing component, use " - "``model.choppers['name'] = new_chopper`` or " - "``model.detectors['name'] = new_detector``." + "``model.components['name'] = new_component``." ) - container = self.choppers if isinstance(component, Chopper) else self.detectors - container[component.name] = component + # results = self.choppers if isinstance(component, Chopper) else self.detectors + self.components[component.name] = component def remove(self, name: str): """ @@ -233,25 +266,26 @@ def remove(self, name: str): name: The name of the component to remove. """ - if name in self.choppers: - del self.choppers[name] - elif name in self.detectors: - del self.detectors[name] - else: - raise KeyError(f"No component with name {name} was found.") - - def __iter__(self): - return chain(self.choppers, self.detectors) - - def __getitem__(self, name) -> Chopper | Detector: - if name not in self: - raise KeyError(f"No component with name {name} was found.") - return self.choppers[name] if name in self.choppers else self.detectors[name] - - def __delitem__(self, name): - self.remove(name) - - def run(self): + del self.components[name] + # if name in self.choppers: + # del self.choppers[name] + # elif name in self.detectors: + # del self.detectors[name] + # else: + # raise KeyError(f"No component with name {name} was found.") + + # def __iter__(self): + # return chain(self.choppers, self.detectors) + + # def __getitem__(self, name) -> Chopper | Detector: + # if name not in self: + # raise KeyError(f"No component with name {name} was found.") + # return self.choppers[name] if name in self.choppers else self.detectors[name] + + # def __delitem__(self, name) -> None: + # self.remove(name) + + def run(self) -> Result: """ Run the simulation. """ @@ -260,10 +294,7 @@ def run(self): "No source has been defined for this model. Please add a source using " "`model.source = Source(...)` before running the simulation." ) - components = sorted( - chain(self.choppers.values(), self.detectors.values()), - key=lambda c: c.distance.value, - ) + components = sorted(self.components.values(), key=lambda c: c.distance.value) if len(components) == 0: raise ValueError("Cannot run model: no components have been defined.") @@ -276,43 +307,59 @@ def run(self): # birth_time = self.source.data.coords['birth_time'] # speed = self.source.data.coords['speed'] - neutrons = self.source.data - initial_mask = sc.ones(sizes=neutrons.sizes, unit=None, dtype=bool) + neutrons = self.source.data.copy(deep=False) + neutrons.masks['blocked_by_others'] = sc.zeros( + sizes=neutrons.sizes, unit=None, dtype=bool + ) + neutrons.coords.update( + distance=self.source.distance, toa=self.coords['birth_time'] + ) - result_choppers = {} - result_detectors = {} + time_unit = neutrons.coords['birth_time'].unit + + results = {} + # result_choppers = {} + # result_detectors = {} time_limit = ( - birth_time - + ((components[-1].distance - self.source.distance) / speed).to( - unit=birth_time.unit - ) + neutrons.coords['birth_time'] + + ( + (components[-1].distance - self.source.distance) + / neutrons.coords['speed'] + ).to(unit=time_unit) ).max() - for c in components: - container = result_detectors if isinstance(c, Detector) else result_choppers - container[c.name] = c.as_dict() - container[c.name]['data'] = self.source.data.copy(deep=False) - tof = ((c.distance - self.source.distance) / speed).to( - unit=birth_time.unit, copy=False + for comp in components: + # results = result_detectors if isinstance(c, Detector) else result_choppers + results[comp.name] = comp.as_dict() + data_at_comp = neutrons.copy(deep=False) + # tof = ((c.distance - self.source.distance) / neutrons.coords['speed']).to( + # unit=time_unit, copy=False + # ) + # t = birth_time + tof + toa = neutrons.coords['toa'] + ( + (comp.distance - neutrons.coords['distance']) / neutrons.coords['speed'] + ).to(unit=time_unit, copy=False) + data_at_comp.coords['toa'] = toa + data_at_comp.coords['eto'] = toa % (1 / self.source.frequency).to( + unit=time_unit, copy=False ) - t = birth_time + tof - container[c.name]['data'].coords['toa'] = t - container[c.name]['data'].coords['eto'] = t % ( - 1 / self.source.frequency - ).to(unit=t.unit, copy=False) - container[c.name]['data'].coords['distance'] = c.distance - container[c.name]['data'].coords['tof'] = tof - if isinstance(c, Detector): - container[c.name]['data'].masks['blocked_by_others'] = ~initial_mask - continue - m = sc.zeros(sizes=t.sizes, unit=None, dtype=bool) - to, tc = c.open_close_times(time_limit=time_limit) - container[c.name].update({'open_times': to, 'close_times': tc}) - for i in range(len(to)): - m |= (t > to[i]) & (t < tc[i]) - combined = initial_mask & m - container[c.name]['data'].masks['blocked_by_others'] = ~initial_mask - container[c.name]['data'].masks['blocked_by_me'] = ~m & initial_mask + data_at_comp.coords['distance'] = comp.distance + + data_at_comp = comp.apply(data_at_comp) + + # # data_at_comp.coords['tof'] = tof + # if isinstance(c, Detector): + # data_at_comp.masks['blocked_by_others'] = ~initial_mask + # continue + # m = sc.zeros(sizes=t.sizes, unit=None, dtype=bool) + # to, tc = c.open_close_times(time_limit=time_limit) + # results[c.name].update({'open_times': to, 'close_times': tc}) + # for i in range(len(to)): + # m |= (t > to[i]) & (t < tc[i]) + # combined = initial_mask & m + # data_at_comp.masks['blocked_by_others'] = ~initial_mask + # data_at_comp.masks['blocked_by_me'] = ~m & initial_mask initial_mask = combined + neutrons = data_at_comp return Result( source=self.source, choppers=result_choppers, detectors=result_detectors From b43a3e3f67ea8e5a766c2043050c5920f1a83b2f Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Sat, 21 Feb 2026 10:43:12 +0100 Subject: [PATCH 03/22] move more logic into components --- src/tof/chopper.py | 123 +++++++++++++++++++++++----------------- src/tof/component.py | 130 +++++++++++++++++++++++++++++++++++++++++++ src/tof/detector.py | 77 +++++++++++++------------ src/tof/model.py | 57 ++++++++++--------- src/tof/result.py | 54 +++++++++--------- 5 files changed, 300 insertions(+), 141 deletions(-) diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 50a255c..38dc258 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -37,6 +37,46 @@ def _array_or_none(container: dict, key: str) -> sc.Variable | None: ) +@dataclass(frozen=True) +class ChopperReading(ComponentReading): + """ + Read-only container for the neutrons that reach the chopper. + """ + + distance: sc.Variable + name: str + frequency: sc.Variable + open: sc.Variable + close: sc.Variable + phase: sc.Variable + open_times: sc.Variable + close_times: sc.Variable + data: sc.DataArray + + def _repr_stats(self) -> str: + return ( + f"visible={int(self.data.sum().value)}, " + f"blocked={int(self.data.masks['blocked_by_me'].sum().value)}" + ) + + def __repr__(self) -> str: + return f"""ChopperReading: '{self.name}' + distance: {self.distance:c} + frequency: {self.frequency:c} + phase: {self.phase:c} + cutouts: {len(self.open)} + neutrons: {self._repr_stats()} +""" + + def __str__(self) -> str: + return self.__repr__() + + def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ChopperReading: + if isinstance(val, int): + val = ('pulse', val) + return replace(self, data=self.data[val]) + + class Chopper: """ A chopper is a rotating device with cutouts that blocks the beam at certain times. @@ -377,58 +417,41 @@ def from_nexus(cls, nexus_chopper, name: str | None = None) -> Chopper: name=name, ) - def apply(self, neutrons: sc.DataArray) -> sc.DataArray: + def make_reading( + self, neutrons: sc.DataGroup, time_limit: sc.Variable + ) -> ChopperReading: """ - Apply the chopper's properties to the given data array. + Create a ChopperReading from the given neutrons that have been processed by this + chopper. """ - # Apply the chopper's open/close times to the data - m = sc.zeros(sizes=t.sizes, unit=None, dtype=bool) - to, tc = c.open_close_times(time_limit=time_limit) - results[c.name].update({'open_times': to, 'close_times': tc}) - for i in range(len(to)): - m |= (t > to[i]) & (t < tc[i]) - combined = initial_mask & m - data_at_comp.masks['blocked_by_others'] = ~initial_mask - data_at_comp.masks['blocked_by_me'] = ~m & initial_mask - - return data - - -@dataclass(frozen=True) -class ChopperReading(ComponentReading): - """ - Read-only container for the neutrons that reach the chopper. - """ - - distance: sc.Variable - name: str - frequency: sc.Variable - open: sc.Variable - close: sc.Variable - phase: sc.Variable - open_times: sc.Variable - close_times: sc.Variable - data: sc.DataArray - - def _repr_stats(self) -> str: - return ( - f"visible={int(self.data.sum().value)}, " - f"blocked={int(self.data.masks['blocked_by_me'].sum().value)}" + to, tc = self.open_close_times(time_limit=time_limit) + return ChopperReading( + distance=self.distance, + name=self.name, + frequency=self.frequency, + phase=self.phase, + open=self.open, + close=self.close, + open_times=to, + close_times=tc, + data=neutrons, ) - def __repr__(self) -> str: - return f"""ChopperReading: '{self.name}' - distance: {self.distance:c} - frequency: {self.frequency:c} - phase: {self.phase:c} - cutouts: {len(self.open)} - neutrons: {self._repr_stats()} -""" - - def __str__(self) -> str: - return self.__repr__() + def apply( + self, neutrons: sc.DataGroup, time_limit: sc.Variable + ) -> tuple[sc.DataGroup, ChopperReading]: + """ + Apply the effect of the chopper to the given neutrons. + """ + # Apply the chopper's open/close times to the data + m = sc.zeros(sizes=neutrons.sizes, unit=None, dtype=bool) + to, tc = self.open_close_times(time_limit=time_limit) + # neutrons.update({'open_times': to, 'close_times': tc}) + for i in range(len(to)): + m |= (neutrons['toa'] > to[i]) & (neutrons['toa'] < tc[i]) + # combined = initial_mask & m + # data_at_comp.masks['blocked_by_others'] = ~initial_mask + neutrons['blocked_by_me'] = (~m) & (~neutrons['blocked_by_others']) - def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ChopperReading: - if isinstance(val, int): - val = ('pulse', val) - return replace(self, data=self.data[val]) + reading = self.make_reading(neutrons, time_limit=time_limit) + return neutrons, reading diff --git a/src/tof/component.py b/src/tof/component.py index 8c7cdcf..a2f4b35 100644 --- a/src/tof/component.py +++ b/src/tof/component.py @@ -1,10 +1,140 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2026 Scipp contributors (https://github.com/scipp) +from __future__ import annotations + from abc import abstractmethod +from dataclasses import dataclass +import plopp as pp import scipp as sc +from .utils import Plot + + +@dataclass(frozen=True) +class ReadingField: + data: sc.DataArray + dim: str + + def plot(self, bins: int = 300, **kwargs): + by_pulse = sc.collapse(self.data, keep="event") + to_plot = {} + color = {} + for key, da in by_pulse.items(): + sel = da[~da.masks["blocked_by_others"]] + if sel.size == 0: + continue + to_plot[key] = sel.hist({self.dim: bins}) + if "blocked_by_me" in self.data.masks: + name = f"blocked-{key}" + to_plot[name] = ( + da[da.masks["blocked_by_me"]] + .drop_masks(list(da.masks.keys())) + .hist({self.dim: to_plot[key].coords[self.dim]}) + ) + color[name] = "gray" + if not to_plot: + raise RuntimeError("Nothing to plot.") + return pp.plot(to_plot, **{**{"color": color}, **kwargs}) + + def min(self): + da = self.data.copy(deep=False) + da.data = da.coords[self.dim] + return da.min().data + + def max(self): + da = self.data.copy(deep=False) + da.data = da.coords[self.dim] + return da.max().data + + def __repr__(self) -> str: + da = self.data.copy(deep=False) + da.data = da.coords[self.dim] + return ( + f"{self.dim}: min={da.min().data:c}, max={da.max().data:c}, " + f"events={int(self.data.sum().value)}" + ) + + def __str__(self) -> str: + return self.__repr__() + + def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ReadingField: + if isinstance(val, int): + val = ('pulse', val) + return self.__class__(data=self.data[val], dim=self.dim) + + +def _make_reading_field(da: sc.DataArray, dim: str) -> ReadingField: + return ReadingField( + data=sc.DataArray( + data=da.data, + coords={dim: da.coords[dim]}, + masks=da.masks, + ), + dim=dim, + ) + + +class ComponentReading: + """ + Data reading for a component placed in the beam path. The reading will have a + record of the arrival times and wavelengths of the neutrons that passed through it. + """ + + @property + def toa(self) -> ReadingField: + """ + Time of arrival of the neutrons at the component. + """ + return _make_reading_field(self.data, dim="toa") + + @property + def tof(self) -> ReadingField: + """ + Time of flight of the neutrons from source to the component. + """ + return _make_reading_field(self.data, dim="tof") + + @property + def eto(self) -> ReadingField: + """ + Event time offset of the neutrons at the component (= toa modulo pulse period). + """ + return _make_reading_field(self.data, dim="eto") + + @property + def wavelength(self) -> ReadingField: + """ + Wavelength of the neutrons at the component. + """ + return _make_reading_field(self.data, dim="wavelength") + + @property + def birth_time(self) -> ReadingField: + """ + Birth time of the neutrons at the source. + """ + return _make_reading_field(self.data, dim="birth_time") + + @property + def speed(self) -> ReadingField: + """ + Speed of the neutrons at the component. + """ + return _make_reading_field(self.data, dim="speed") + + def plot(self, bins: int = 300) -> Plot: + """ + Plot both the toa and wavelength data side by side. + + Parameters + ---------- + bins: + Number of bins to use for histogramming the neutrons. + """ + return self.toa.plot(bins=bins) + self.wavelength.plot(bins=bins) + class Component: kind: str diff --git a/src/tof/detector.py b/src/tof/detector.py index c2d24a4..132a360 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -11,6 +11,36 @@ from .utils import var_to_dict +@dataclass(frozen=True) +class DetectorReading(ComponentReading): + """ + Read-only container for the neutrons that reach the detector. + """ + + distance: sc.Variable + name: str + data: sc.DataArray + + def _repr_stats(self) -> str: + return f"visible={int(self.data.sum().value)}" + + def __repr__(self) -> str: + return f"""DetectorReading: '{self.name}' + distance: {self.distance:c} + neutrons: {self._repr_stats()} +""" + + def __str__(self) -> str: + return self.__repr__() + + def __getitem__( + self, val: int | slice | tuple[str, int | slice] + ) -> DetectorReading: + if isinstance(val, int): + val = ('pulse', val) + return replace(self, data=self.data[val]) + + class Detector(Component): """ A detector component does not block any neutrons, it sees all neutrons passing @@ -67,48 +97,21 @@ def as_json(self) -> dict: 'name': self.name, } - def apply(self, neutrons: sc.DataArray) -> sc.DataArray: + def make_reading(self, neutrons: sc.DataGroup) -> DetectorReading: + return DetectorReading(distance=self.distance, name=self.name, data=neutrons) + + def apply(self, neutrons: sc.DataGroup, time_limit: sc.Variable) -> sc.DataGroup: """ Apply the detector to the given neutrons. + A detector does not modify the neutrons, it simply records them without + blocking any. Parameters ---------- neutrons: The neutrons to which the detector will be applied. - - Returns - ------- - The modified neutrons. + time_limit: + The time limit for the neutrons to be considered as reaching the detector. """ - # A detector does not modify the neutrons, it simply records them. - return neutrons - - -@dataclass(frozen=True) -class DetectorReading(ComponentReading): - """ - Read-only container for the neutrons that reach the detector. - """ - - distance: sc.Variable - name: str - data: sc.DataArray - - def _repr_stats(self) -> str: - return f"visible={int(self.data.sum().value)}" - - def __repr__(self) -> str: - return f"""DetectorReading: '{self.name}' - distance: {self.distance:c} - neutrons: {self._repr_stats()} -""" - - def __str__(self) -> str: - return self.__repr__() - - def __getitem__( - self, val: int | slice | tuple[str, int | slice] - ) -> DetectorReading: - if isinstance(val, int): - val = ('pulse', val) - return replace(self, data=self.data[val]) + # neutrons.pop("blocked_by_me", None) + return neutrons, self.make_reading(neutrons) diff --git a/src/tof/model.py b/src/tof/model.py index 30fd285..c696916 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -307,44 +307,51 @@ def run(self) -> Result: # birth_time = self.source.data.coords['birth_time'] # speed = self.source.data.coords['speed'] - neutrons = self.source.data.copy(deep=False) - neutrons.masks['blocked_by_others'] = sc.zeros( - sizes=neutrons.sizes, unit=None, dtype=bool - ) - neutrons.coords.update( - distance=self.source.distance, toa=self.coords['birth_time'] + neutrons = sc.DataGroup(self.source.data.coords) + neutrons.update( + blocked_by_others=sc.zeros(sizes=neutrons.sizes, unit=None, dtype=bool), + distance=self.source.distance, + toa=neutrons['birth_time'], ) - time_unit = neutrons.coords['birth_time'].unit + time_unit = neutrons['birth_time'].unit - results = {} + readings = {} # result_choppers = {} # result_detectors = {} time_limit = ( - neutrons.coords['birth_time'] - + ( - (components[-1].distance - self.source.distance) - / neutrons.coords['speed'] - ).to(unit=time_unit) + neutrons['birth_time'] + + ((components[-1].distance - self.source.distance) / neutrons['speed']).to( + unit=time_unit + ) ).max() for comp in components: # results = result_detectors if isinstance(c, Detector) else result_choppers - results[comp.name] = comp.as_dict() - data_at_comp = neutrons.copy(deep=False) + # results[comp.name] = comp.as_dict() + neutrons = neutrons.copy(deep=False) # tof = ((c.distance - self.source.distance) / neutrons.coords['speed']).to( # unit=time_unit, copy=False # ) # t = birth_time + tof - toa = neutrons.coords['toa'] + ( - (comp.distance - neutrons.coords['distance']) / neutrons.coords['speed'] + toa = neutrons['toa'] + ( + (comp.distance - neutrons['distance']) / neutrons['speed'] ).to(unit=time_unit, copy=False) - data_at_comp.coords['toa'] = toa - data_at_comp.coords['eto'] = toa % (1 / self.source.frequency).to( + neutrons['toa'] = toa + neutrons['eto'] = toa % (1 / self.source.frequency).to( unit=time_unit, copy=False ) - data_at_comp.coords['distance'] = comp.distance + neutrons['distance'] = comp.distance + + if "blocked_by_me" in neutrons: + # Because we use shallow copies, we do not want to do an in-place |= + # operation here + neutrons['blocked_by_others'] = neutrons[ + 'blocked_by_others' + ] | neutrons.pop('blocked_by_me') - data_at_comp = comp.apply(data_at_comp) + neutrons, reading = comp.apply(neutrons=neutrons, time_limit=time_limit) + + readings[comp.name] = reading # # data_at_comp.coords['tof'] = tof # if isinstance(c, Detector): @@ -358,12 +365,10 @@ def run(self) -> Result: # combined = initial_mask & m # data_at_comp.masks['blocked_by_others'] = ~initial_mask # data_at_comp.masks['blocked_by_me'] = ~m & initial_mask - initial_mask = combined - neutrons = data_at_comp + # initial_mask = combined + # neutrons = data_at_comp - return Result( - source=self.source, choppers=result_choppers, detectors=result_detectors - ) + return Result(source=self.source, readings=readings) def __repr__(self) -> str: out = f"Model:\n Source: {self.source}\n Choppers:\n" diff --git a/src/tof/result.py b/src/tof/result.py index 0e3ad2b..0f3dcaa 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -51,41 +51,39 @@ class Result: ---------- source: The source of neutrons. - choppers: - The choppers in the model. - detectors: - The detectors in the model. + results: + The state of neutrons at each component in the model. """ def __init__( self, source: Source, - choppers: dict[str, Chopper], - detectors: dict[str, Detector], + readings: dict[str, dict], ): self._source = source.as_readonly() - self._choppers = {} - for name, chopper in choppers.items(): - self._choppers[name] = ChopperReading( - distance=chopper["distance"], - name=chopper["name"], - frequency=chopper["frequency"], - open=chopper["open"], - close=chopper["close"], - phase=chopper["phase"], - open_times=chopper["open_times"], - close_times=chopper["close_times"], - data=chopper["data"], - ) - - self._detectors = {} - for name, det in detectors.items(): - self._detectors[name] = DetectorReading( - distance=det["distance"], name=det["name"], data=det["data"] - ) - - self._choppers = MappingProxyType(self._choppers) - self._detectors = MappingProxyType(self._detectors) + self._components = MappingProxyType(readings) + # self._choppers = {} + # for name, chopper in choppers.items(): + # self._choppers[name] = ChopperReading( + # distance=chopper["distance"], + # name=chopper["name"], + # frequency=chopper["frequency"], + # open=chopper["open"], + # close=chopper["close"], + # phase=chopper["phase"], + # open_times=chopper["open_times"], + # close_times=chopper["close_times"], + # data=chopper["data"], + # ) + + # self._detectors = {} + # for name, det in detectors.items(): + # self._detectors[name] = DetectorReading( + # distance=det["distance"], name=det["name"], data=det["data"] + # ) + + # self._choppers = MappingProxyType(self._choppers) + # self._detectors = MappingProxyType(self._detectors) @property def choppers(self) -> MappingProxyType[str, ChopperReading]: From e23e75362fb6ce353daa19486accf219d4ebe44e Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Sat, 21 Feb 2026 22:15:40 +0100 Subject: [PATCH 04/22] fix imports --- src/tof/chopper.py | 5 +++-- src/tof/detector.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 38dc258..6fb5052 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -8,7 +8,8 @@ import scipp as sc -from .reading import ComponentReading +# from .reading import ComponentReading +from .component import Component, ComponentReading from .utils import two_pi, var_to_dict if TYPE_CHECKING: @@ -77,7 +78,7 @@ def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ChopperRead return replace(self, data=self.data[val]) -class Chopper: +class Chopper(Component): """ A chopper is a rotating device with cutouts that blocks the beam at certain times. diff --git a/src/tof/detector.py b/src/tof/detector.py index 132a360..bb0f463 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -6,8 +6,9 @@ import scipp as sc -from .component import Component -from .reading import ComponentReading +from .component import Component, ComponentReading + +# from .reading import ComponentReading from .utils import var_to_dict From 50da6f88ccbfb6b0491e378077e174927915e86f Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Sun, 22 Feb 2026 10:10:52 +0100 Subject: [PATCH 05/22] model runs and results are plotting --- src/tof/chopper.py | 37 ++++++++++++-- src/tof/component.py | 4 +- src/tof/detector.py | 14 +++++- src/tof/model.py | 64 +++++++++++++++--------- src/tof/result.py | 114 ++++++++++++++++++++++++++----------------- 5 files changed, 157 insertions(+), 76 deletions(-) diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 6fb5052..4a8578f 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -6,6 +6,7 @@ from enum import Enum from typing import TYPE_CHECKING +import numpy as np import scipp as sc # from .reading import ComponentReading @@ -54,6 +55,10 @@ class ChopperReading(ComponentReading): close_times: sc.Variable data: sc.DataArray + @property + def kind(self) -> str: + return "chopper" + def _repr_stats(self) -> str: return ( f"visible={int(self.data.sum().value)}, " @@ -77,6 +82,27 @@ def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ChopperRead val = ('pulse', val) return replace(self, data=self.data[val]) + def plot(self, ax, tmax) -> None: + dx = 0.05 * tmax + x0 = self.open_times.values + x1 = self.close_times.values + x = np.empty(3 * x0.size, dtype=x0.dtype) + x[0::3] = x0 + x[1::3] = 0.5 * (x0 + x1) + x[2::3] = x1 + x = np.concatenate( + ([[0]] if x[0] > 0 else [x[0:1]]) + + [x] + + ([[tmax + dx]] if x[-1] < tmax else []) + ) + y = np.full_like(x, self.distance.value) + y[2::3] = None + inds = np.argsort(x) + ax.plot(x[inds], y[inds], color="k") + ax.text( + tmax, self.distance.value, self.name, ha="right", va="bottom", color="k" + ) + class Chopper(Component): """ @@ -157,6 +183,7 @@ def __init__( self.distance = distance.to(dtype=float, copy=False) self.phase = phase.to(dtype=float, copy=False) self.name = name + self.kind = "chopper" super().__init__() @property @@ -419,7 +446,7 @@ def from_nexus(cls, nexus_chopper, name: str | None = None) -> Chopper: ) def make_reading( - self, neutrons: sc.DataGroup, time_limit: sc.Variable + self, neutrons: sc.DataArray, time_limit: sc.Variable ) -> ChopperReading: """ Create a ChopperReading from the given neutrons that have been processed by this @@ -439,8 +466,8 @@ def make_reading( ) def apply( - self, neutrons: sc.DataGroup, time_limit: sc.Variable - ) -> tuple[sc.DataGroup, ChopperReading]: + self, neutrons: sc.DataArray, time_limit: sc.Variable + ) -> tuple[sc.DataArray, ChopperReading]: """ Apply the effect of the chopper to the given neutrons. """ @@ -449,10 +476,10 @@ def apply( to, tc = self.open_close_times(time_limit=time_limit) # neutrons.update({'open_times': to, 'close_times': tc}) for i in range(len(to)): - m |= (neutrons['toa'] > to[i]) & (neutrons['toa'] < tc[i]) + m |= (neutrons.coords['toa'] > to[i]) & (neutrons.coords['toa'] < tc[i]) # combined = initial_mask & m # data_at_comp.masks['blocked_by_others'] = ~initial_mask - neutrons['blocked_by_me'] = (~m) & (~neutrons['blocked_by_others']) + neutrons.masks['blocked_by_me'] = (~m) & (~neutrons.masks['blocked_by_others']) reading = self.make_reading(neutrons, time_limit=time_limit) return neutrons, reading diff --git a/src/tof/component.py b/src/tof/component.py index a2f4b35..ac84e0e 100644 --- a/src/tof/component.py +++ b/src/tof/component.py @@ -140,7 +140,9 @@ class Component: kind: str @abstractmethod - def apply(self, neutrons: sc.DataArray) -> sc.DataArray: + def apply( + self, neutrons: sc.DataArray, time_limit: sc.Variable + ) -> tuple[sc.DataArray, ComponentReading]: """ Apply the component to the given neutrons. diff --git a/src/tof/detector.py b/src/tof/detector.py index bb0f463..f601ae4 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -22,6 +22,10 @@ class DetectorReading(ComponentReading): name: str data: sc.DataArray + @property + def kind(self) -> str: + return "detector" + def _repr_stats(self) -> str: return f"visible={int(self.data.sum().value)}" @@ -41,6 +45,10 @@ def __getitem__( val = ('pulse', val) return replace(self, data=self.data[val]) + def plot(self, ax, tmax) -> None: + ax.plot([0, tmax], [self.distance.value] * 2, color="gray", lw=3) + ax.text(0, self.distance.value, self.name, ha="left", va="bottom", color="gray") + class Detector(Component): """ @@ -98,10 +106,12 @@ def as_json(self) -> dict: 'name': self.name, } - def make_reading(self, neutrons: sc.DataGroup) -> DetectorReading: + def make_reading(self, neutrons: sc.DataArray) -> DetectorReading: return DetectorReading(distance=self.distance, name=self.name, data=neutrons) - def apply(self, neutrons: sc.DataGroup, time_limit: sc.Variable) -> sc.DataGroup: + def apply( + self, neutrons: sc.DataArray, time_limit: sc.Variable + ) -> tuple[sc.DataArray, DetectorReading]: """ Apply the detector to the given neutrons. A detector does not modify the neutrons, it simply records them without diff --git a/src/tof/model.py b/src/tof/model.py index c696916..d701186 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -128,6 +128,7 @@ def __init__( # self.choppers = {} # self.detectors = {} self.source = source + self.components = {} # for components, kind in ((choppers, Chopper), (detectors, Detector)): for comp in chain((choppers or ()), (detectors or ()), (components or ())): self.add(comp) @@ -307,23 +308,26 @@ def run(self) -> Result: # birth_time = self.source.data.coords['birth_time'] # speed = self.source.data.coords['speed'] - neutrons = sc.DataGroup(self.source.data.coords) - neutrons.update( - blocked_by_others=sc.zeros(sizes=neutrons.sizes, unit=None, dtype=bool), - distance=self.source.distance, - toa=neutrons['birth_time'], + # neutrons = sc.DataGroup(self.source.data.coords) + neutrons = self.source.data.copy(deep=False) + neutrons.masks["blocked_by_others"] = sc.zeros( + sizes=neutrons.sizes, unit=None, dtype=bool + ) + neutrons.coords.update( + distance=self.source.distance, toa=neutrons.coords['birth_time'] ) - time_unit = neutrons['birth_time'].unit + time_unit = neutrons.coords['birth_time'].unit readings = {} # result_choppers = {} # result_detectors = {} time_limit = ( - neutrons['birth_time'] - + ((components[-1].distance - self.source.distance) / neutrons['speed']).to( - unit=time_unit - ) + neutrons.coords['birth_time'] + + ( + (components[-1].distance - self.source.distance) + / neutrons.coords['speed'] + ).to(unit=time_unit) ).max() for comp in components: # results = result_detectors if isinstance(c, Detector) else result_choppers @@ -333,21 +337,21 @@ def run(self) -> Result: # unit=time_unit, copy=False # ) # t = birth_time + tof - toa = neutrons['toa'] + ( - (comp.distance - neutrons['distance']) / neutrons['speed'] + toa = neutrons.coords['toa'] + ( + (comp.distance - neutrons.coords['distance']) / neutrons.coords['speed'] ).to(unit=time_unit, copy=False) - neutrons['toa'] = toa - neutrons['eto'] = toa % (1 / self.source.frequency).to( + neutrons.coords['toa'] = toa + neutrons.coords['eto'] = toa % (1 / self.source.frequency).to( unit=time_unit, copy=False ) - neutrons['distance'] = comp.distance + neutrons.coords['distance'] = comp.distance - if "blocked_by_me" in neutrons: + if "blocked_by_me" in neutrons.masks: # Because we use shallow copies, we do not want to do an in-place |= # operation here - neutrons['blocked_by_others'] = neutrons[ + neutrons.masks['blocked_by_others'] = neutrons.masks[ 'blocked_by_others' - ] | neutrons.pop('blocked_by_me') + ] | neutrons.masks.pop('blocked_by_me') neutrons, reading = comp.apply(neutrons=neutrons, time_limit=time_limit) @@ -371,12 +375,24 @@ def run(self) -> Result: return Result(source=self.source, readings=readings) def __repr__(self) -> str: - out = f"Model:\n Source: {self.source}\n Choppers:\n" - for name, ch in self.choppers.items(): - out += f" {name}: {ch}\n" - out += " Detectors:\n" - for name, det in self.detectors.items(): - out += f" {name}: {det}\n" + # out = f"Model:\n Source: {self.source}\n Choppers:\n" + # for name, ch in self.choppers.items(): + # out += f" {name}: {ch}\n" + # out += " Detectors:\n" + # for name, det in self.detectors.items(): + # out += f" {name}: {det}\n" + # return out + out = f"Model:\n Source: {self.source}\n" + groups = {} + for comp in self.components.values(): + if comp.kind not in groups: + groups[comp.kind] = [] + groups[comp.kind].append(comp) + + for group, comps in groups.items(): + out += f" {group.capitalize()}s:\n" + for comp in sorted(comps, key=lambda c: c.distance): + out += f" {comp.name}: {comp}\n" return out def __str__(self) -> str: diff --git a/src/tof/result.py b/src/tof/result.py index 0f3dcaa..82a76c2 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -12,6 +12,7 @@ from matplotlib.collections import LineCollection from .chopper import Chopper, ChopperReading +from .component import ComponentReading from .detector import Detector, DetectorReading from .source import Source, SourceParameters from .utils import Plot, one_mask @@ -88,12 +89,24 @@ def __init__( @property def choppers(self) -> MappingProxyType[str, ChopperReading]: """The choppers in the model.""" - return self._choppers + return MappingProxyType( + { + key: comp + for key, comp in self._components.items() + if comp.kind == "chopper" + } + ) @property def detectors(self) -> MappingProxyType[str, DetectorReading]: """The detectors in the model.""" - return self._detectors + return MappingProxyType( + { + key: comp + for key, comp in self._components.items() + if comp.kind == "detector" + } + ) @property def source(self) -> SourceParameters: @@ -101,12 +114,12 @@ def source(self) -> SourceParameters: return self._source def __iter__(self): - return chain(self._choppers, self._detectors) + return iter(self._components) - def __getitem__(self, name: str) -> ChopperReading | DetectorReading: - if name not in self: - raise KeyError(f"No component with name {name} was found.") - return self._choppers[name] if name in self._choppers else self._detectors[name] + def __getitem__(self, name: str) -> ComponentReading: + # if name not in self: + # raise KeyError(f"No component with name {name} was found.") + return self._components[name] def plot( self, @@ -156,10 +169,7 @@ def plot( fig, ax = plt.subplots(figsize=figsize) else: fig = ax.get_figure() - components = sorted( - chain(self.choppers.values(), self.detectors.values()), - key=lambda c: c.distance, - ) + components = sorted(self._components.values(), key=lambda c: c.distance) furthest_component = components[-1] source_dist = self.source.distance.value repeats = [1] + [2] * len(components) @@ -248,35 +258,40 @@ def plot( toa_max = furthest_component.toa.max().value else: toa_max = furthest_component.toa.data.coords["toa"].max().value - dx = 0.05 * toa_max - # Plot choppers - for ch in self._choppers.values(): - x0 = ch.open_times.values - x1 = ch.close_times.values - x = np.empty(3 * x0.size, dtype=x0.dtype) - x[0::3] = x0 - x[1::3] = 0.5 * (x0 + x1) - x[2::3] = x1 - x = np.concatenate( - ([[0]] if x[0] > 0 else [x[0:1]]) - + [x] - + ([[toa_max + dx]] if x[-1] < toa_max else []) - ) - y = np.full_like(x, ch.distance.value) - y[2::3] = None - inds = np.argsort(x) - ax.plot(x[inds], y[inds], color="k") - ax.text( - toa_max, ch.distance.value, ch.name, ha="right", va="bottom", color="k" - ) - # Plot detectors - for det in self._detectors.values(): - ax.plot([0, toa_max], [det.distance.value] * 2, color="gray", lw=3) - ax.text( - 0, det.distance.value, det.name, ha="left", va="bottom", color="gray" - ) + for comp in self._components.values(): + comp.plot(ax=ax, tmax=toa_max) + + # dx = 0.05 * toa_max + # # Plot choppers + # for ch in self._choppers.values(): + # x0 = ch.open_times.values + # x1 = ch.close_times.values + # x = np.empty(3 * x0.size, dtype=x0.dtype) + # x[0::3] = x0 + # x[1::3] = 0.5 * (x0 + x1) + # x[2::3] = x1 + # x = np.concatenate( + # ([[0]] if x[0] > 0 else [x[0:1]]) + # + [x] + # + ([[toa_max + dx]] if x[-1] < toa_max else []) + # ) + # y = np.full_like(x, ch.distance.value) + # y[2::3] = None + # inds = np.argsort(x) + # ax.plot(x[inds], y[inds], color="k") + # ax.text( + # toa_max, ch.distance.value, ch.name, ha="right", va="bottom", color="k" + # ) + # # Plot detectors + # for det in self._detectors.values(): + # ax.plot([0, toa_max], [det.distance.value] * 2, color="gray", lw=3) + # ax.text( + # 0, det.distance.value, det.name, ha="left", va="bottom", color="gray" + # ) + + dx = 0.05 * toa_max ax.set(xlabel="Time [μs]", ylabel="Distance [m]") ax.set_xlim(0 - dx, toa_max + dx) if figsize is None: @@ -288,13 +303,24 @@ def plot( def __repr__(self) -> str: out = ( f"Result:\n Source: {self.source.pulses} pulses, " - f"{self.source.neutrons} neutrons per pulse.\n Choppers:\n" + f"{self.source.neutrons} neutrons per pulse.\n" ) - for name, ch in self._choppers.items(): - out += f" {name}: {ch._repr_stats()}\n" - out += " Detectors:\n" - for name, det in self._detectors.items(): - out += f" {name}: {det._repr_stats()}\n" + groups = {} + for comp in self._components.values(): + if comp.kind not in groups: + groups[comp.kind] = [] + groups[comp.kind].append(comp) + + for group, comps in groups.items(): + out += f" {group.capitalize()}s:\n" + for comp in sorted(comps, key=lambda c: c.distance): + out += f" {comp.name}: {comp._repr_stats()}\n" + + # for name, ch in self._choppers.items(): + # out += f" {name}: {ch._repr_stats()}\n" + # out += " Detectors:\n" + # for name, det in self._detectors.items(): + # out += f" {name}: {det._repr_stats()}\n" return out def __str__(self) -> str: From 58f0b256d76113c2352f1ba786d1dac698a883c0 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Sun, 22 Feb 2026 23:44:56 +0100 Subject: [PATCH 06/22] fix plotting of rays --- src/tof/chopper.py | 7 +-- src/tof/component.py | 6 +- src/tof/detector.py | 6 +- src/tof/model.py | 2 +- src/tof/result.py | 128 +++++++++++++------------------------------ src/tof/source.py | 20 ++++++- 6 files changed, 64 insertions(+), 105 deletions(-) diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 4a8578f..151894b 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -82,7 +82,7 @@ def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ChopperRead val = ('pulse', val) return replace(self, data=self.data[val]) - def plot(self, ax, tmax) -> None: + def plot_on_time_distance_diagram(self, ax, tmax) -> None: dx = 0.05 * tmax x0 = self.open_times.values x1 = self.close_times.values @@ -445,7 +445,7 @@ def from_nexus(cls, nexus_chopper, name: str | None = None) -> Chopper: name=name, ) - def make_reading( + def as_readonly( self, neutrons: sc.DataArray, time_limit: sc.Variable ) -> ChopperReading: """ @@ -481,5 +481,4 @@ def apply( # data_at_comp.masks['blocked_by_others'] = ~initial_mask neutrons.masks['blocked_by_me'] = (~m) & (~neutrons.masks['blocked_by_others']) - reading = self.make_reading(neutrons, time_limit=time_limit) - return neutrons, reading + return neutrons, self.as_readonly(neutrons, time_limit=time_limit) diff --git a/src/tof/component.py b/src/tof/component.py index ac84e0e..628ff85 100644 --- a/src/tof/component.py +++ b/src/tof/component.py @@ -67,11 +67,7 @@ def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ReadingFiel def _make_reading_field(da: sc.DataArray, dim: str) -> ReadingField: return ReadingField( - data=sc.DataArray( - data=da.data, - coords={dim: da.coords[dim]}, - masks=da.masks, - ), + data=sc.DataArray(data=da.data, coords={dim: da.coords[dim]}, masks=da.masks), dim=dim, ) diff --git a/src/tof/detector.py b/src/tof/detector.py index f601ae4..7302bab 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -45,7 +45,7 @@ def __getitem__( val = ('pulse', val) return replace(self, data=self.data[val]) - def plot(self, ax, tmax) -> None: + def plot_on_time_distance_diagram(self, ax, tmax) -> None: ax.plot([0, tmax], [self.distance.value] * 2, color="gray", lw=3) ax.text(0, self.distance.value, self.name, ha="left", va="bottom", color="gray") @@ -106,7 +106,7 @@ def as_json(self) -> dict: 'name': self.name, } - def make_reading(self, neutrons: sc.DataArray) -> DetectorReading: + def as_readonly(self, neutrons: sc.DataArray) -> DetectorReading: return DetectorReading(distance=self.distance, name=self.name, data=neutrons) def apply( @@ -125,4 +125,4 @@ def apply( The time limit for the neutrons to be considered as reaching the detector. """ # neutrons.pop("blocked_by_me", None) - return neutrons, self.make_reading(neutrons) + return neutrons, self.as_readonly(neutrons) diff --git a/src/tof/model.py b/src/tof/model.py index d701186..46d0f99 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -372,7 +372,7 @@ def run(self) -> Result: # initial_mask = combined # neutrons = data_at_comp - return Result(source=self.source, readings=readings) + return Result(source=self.source.as_readonly(), readings=readings) def __repr__(self) -> str: # out = f"Model:\n Source: {self.source}\n Choppers:\n" diff --git a/src/tof/result.py b/src/tof/result.py index 82a76c2..dd1b47b 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -14,10 +14,24 @@ from .chopper import Chopper, ChopperReading from .component import ComponentReading from .detector import Detector, DetectorReading -from .source import Source, SourceParameters +from .source import SourceReading from .utils import Plot, one_mask +def _get_rays( + components: list[ComponentReading], pulse: int, inds: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + x = np.stack( + [comp.data["pulse", pulse].coords["toa"].values[inds] for comp in components], + axis=1, + ) + y = np.stack( + [np.full_like(x[:, 0], comp.distance.value) for comp in components], + axis=1, + ) + return x, y + + def _add_rays( ax: plt.Axes, x: np.ndarray, @@ -56,12 +70,8 @@ class Result: The state of neutrons at each component in the model. """ - def __init__( - self, - source: Source, - readings: dict[str, dict], - ): - self._source = source.as_readonly() + def __init__(self, source: SourceReading, readings: dict[str, dict]): + self._source = source self._components = MappingProxyType(readings) # self._choppers = {} # for name, chopper in choppers.items(): @@ -109,7 +119,7 @@ def detectors(self) -> MappingProxyType[str, DetectorReading]: ) @property - def source(self) -> SourceParameters: + def source(self) -> SourceReading: """The source of neutrons.""" return self._source @@ -169,10 +179,11 @@ def plot( fig, ax = plt.subplots(figsize=figsize) else: fig = ax.get_figure() - components = sorted(self._components.values(), key=lambda c: c.distance) + + components = sorted( + chain((self.source,), self._components.values()), key=lambda c: c.distance + ) furthest_component = components[-1] - source_dist = self.source.distance.value - repeats = [1] + [2] * len(components) wavelengths = sc.DataArray( data=furthest_component.data.coords["wavelength"], @@ -183,9 +194,9 @@ def plot( rng = np.random.default_rng(seed) for i in range(self._source.data.sizes["pulse"]): - source_data = self.source.data["pulse", i] component_data = furthest_component.data["pulse", i] - ids = np.arange(self.source.neutrons) + ids = component_data.coords["id"].values + # Plot visible rays blocked = one_mask(component_data.masks).values nblocked = int(blocked.sum()) @@ -195,17 +206,12 @@ def plot( size=min(self.source.neutrons - nblocked, visible_rays), replace=False, ) - - xstart = source_data.coords["birth_time"].values[inds] - xend = component_data.coords["toa"].values[inds] - ystart = np.full_like(xstart, source_dist) - yend = np.full_like(ystart, furthest_component.distance.value) - + x, y = _get_rays(components, pulse=i, inds=inds) _add_rays( ax=ax, - x=np.stack((xstart, xend), axis=1), - y=np.stack((ystart, yend), axis=1), - color=source_data.coords["wavelength"].values[inds], + x=x, + y=y, + color=component_data.coords["wavelength"].values[inds], cbar=cbar and (i == 0), cmap=cmap, vmin=wmin.value if vmin is None else vmin, @@ -217,79 +223,28 @@ def plot( inds = rng.choice( ids[blocked], size=min(blocked_rays, nblocked), replace=False ) - x = np.repeat( - np.stack( - [source_data.coords["birth_time"].values[inds]] - + [ - c.data.coords["toa"]["pulse", i].values[inds] - for c in components - ], - axis=1, - ), - repeats, - axis=1, - ) - y = np.repeat( - np.stack( - [np.full_like(x[:, 0], source_dist)] - + [np.full_like(x[:, 0], c.distance.value) for c in components], - axis=1, - ), - repeats, + x, y = _get_rays(components, pulse=i, inds=inds) + blocked_by_others = np.stack( + [ + comp.data["pulse", i].masks["blocked_by_others"].values[inds] + for comp in components + ], axis=1, ) - for j, c in enumerate(components): - comp_data = c.data["pulse", i] - m_others = comp_data.masks["blocked_by_others"].values[inds] - x[:, 2 * j + 1][m_others] = np.nan - y[:, 2 * j + 1][m_others] = np.nan - if "blocked_by_me" in comp_data.masks: - m_me = comp_data.masks["blocked_by_me"].values[inds] - x[:, 2 * j + 2][m_me] = np.nan - y[:, 2 * j + 2][m_me] = np.nan + x[blocked_by_others] = np.nan + y[blocked_by_others] = np.nan _add_rays(ax=ax, x=x, y=y, color="lightgray", zorder=-1) # Plot pulse - time_coord = source_data.coords["birth_time"].values - tmin = time_coord.min() - ax.plot([tmin, time_coord.max()], [source_dist] * 2, color="gray", lw=3) - ax.text(tmin, source_dist, "Pulse", ha="left", va="top", color="gray") + self.source.plot_on_time_distance_diagram(ax, pulse=i) if furthest_component.toa.data.sum().value > 0: toa_max = furthest_component.toa.max().value else: toa_max = furthest_component.toa.data.coords["toa"].max().value + # Plot components for comp in self._components.values(): - comp.plot(ax=ax, tmax=toa_max) - - # dx = 0.05 * toa_max - # # Plot choppers - # for ch in self._choppers.values(): - # x0 = ch.open_times.values - # x1 = ch.close_times.values - # x = np.empty(3 * x0.size, dtype=x0.dtype) - # x[0::3] = x0 - # x[1::3] = 0.5 * (x0 + x1) - # x[2::3] = x1 - # x = np.concatenate( - # ([[0]] if x[0] > 0 else [x[0:1]]) - # + [x] - # + ([[toa_max + dx]] if x[-1] < toa_max else []) - # ) - # y = np.full_like(x, ch.distance.value) - # y[2::3] = None - # inds = np.argsort(x) - # ax.plot(x[inds], y[inds], color="k") - # ax.text( - # toa_max, ch.distance.value, ch.name, ha="right", va="bottom", color="k" - # ) - - # # Plot detectors - # for det in self._detectors.values(): - # ax.plot([0, toa_max], [det.distance.value] * 2, color="gray", lw=3) - # ax.text( - # 0, det.distance.value, det.name, ha="left", va="bottom", color="gray" - # ) + comp.plot_on_time_distance_diagram(ax=ax, tmax=toa_max) dx = 0.05 * toa_max ax.set(xlabel="Time [μs]", ylabel="Distance [m]") @@ -316,11 +271,6 @@ def __repr__(self) -> str: for comp in sorted(comps, key=lambda c: c.distance): out += f" {comp.name}: {comp._repr_stats()}\n" - # for name, ch in self._choppers.items(): - # out += f" {name}: {ch._repr_stats()}\n" - # out += " Detectors:\n" - # for name, det in self._detectors.items(): - # out += f" {name}: {det._repr_stats()}\n" return out def __str__(self) -> str: diff --git a/src/tof/source.py b/src/tof/source.py index 504d4c3..49d40cd 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -8,6 +8,7 @@ import plopp as pp import scipp as sc +from .component import ComponentReading from .utils import wavelength_to_speed TIME_UNIT = "us" @@ -452,8 +453,10 @@ def plot(self, bins: int = 300) -> tuple: return f1 + f2 def as_readonly(self): - return SourceParameters( - data=self.data, + return SourceReading( + data=self.data.assign_masks( + blocked_by_others=sc.zeros_like(self.data.data, dtype=bool, unit=None) + ), facility=self.facility, neutrons=self.neutrons, frequency=self.frequency, @@ -485,7 +488,7 @@ def as_json(self) -> dict: @dataclass(frozen=True) -class SourceParameters: +class SourceReading(ComponentReading): """ Read-only container for the parameters of a source. """ @@ -496,3 +499,14 @@ class SourceParameters: 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 ddb0eb89009235141ad2ab49a35e7bb4a23fc6d4 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Sun, 22 Feb 2026 23:45:47 +0100 Subject: [PATCH 07/22] remove old reading file --- src/tof/chopper.py | 1 - src/tof/detector.py | 2 - src/tof/reading.py | 135 -------------------------------------------- 3 files changed, 138 deletions(-) delete mode 100644 src/tof/reading.py diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 151894b..fce799e 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -9,7 +9,6 @@ import numpy as np import scipp as sc -# from .reading import ComponentReading from .component import Component, ComponentReading from .utils import two_pi, var_to_dict diff --git a/src/tof/detector.py b/src/tof/detector.py index 7302bab..5e55275 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -7,8 +7,6 @@ import scipp as sc from .component import Component, ComponentReading - -# from .reading import ComponentReading from .utils import var_to_dict diff --git a/src/tof/reading.py b/src/tof/reading.py deleted file mode 100644 index 62887de..0000000 --- a/src/tof/reading.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2025 Scipp contributors (https://github.com/scipp) - -from __future__ import annotations - -from dataclasses import dataclass - -import plopp as pp -import scipp as sc - -from .utils import Plot - - -@dataclass(frozen=True) -class ReadingField: - data: sc.DataArray - dim: str - - def plot(self, bins: int = 300, **kwargs): - by_pulse = sc.collapse(self.data, keep="event") - to_plot = {} - color = {} - for key, da in by_pulse.items(): - sel = da[~da.masks["blocked_by_others"]] - if sel.size == 0: - continue - to_plot[key] = sel.hist({self.dim: bins}) - if "blocked_by_me" in self.data.masks: - name = f"blocked-{key}" - to_plot[name] = ( - da[da.masks["blocked_by_me"]] - .drop_masks(list(da.masks.keys())) - .hist({self.dim: to_plot[key].coords[self.dim]}) - ) - color[name] = "gray" - if not to_plot: - raise RuntimeError("Nothing to plot.") - return pp.plot(to_plot, **{**{"color": color}, **kwargs}) - - def min(self): - da = self.data.copy(deep=False) - da.data = da.coords[self.dim] - return da.min().data - - def max(self): - da = self.data.copy(deep=False) - da.data = da.coords[self.dim] - return da.max().data - - def __repr__(self) -> str: - da = self.data.copy(deep=False) - da.data = da.coords[self.dim] - return ( - f"{self.dim}: min={da.min().data:c}, max={da.max().data:c}, " - f"events={int(self.data.sum().value)}" - ) - - def __str__(self) -> str: - return self.__repr__() - - def __getitem__(self, val: int | slice | tuple[str, int | slice]) -> ReadingField: - if isinstance(val, int): - val = ('pulse', val) - return self.__class__(data=self.data[val], dim=self.dim) - - -def _make_reading_field(da: sc.DataArray, dim: str) -> ReadingField: - return ReadingField( - data=sc.DataArray( - data=da.data, - coords={dim: da.coords[dim]}, - masks=da.masks, - ), - dim=dim, - ) - - -class ComponentReading: - """ - Data reading for a component placed in the beam path. The reading will have a - record of the arrival times and wavelengths of the neutrons that passed through it. - """ - - @property - def toa(self) -> ReadingField: - """ - Time of arrival of the neutrons at the component. - """ - return _make_reading_field(self.data, dim="toa") - - @property - def tof(self) -> ReadingField: - """ - Time of flight of the neutrons from source to the component. - """ - return _make_reading_field(self.data, dim="tof") - - @property - def eto(self) -> ReadingField: - """ - Event time offset of the neutrons at the component (= toa modulo pulse period). - """ - return _make_reading_field(self.data, dim="eto") - - @property - def wavelength(self) -> ReadingField: - """ - Wavelength of the neutrons at the component. - """ - return _make_reading_field(self.data, dim="wavelength") - - @property - def birth_time(self) -> ReadingField: - """ - Birth time of the neutrons at the source. - """ - return _make_reading_field(self.data, dim="birth_time") - - @property - def speed(self) -> ReadingField: - """ - Speed of the neutrons at the component. - """ - return _make_reading_field(self.data, dim="speed") - - def plot(self, bins: int = 300) -> Plot: - """ - Plot both the toa and wavelength data side by side. - - Parameters - ---------- - bins: - Number of bins to use for histogramming the neutrons. - """ - return self.toa.plot(bins=bins) + self.wavelength.plot(bins=bins) From 7a8101dc06c15e26aecaf0478baa971b76569597 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Mon, 23 Feb 2026 00:30:17 +0100 Subject: [PATCH 08/22] start adding inelastic sample --- src/tof/__init__.py | 5 +- src/tof/inelastic.py | 170 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 src/tof/inelastic.py diff --git a/src/tof/__init__.py b/src/tof/__init__.py index d3a75af..93f8699 100644 --- a/src/tof/__init__.py +++ b/src/tof/__init__.py @@ -17,12 +17,13 @@ submodules=['facilities'], submod_attrs={ 'chopper': ['AntiClockwise', 'Chopper', 'ChopperReading', 'Clockwise'], + 'component': ['Component', 'ComponentReading', 'ReadingField'], 'dashboard': ['Dashboard'], 'detector': ['Detector', 'DetectorReading'], + 'inelastic': ['InelasticSample', 'InelasticSampleReading'], 'model': ['Model'], - 'reading': ['ComponentReading', 'ReadingField'], 'result': ['Result'], - 'source': ['Source', 'SourceParameters'], + 'source': ['Source', 'SourceReading'], }, ) diff --git a/src/tof/inelastic.py b/src/tof/inelastic.py new file mode 100644 index 0000000..00a007b --- /dev/null +++ b/src/tof/inelastic.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2026 Scipp contributors (https://github.com/scipp) +from __future__ import annotations + +from dataclasses import dataclass, replace + +import numpy as np +import plopp as pp +import scipp as sc + +from .component import Component, ComponentReading +from .utils import var_to_dict, wavelength_to_speed + + +@dataclass(frozen=True) +class InelasticSampleReading(ComponentReading): + """ + Read-only container for the neutrons that reach the inelastic sample. + """ + + distance: sc.Variable + name: str + data: sc.DataArray + + @property + def kind(self) -> str: + return "inelastic_sample" + + def _repr_stats(self) -> str: + return f"visible={int(self.data.sum().value)}" + + def __repr__(self) -> str: + return f"""InelasticSampleReading: '{self.name}' + distance: {self.distance:c} + neutrons: {self._repr_stats()} +""" + + def __str__(self) -> str: + return self.__repr__() + + def __getitem__( + self, val: int | slice | tuple[str, int | slice] + ) -> InelasticSampleReading: + if isinstance(val, int): + val = ('pulse', val) + return replace(self, data=self.data[val]) + + def plot_on_time_distance_diagram(self, ax, tmax) -> None: + ax.plot([0, tmax], [self.distance.value] * 2, color="gray", lw=3) + ax.text(0, self.distance.value, self.name, ha="left", va="bottom", color="gray") + + +class InelasticSample(Component): + """ + An inelastic sample component changes the energy of the neutrons that pass through + it, but does not block any. + + Parameters + ---------- + distance: + The distance from the source to the inelastic sample. + name: + The name of the inelastic sample. + """ + + def __init__( + self, + distance: sc.Variable, + name: str, + delta_e: sc.DataArray, + seed: int | None = None, + ): + self.distance = distance.to(dtype=float, copy=False) + self.name = name + if delta_e.ndim != 1: + raise ValueError("delta_e must be a 1D array.") + self.probabilities = delta_e.data / delta_e.sum().value + dim = delta_e.dim + self.energies = delta_e.coords[dim] + # TODO: check for bin edges + self._noise_scale = ( + 0.5 * (self.energies.max() - self.energies.min()).value / (len(delta_e) - 1) + ) + self.kind = "inelastic_sample" + self._rng = np.random.default_rng(seed) + + def __repr__(self) -> str: + return f"InelasticSample(name={self.name}, distance={self.distance:c})" + + def plot(self, **kwargs) -> pp.FigureLike: + return pp.xyplot(self.energies, self.probabilities, **kwargs) + + # def __eq__(self, other: object) -> bool: + # if not isinstance(other, Detector): + # return NotImplemented + # return self.name == other.name and sc.identical(self.distance, other.distance) + + def as_dict(self) -> dict: + """ + Return the inelastic sample as a dictionary. + """ + return {'distance': self.distance, 'name': self.name, 'delta_e': self.delta_e} + + @classmethod + def from_json(cls, name: str, params: dict) -> InelasticSample: + """ + Create an inelastic sample from a JSON-serializable dictionary. + """ + return cls( + distance=sc.scalar( + params["distance"]["value"], unit=params["distance"]["unit"] + ), + name=name, + delta_e=sc.scalar( + params["delta_e"]["value"], unit=params["delta_e"]["unit"] + ), + ) + + def as_json(self) -> dict: + """ + Return the inelastic sample as a JSON-serializable dictionary. + .. versionadded:: 26.03.0 + """ + return { + 'type': 'inelastic_sample', + 'distance': var_to_dict(self.distance), + 'name': self.name, + 'delta_e': var_to_dict(self.delta_e), + } + + def as_readonly(self, neutrons: sc.DataArray) -> InelasticSampleReading: + return InelasticSampleReading( + distance=self.distance, name=self.name, data=neutrons + ) + + def apply( + self, neutrons: sc.DataArray, time_limit: sc.Variable + ) -> tuple[sc.DataArray, InelasticSampleReading]: + """ + Apply the change in energy to the given neutrons. + + Parameters + ---------- + neutrons: + The neutrons to which the inelastic sample will be applied. + time_limit: + The time limit for the neutrons to be considered as reaching the inelastic + sample. + """ + # neutrons.pop("blocked_by_me", None) + # return neutrons, self.as_readonly(neutrons) + n = neutrons.sizes["pulse"] + inds = self._rng.choice(len(self.delta_e), size=n, p=self.probabilities.values) + de = self.energies.values[inds] + self._rng.normal( + scale=self._noise_scale, size=n + ) + # Convert energy change to wavelength change + w_initial = sc.reciprocal(neutrons.coords["wavelength"] ** 2) + w_final = sc.reciprocal( + sc.sqrt( + ((2 * sc.constants.m_n / (sc.constants.h**2)) * de).to( + unit=w_initial.unit + ) + + w_initial + ) + ) + neutrons = neutrons.assign_coords( + wavelength=w_final, speed=wavelength_to_speed(w_final) + ) + return neutrons, self.as_readonly(neutrons) From 1bf824c56ce12dd607800d3629380741b22aafd1 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Mon, 23 Feb 2026 16:15:59 +0100 Subject: [PATCH 09/22] fix deltae --- src/tof/inelastic.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/tof/inelastic.py b/src/tof/inelastic.py index 00a007b..037f2b6 100644 --- a/src/tof/inelastic.py +++ b/src/tof/inelastic.py @@ -74,7 +74,8 @@ def __init__( self.name = name if delta_e.ndim != 1: raise ValueError("delta_e must be a 1D array.") - self.probabilities = delta_e.data / delta_e.sum().value + self.probabilities = delta_e.values + self.probabilities = self.probabilities / self.probabilities.sum() dim = delta_e.dim self.energies = delta_e.coords[dim] # TODO: check for bin edges @@ -114,6 +115,7 @@ def from_json(cls, name: str, params: dict) -> InelasticSample: delta_e=sc.scalar( params["delta_e"]["value"], unit=params["delta_e"]["unit"] ), + seed=params.get("seed"), ) def as_json(self) -> dict: @@ -149,13 +151,18 @@ def apply( """ # neutrons.pop("blocked_by_me", None) # return neutrons, self.as_readonly(neutrons) - n = neutrons.sizes["pulse"] - inds = self._rng.choice(len(self.delta_e), size=n, p=self.probabilities.values) - de = self.energies.values[inds] + self._rng.normal( - scale=self._noise_scale, size=n + w_initial = sc.reciprocal(neutrons.coords["wavelength"] ** 2) + + n = neutrons.shape + inds = self._rng.choice(len(self.energies), size=n, p=self.probabilities) + de = sc.array( + dims=w_initial.dims, + values=self.energies.values[inds] + + self._rng.normal(scale=self._noise_scale, size=n), + unit=self.energies.unit, ) # Convert energy change to wavelength change - w_initial = sc.reciprocal(neutrons.coords["wavelength"] ** 2) + # w_initial = sc.reciprocal(neutrons.coords["wavelength"] ** 2) w_final = sc.reciprocal( sc.sqrt( ((2 * sc.constants.m_n / (sc.constants.h**2)) * de).to( From 0933abb0afa7f2f86ab8aa71b39786686473667d Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Mon, 23 Feb 2026 17:50:26 +0100 Subject: [PATCH 10/22] plot wavelength color for each section between components --- src/tof/result.py | 75 ++++++++++++++++++++++++++++++----------------- 1 file changed, 48 insertions(+), 27 deletions(-) diff --git a/src/tof/result.py b/src/tof/result.py index dd1b47b..d3c260c 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -21,15 +21,32 @@ def _get_rays( components: list[ComponentReading], pulse: int, inds: np.ndarray ) -> tuple[np.ndarray, np.ndarray]: - x = np.stack( - [comp.data["pulse", pulse].coords["toa"].values[inds] for comp in components], - axis=1, - ) - y = np.stack( - [np.full_like(x[:, 0], comp.distance.value) for comp in components], - axis=1, - ) - return x, y + x = [] + y = [] + c = [] + xstart = components[0].data["pulse", pulse].coords["toa"].values[inds] + ystart = np.full_like(xstart, components[0].distance.value) + for comp in components[1:]: + xend = comp.data["pulse", pulse].coords["toa"].values[inds] + yend = np.full_like(xend, comp.distance.value) + x.append([xstart, xend]) + y.append([ystart, yend]) + c.append(comp.data["pulse", pulse].coords["wavelength"].values[inds]) + xstart, ystart = xend, yend + + # comp_data = compo.data["pulse", pulse] + # x.append(comp_data.coords["toa"].values[inds]) + # y.append(np.full_like(x[-1], comp.distance.value)) + # c.append(comp_data.coords["wavelength"].values[inds]) + # x = np.stack( + # [comp.data["pulse", pulse].coords["toa"].values[inds] for comp in components], + # axis=1, + # ) + # y = np.stack( + # [np.full_like(x[:, 0], comp.distance.value) for comp in components], + # axis=1, + # ) + return np.array(x), np.array(y), np.array(c) def _add_rays( @@ -44,12 +61,15 @@ def _add_rays( cax: plt.Axes | None = None, zorder: int = 1, ): + print(x.shape, y.shape) + # print(np.stack((x, y), axis=2).shape) + x, y = (np.array(a).transpose((0, 2, 1)).reshape((-1, 2)) for a in (x, y)) coll = LineCollection(np.stack((x, y), axis=2), zorder=zorder) if isinstance(color, str): coll.set_color(color) else: coll.set_cmap(plt.colormaps[cmap]) - coll.set_array(color) + coll.set_array(color.ravel()) coll.set_norm(plt.Normalize(vmin, vmax)) if cbar: cb = plt.colorbar(coll, ax=ax, cax=cax) @@ -206,12 +226,13 @@ def plot( size=min(self.source.neutrons - nblocked, visible_rays), replace=False, ) - x, y = _get_rays(components, pulse=i, inds=inds) + x, y, c = _get_rays(components, pulse=i, inds=inds) + print(x.shape, y.shape, c.shape) _add_rays( ax=ax, x=x, y=y, - color=component_data.coords["wavelength"].values[inds], + color=c, cbar=cbar and (i == 0), cmap=cmap, vmin=wmin.value if vmin is None else vmin, @@ -219,21 +240,21 @@ def plot( cax=cax, ) - # Plot blocked rays - inds = rng.choice( - ids[blocked], size=min(blocked_rays, nblocked), replace=False - ) - x, y = _get_rays(components, pulse=i, inds=inds) - blocked_by_others = np.stack( - [ - comp.data["pulse", i].masks["blocked_by_others"].values[inds] - for comp in components - ], - axis=1, - ) - x[blocked_by_others] = np.nan - y[blocked_by_others] = np.nan - _add_rays(ax=ax, x=x, y=y, color="lightgray", zorder=-1) + # # Plot blocked rays + # inds = rng.choice( + # ids[blocked], size=min(blocked_rays, nblocked), replace=False + # ) + # x, y = _get_rays(components, pulse=i, inds=inds) + # blocked_by_others = np.stack( + # [ + # comp.data["pulse", i].masks["blocked_by_others"].values[inds] + # for comp in components + # ], + # axis=1, + # ) + # x[blocked_by_others] = np.nan + # y[blocked_by_others] = np.nan + # _add_rays(ax=ax, x=x, y=y, color="lightgray", zorder=-1) # Plot pulse self.source.plot_on_time_distance_diagram(ax, pulse=i) From 4029e4926dd9792970d27e32bf59f958baf76819 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Tue, 24 Feb 2026 11:49:12 +0100 Subject: [PATCH 11/22] fix ray colors --- src/tof/result.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/tof/result.py b/src/tof/result.py index d3c260c..ac77c90 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -24,15 +24,19 @@ def _get_rays( x = [] y = [] c = [] - xstart = components[0].data["pulse", pulse].coords["toa"].values[inds] + data = components[0].data["pulse", pulse] + # TODO optimize: should not index multiple times + xstart = data.coords["toa"].values[inds] ystart = np.full_like(xstart, components[0].distance.value) + color = data.coords["wavelength"].values[inds] for comp in components[1:]: xend = comp.data["pulse", pulse].coords["toa"].values[inds] yend = np.full_like(xend, comp.distance.value) x.append([xstart, xend]) y.append([ystart, yend]) - c.append(comp.data["pulse", pulse].coords["wavelength"].values[inds]) + c.append(color) xstart, ystart = xend, yend + color = comp.data["pulse", pulse].coords["wavelength"].values[inds] # comp_data = compo.data["pulse", pulse] # x.append(comp_data.coords["toa"].values[inds]) @@ -61,7 +65,7 @@ def _add_rays( cax: plt.Axes | None = None, zorder: int = 1, ): - print(x.shape, y.shape) + # print(x.shape, y.shape) # print(np.stack((x, y), axis=2).shape) x, y = (np.array(a).transpose((0, 2, 1)).reshape((-1, 2)) for a in (x, y)) coll = LineCollection(np.stack((x, y), axis=2), zorder=zorder) From b2c6bed8fe8557b48a0d4a9d36d5434fb8da90f6 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Tue, 24 Feb 2026 17:39:04 +0100 Subject: [PATCH 12/22] add inelastic sample section to components notebooks --- docs/components.ipynb | 208 ++++++++++++++++++++++++++++++++++++++---- src/tof/inelastic.py | 9 +- src/tof/model.py | 2 + src/tof/result.py | 3 + src/tof/source.py | 20 ++++ 5 files changed, 223 insertions(+), 19 deletions(-) diff --git a/docs/components.ipynb b/docs/components.ipynb index 9e5a463..7a2185d 100644 --- a/docs/components.ipynb +++ b/docs/components.ipynb @@ -18,7 +18,10 @@ "metadata": {}, "outputs": [], "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", "import scipp as sc\n", + "import plopp as pp\n", "import tof\n", "\n", "meter = sc.Unit('m')\n", @@ -412,7 +415,183 @@ "id": "33", "metadata": {}, "source": [ - "## Loading from a JSON file\n", + "## Inelastic sample\n", + "\n", + "Placing an `InelasticSample` in the instrument will change the energy of the incoming neutrons by a $\\Delta E$ defined by a probability distribution function.\n", + "It defines the likeliness of what energy-shift will be applied to each neutron.\n", + "\n", + "To give an equal chance for a random $\\Delta E$ between -0.2 and 0.2 meV, we can create a flat distribution (all ones):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "sample = tof.InelasticSample(\n", + " distance=28.0 * meter,\n", + " name=\"sample\",\n", + " delta_e=sc.DataArray(\n", + " data=sc.ones(sizes={'e': 100}),\n", + " coords={'e': sc.linspace('e', -0.2, 0.2, 100, unit='meV')},\n", + " ),\n", + ")\n", + "sample" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": {}, + "source": [ + "We then make a single fast-rotating chopper with one small opening,\n", + "to select a narrow wavelength range at every rotation.\n", + "\n", + "We also add a monitor before the sample, and a detector after the sample so we can follow the changes in energies/wavelengths." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36", + "metadata": {}, + "outputs": [], + "source": [ + "choppers = [\n", + " tof.Chopper(\n", + " frequency=70.0 * Hz,\n", + " open=sc.array(dims=['cutout'], values=[0.0], unit='deg'),\n", + " close=sc.array(dims=['cutout'], values=[1.0], unit='deg'),\n", + " phase=0.0 * deg,\n", + " distance=20.0 * meter,\n", + " name=\"fastchopper\",\n", + " ),\n", + "]\n", + "\n", + "detectors = [\n", + " tof.Detector(distance=26.0 * meter, name='monitor'),\n", + " tof.Detector(distance=32.0 * meter, name='detector'),\n", + "]\n", + "\n", + "source = tof.Source(facility='ess', neutrons=5_000_000)\n", + "\n", + "model = tof.Model(source=source, components=choppers + detectors + [sample])\n", + "res = model.run()\n", + "\n", + "\n", + "fig, ax = plt.subplots(1, 2, figsize=(12, 4.5))\n", + "\n", + "dw = sc.scalar(0.1, unit='angstrom')\n", + "pp.plot(\n", + " {\n", + " 'monitor': res['monitor'].data.hist(wavelength=dw),\n", + " 'detector': res['detector'].data.hist(wavelength=dw),\n", + " },\n", + " title=\"With inelastic sample\",\n", + " xmin=4,\n", + " xmax=20,\n", + " ymin=-20,\n", + " ymax=400,\n", + " ax=ax[1],\n", + ")\n", + "\n", + "res.plot(visible_rays=10000, ax=ax[0])" + ] + }, + { + "cell_type": "markdown", + "id": "37", + "metadata": {}, + "source": [ + "### Non-uniform energy-transfer distributions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "38", + "metadata": {}, + "outputs": [], + "source": [ + "# Sample 1: double-peak at min and max\n", + "delta_e = sc.DataArray(\n", + " data=sc.zeros(sizes={'e': 100}),\n", + " coords={'e': sc.linspace('e', -0.2, 0.2, 100, unit='meV')},\n", + ")\n", + "delta_e.values[[0, -1]] = 1.0\n", + "sample1 = tof.InelasticSample(\n", + " distance=28.0 * meter,\n", + " name=\"sample\",\n", + " delta_e=delta_e,\n", + ")\n", + "\n", + "# Sample 2: normal distribution\n", + "x = sc.linspace('e', -0.2, 0.2, 100, unit='meV')\n", + "sig = sc.scalar(0.03, unit='meV')\n", + "y = 1.0 / (np.sqrt(2.0 * np.pi) * sig) * sc.exp(-((x / sig) ** 2) / 2)\n", + "y.unit = \"\"\n", + "\n", + "sample2 = tof.InelasticSample(\n", + " distance=28.0 * meter,\n", + " name=\"sample\",\n", + " delta_e=sc.DataArray(data=y, coords={'e': x}),\n", + ")\n", + "\n", + "sample1.plot() + sample2.plot()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "39", + "metadata": {}, + "outputs": [], + "source": [ + "model1 = tof.Model(source=source, components=choppers + detectors + [sample1])\n", + "model2 = tof.Model(source=source, components=choppers + detectors + [sample2])\n", + "\n", + "res1 = model1.run()\n", + "res2 = model2.run()\n", + "\n", + "fig, ax = plt.subplots(2, 2, figsize=(12, 9))\n", + "\n", + "res1.plot(ax=ax[0, 0], title=\"Sample 1\")\n", + "pp.plot(\n", + " {\n", + " 'monitor': res1['monitor'].data.hist(wavelength=dw),\n", + " 'detector': res1['detector'].data.hist(wavelength=dw),\n", + " },\n", + " title=\"Sample 1\",\n", + " xmin=4,\n", + " xmax=20,\n", + " ymin=-20,\n", + " ymax=400,\n", + " ax=ax[0, 1],\n", + ")\n", + "\n", + "res2.plot(ax=ax[1, 0], title=\"Sample 2\")\n", + "_ = pp.plot(\n", + " {\n", + " 'monitor': res2['monitor'].data.hist(wavelength=dw),\n", + " 'detector': res2['detector'].data.hist(wavelength=dw),\n", + " },\n", + " title=\"Sample 2\",\n", + " xmin=4,\n", + " xmax=20,\n", + " ymin=-20,\n", + " ymax=400,\n", + " ax=ax[1, 1],\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "40", + "metadata": {}, + "source": [ + "## Loading components from a JSON file\n", "\n", "It is also possible to load components from a JSON file,\n", "which can be very useful to quickly load a pre-configured instrument.\n", @@ -423,19 +602,14 @@ { "cell_type": "code", "execution_count": null, - "id": "34", + "id": "41", "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "params = {\n", - " \"source\": {\n", - " \"type\": \"source\",\n", - " \"facility\": \"ess\",\n", - " \"neutrons\": 1e6,\n", - " \"pulses\": 1\n", - " },\n", + " \"source\": {\"type\": \"source\", \"facility\": \"ess\", \"neutrons\": 1e6, \"pulses\": 1},\n", " \"chopper1\": {\n", " \"type\": \"chopper\",\n", " \"frequency\": {\"value\": 56.0, \"unit\": \"Hz\"},\n", @@ -470,7 +644,7 @@ { "cell_type": "code", "execution_count": null, - "id": "35", + "id": "42", "metadata": {}, "outputs": [], "source": [ @@ -479,7 +653,7 @@ }, { "cell_type": "markdown", - "id": "36", + "id": "43", "metadata": {}, "source": [ "We now use the `tof.Model.from_json()` method to load our instrument:" @@ -488,7 +662,7 @@ { "cell_type": "code", "execution_count": null, - "id": "37", + "id": "44", "metadata": {}, "outputs": [], "source": [ @@ -498,7 +672,7 @@ }, { "cell_type": "markdown", - "id": "38", + "id": "45", "metadata": {}, "source": [ "We can see that all components have been read in correctly.\n", @@ -508,7 +682,7 @@ { "cell_type": "code", "execution_count": null, - "id": "39", + "id": "46", "metadata": {}, "outputs": [], "source": [ @@ -517,7 +691,7 @@ }, { "cell_type": "markdown", - "id": "40", + "id": "47", "metadata": {}, "source": [ "### Modifying the source\n", @@ -531,7 +705,7 @@ { "cell_type": "code", "execution_count": null, - "id": "41", + "id": "48", "metadata": {}, "outputs": [], "source": [ @@ -542,7 +716,7 @@ }, { "cell_type": "markdown", - "id": "42", + "id": "49", "metadata": {}, "source": [ "## Saving to JSON\n", @@ -553,7 +727,7 @@ { "cell_type": "code", "execution_count": null, - "id": "43", + "id": "50", "metadata": {}, "outputs": [], "source": [ diff --git a/src/tof/inelastic.py b/src/tof/inelastic.py index 037f2b6..82061b1 100644 --- a/src/tof/inelastic.py +++ b/src/tof/inelastic.py @@ -46,8 +46,13 @@ def __getitem__( return replace(self, data=self.data[val]) def plot_on_time_distance_diagram(self, ax, tmax) -> None: - ax.plot([0, tmax], [self.distance.value] * 2, color="gray", lw=3) - ax.text(0, self.distance.value, self.name, ha="left", va="bottom", color="gray") + ax.plot([0, tmax], [self.distance.value] * 2, color="tab:brown", lw=4) + # ax.plot( + # [0, tmax], [self.distance.value] * 2, color="tab:cyan", lw=3, ls='dashed' + # ) + ax.text( + 0, self.distance.value, self.name, ha="left", va="bottom", color="tab:brown" + ) class InelasticSample(Component): diff --git a/src/tof/model.py b/src/tof/model.py index 46d0f99..09d54c0 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -49,6 +49,8 @@ def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]] # detectors = [] mapping = {"chopper": Chopper, "detector": Detector} for name, comp in instrument.items(): + if comp["type"] == "source": + continue if comp["type"] not in mapping: raise ValueError( f"Unknown component type: {comp['type']} for component {name}. " diff --git a/src/tof/result.py b/src/tof/result.py index ac77c90..80da564 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -167,6 +167,7 @@ def plot( seed: int | None = None, vmin: float | None = None, vmax: float | None = None, + title: str | None = None, ) -> Plot: """ Plot the time-distance diagram for the instrument, including the rays of @@ -278,6 +279,8 @@ def plot( inches = fig.get_size_inches() fig.set_size_inches((min(inches[0] * self.source.pulses, 12.0), inches[1])) fig.tight_layout() + if title is not None: + ax.set_title(title) return Plot(fig=fig, ax=ax) def __repr__(self) -> str: diff --git a/src/tof/source.py b/src/tof/source.py index 49d40cd..5dd0599 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -25,6 +25,18 @@ def _default_frequency(frequency: sc.Variable | None, pulses: int) -> sc.Variabl return frequency +def _bin_edges_to_midpoints( + da: sc.DataArray, dims: list[str] | tuple[str] +) -> sc.DataArray: + return da.assign_coords( + { + dim: sc.midpoints(da.coords[dim], dim) + for dim in dims + if da.coords.is_edges(dim) + } + ) + + def _make_pulses( neutrons: int, frequency: sc.Variable, @@ -408,6 +420,14 @@ def from_distribution( 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"]) + pulse_params = _make_pulses( neutrons=neutrons, p=p, From 047d6852c719111b1bd125d84a08024e6085651f Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Tue, 24 Feb 2026 18:06:09 +0100 Subject: [PATCH 13/22] cleanup --- src/tof/chopper.py | 3 - src/tof/detector.py | 1 - src/tof/inelastic.py | 36 +++++------ src/tof/model.py | 143 +------------------------------------------ src/tof/result.py | 43 +------------ src/tof/utils.py | 28 +++++++++ 6 files changed, 47 insertions(+), 207 deletions(-) diff --git a/src/tof/chopper.py b/src/tof/chopper.py index fce799e..2a50484 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -473,11 +473,8 @@ def apply( # Apply the chopper's open/close times to the data m = sc.zeros(sizes=neutrons.sizes, unit=None, dtype=bool) to, tc = self.open_close_times(time_limit=time_limit) - # neutrons.update({'open_times': to, 'close_times': tc}) for i in range(len(to)): m |= (neutrons.coords['toa'] > to[i]) & (neutrons.coords['toa'] < tc[i]) - # combined = initial_mask & m - # data_at_comp.masks['blocked_by_others'] = ~initial_mask neutrons.masks['blocked_by_me'] = (~m) & (~neutrons.masks['blocked_by_others']) return neutrons, self.as_readonly(neutrons, time_limit=time_limit) diff --git a/src/tof/detector.py b/src/tof/detector.py index 5e55275..bf9b136 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -122,5 +122,4 @@ def apply( time_limit: The time limit for the neutrons to be considered as reaching the detector. """ - # neutrons.pop("blocked_by_me", None) return neutrons, self.as_readonly(neutrons) diff --git a/src/tof/inelastic.py b/src/tof/inelastic.py index 82061b1..5cf2102 100644 --- a/src/tof/inelastic.py +++ b/src/tof/inelastic.py @@ -9,7 +9,12 @@ import scipp as sc from .component import Component, ComponentReading -from .utils import var_to_dict, wavelength_to_speed +from .utils import ( + energy_to_wavelength, + var_to_dict, + wavelength_to_energy, + wavelength_to_speed, +) @dataclass(frozen=True) @@ -47,9 +52,6 @@ def __getitem__( def plot_on_time_distance_diagram(self, ax, tmax) -> None: ax.plot([0, tmax], [self.distance.value] * 2, color="tab:brown", lw=4) - # ax.plot( - # [0, tmax], [self.distance.value] * 2, color="tab:cyan", lw=3, ls='dashed' - # ) ax.text( 0, self.distance.value, self.name, ha="left", va="bottom", color="tab:brown" ) @@ -66,6 +68,13 @@ class InelasticSample(Component): The distance from the source to the inelastic sample. name: The name of the inelastic sample. + delta_e: + The change in energy of the neutrons when they pass through the inelastic + sample. The values of the array represent the probability of a neutron to have + its energy changed by the corresponding amount in the coordinates. The + coordinate values should be in energy units, and the array should be 1D. + seed: + The seed for the random number generator used to apply the energy change. """ def __init__( @@ -96,11 +105,6 @@ def __repr__(self) -> str: def plot(self, **kwargs) -> pp.FigureLike: return pp.xyplot(self.energies, self.probabilities, **kwargs) - # def __eq__(self, other: object) -> bool: - # if not isinstance(other, Detector): - # return NotImplemented - # return self.name == other.name and sc.identical(self.distance, other.distance) - def as_dict(self) -> dict: """ Return the inelastic sample as a dictionary. @@ -154,9 +158,7 @@ def apply( The time limit for the neutrons to be considered as reaching the inelastic sample. """ - # neutrons.pop("blocked_by_me", None) - # return neutrons, self.as_readonly(neutrons) - w_initial = sc.reciprocal(neutrons.coords["wavelength"] ** 2) + w_initial = neutrons.coords["wavelength"] n = neutrons.shape inds = self._rng.choice(len(self.energies), size=n, p=self.probabilities) @@ -167,14 +169,8 @@ def apply( unit=self.energies.unit, ) # Convert energy change to wavelength change - # w_initial = sc.reciprocal(neutrons.coords["wavelength"] ** 2) - w_final = sc.reciprocal( - sc.sqrt( - ((2 * sc.constants.m_n / (sc.constants.h**2)) * de).to( - unit=w_initial.unit - ) - + w_initial - ) + w_final = energy_to_wavelength( + wavelength_to_energy(w_initial, unit=de.unit) + de, unit=w_initial.unit ) neutrons = neutrons.assign_coords( wavelength=w_final, speed=wavelength_to_speed(w_final) diff --git a/src/tof/model.py b/src/tof/model.py index 09d54c0..1bf4895 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -4,13 +4,11 @@ from __future__ import annotations import warnings -from collections import defaultdict from itertools import chain import scipp as sc -# from tof.component import Component -from .chopper import AntiClockwise, Chopper, Clockwise +from .chopper import Chopper from .component import Component from .detector import Detector from .result import Result @@ -19,16 +17,6 @@ ComponentType = Chopper | Detector -# def _array_or_none(results: dict, key: str) -> sc.Variable | None: -# return ( -# sc.array( -# dims=["cutout"], values=results[key]["value"], unit=results[key]["unit"] -# ) -# if key in results -# else None -# ) - - def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]]: """ Create choppers and detectors from a dictionary. @@ -46,7 +34,6 @@ def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]] classes for details. """ components = [] - # detectors = [] mapping = {"chopper": Chopper, "detector": Detector} for name, comp in instrument.items(): if comp["type"] == "source": @@ -56,48 +43,6 @@ def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]] f"Unknown component type: {comp['type']} for component {name}. " ) components.append(mapping[comp["type"]].from_json(name=name, params=comp)) - # component_class = comp_mapping[comp["type"]] - # if comp["type"] == "chopper": - # direction = comp["direction"].lower() - # if direction == "clockwise": - # _dir = Clockwise - # elif any(x in direction for x in ("anti", "counter")): - # _dir = AntiClockwise - # else: - # raise ValueError( - # f"Chopper direction must be 'clockwise' or 'anti-clockwise', got " - # f"'{comp['direction']}' for component {name}." - # ) - # choppers.append( - # Chopper( - # frequency=comp["frequency"]["value"] - # * sc.Unit(comp["frequency"]["unit"]), - # direction=_dir, - # open=_array_or_none(comp, "open"), - # close=_array_or_none(comp, "close"), - # centers=_array_or_none(comp, "centers"), - # widths=_array_or_none(comp, "widths"), - # phase=comp["phase"]["value"] * sc.Unit(comp["phase"]["unit"]), - # distance=comp["distance"]["value"] - # * sc.Unit(comp["distance"]["unit"]), - # name=name, - # ) - # ) - # elif comp["type"] == "detector": - # detectors.append( - # Detector( - # distance=comp["distance"]["value"] - # * sc.Unit(comp["distance"]["unit"]), - # name=name, - # ) - # ) - # elif comp["type"] == "source": - # continue - # else: - # raise ValueError( - # f"Unknown component type: {comp['type']} for component {name}. " - # "Supported types are 'chopper', 'detector', and 'source'." - # ) return components @@ -127,40 +72,11 @@ def __init__( choppers: list[Chopper] | tuple[Chopper, ...] | None = None, detectors: list[Detector] | tuple[Detector, ...] | None = None, ): - # self.choppers = {} - # self.detectors = {} self.source = source self.components = {} - # for components, kind in ((choppers, Chopper), (detectors, Detector)): for comp in chain((choppers or ()), (detectors or ()), (components or ())): self.add(comp) - # @property - # def choppers(self) -> dict[str, Chopper]: - # """ - # Return a dictionary of all choppers in the model. - # This is meant for retro-compatibility with older versions of the code, and will - # be removed in a future version. - # """ - # return { - # name: comp - # for name, comp in self.components.items() - # if comp.kind == "chopper" - # } - - # @property - # def detectors(self) -> dict[str, Detector]: - # """ - # Return a dictionary of all detectors in the model. - # This is meant for retro-compatibility with older versions of the code, and will - # be removed in a future version. - # """ - # return { - # name: comp - # for name, comp in self.components.items() - # if comp.kind == "detector" - # } - @classmethod def from_json(cls, filename: str) -> Model: """ @@ -245,11 +161,6 @@ def add(self, component: Component) -> None: component: A chopper or detector. """ - # if not isinstance(component, (Chopper | Detector)): - # raise TypeError( - # f"Cannot add component of type {type(component)} to the model. " - # "Only Chopper and Detector instances are allowed." - # ) # Note that the name "source" is reserved for the source. if component.name in (*self.components, "source"): raise KeyError( @@ -257,7 +168,6 @@ def add(self, component: Component) -> None: "If you wish to replace/update an existing component, use " "``model.components['name'] = new_component``." ) - # results = self.choppers if isinstance(component, Chopper) else self.detectors self.components[component.name] = component def remove(self, name: str): @@ -270,23 +180,6 @@ def remove(self, name: str): The name of the component to remove. """ del self.components[name] - # if name in self.choppers: - # del self.choppers[name] - # elif name in self.detectors: - # del self.detectors[name] - # else: - # raise KeyError(f"No component with name {name} was found.") - - # def __iter__(self): - # return chain(self.choppers, self.detectors) - - # def __getitem__(self, name) -> Chopper | Detector: - # if name not in self: - # raise KeyError(f"No component with name {name} was found.") - # return self.choppers[name] if name in self.choppers else self.detectors[name] - - # def __delitem__(self, name) -> None: - # self.remove(name) def run(self) -> Result: """ @@ -308,9 +201,6 @@ def run(self) -> Result: "itself. Please check the distances of the components." ) - # birth_time = self.source.data.coords['birth_time'] - # speed = self.source.data.coords['speed'] - # neutrons = sc.DataGroup(self.source.data.coords) neutrons = self.source.data.copy(deep=False) neutrons.masks["blocked_by_others"] = sc.zeros( sizes=neutrons.sizes, unit=None, dtype=bool @@ -322,8 +212,6 @@ def run(self) -> Result: time_unit = neutrons.coords['birth_time'].unit readings = {} - # result_choppers = {} - # result_detectors = {} time_limit = ( neutrons.coords['birth_time'] + ( @@ -332,13 +220,7 @@ def run(self) -> Result: ).to(unit=time_unit) ).max() for comp in components: - # results = result_detectors if isinstance(c, Detector) else result_choppers - # results[comp.name] = comp.as_dict() neutrons = neutrons.copy(deep=False) - # tof = ((c.distance - self.source.distance) / neutrons.coords['speed']).to( - # unit=time_unit, copy=False - # ) - # t = birth_time + tof toa = neutrons.coords['toa'] + ( (comp.distance - neutrons.coords['distance']) / neutrons.coords['speed'] ).to(unit=time_unit, copy=False) @@ -356,34 +238,11 @@ def run(self) -> Result: ] | neutrons.masks.pop('blocked_by_me') neutrons, reading = comp.apply(neutrons=neutrons, time_limit=time_limit) - readings[comp.name] = reading - # # data_at_comp.coords['tof'] = tof - # if isinstance(c, Detector): - # data_at_comp.masks['blocked_by_others'] = ~initial_mask - # continue - # m = sc.zeros(sizes=t.sizes, unit=None, dtype=bool) - # to, tc = c.open_close_times(time_limit=time_limit) - # results[c.name].update({'open_times': to, 'close_times': tc}) - # for i in range(len(to)): - # m |= (t > to[i]) & (t < tc[i]) - # combined = initial_mask & m - # data_at_comp.masks['blocked_by_others'] = ~initial_mask - # data_at_comp.masks['blocked_by_me'] = ~m & initial_mask - # initial_mask = combined - # neutrons = data_at_comp - return Result(source=self.source.as_readonly(), readings=readings) def __repr__(self) -> str: - # out = f"Model:\n Source: {self.source}\n Choppers:\n" - # for name, ch in self.choppers.items(): - # out += f" {name}: {ch}\n" - # out += " Detectors:\n" - # for name, det in self.detectors.items(): - # out += f" {name}: {det}\n" - # return out out = f"Model:\n Source: {self.source}\n" groups = {} for comp in self.components.values(): diff --git a/src/tof/result.py b/src/tof/result.py index 80da564..76df04d 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -11,9 +11,9 @@ import scipp as sc from matplotlib.collections import LineCollection -from .chopper import Chopper, ChopperReading +from .chopper import ChopperReading from .component import ComponentReading -from .detector import Detector, DetectorReading +from .detector import DetectorReading from .source import SourceReading from .utils import Plot, one_mask @@ -38,18 +38,6 @@ def _get_rays( xstart, ystart = xend, yend color = comp.data["pulse", pulse].coords["wavelength"].values[inds] - # comp_data = compo.data["pulse", pulse] - # x.append(comp_data.coords["toa"].values[inds]) - # y.append(np.full_like(x[-1], comp.distance.value)) - # c.append(comp_data.coords["wavelength"].values[inds]) - # x = np.stack( - # [comp.data["pulse", pulse].coords["toa"].values[inds] for comp in components], - # axis=1, - # ) - # y = np.stack( - # [np.full_like(x[:, 0], comp.distance.value) for comp in components], - # axis=1, - # ) return np.array(x), np.array(y), np.array(c) @@ -65,8 +53,6 @@ def _add_rays( cax: plt.Axes | None = None, zorder: int = 1, ): - # print(x.shape, y.shape) - # print(np.stack((x, y), axis=2).shape) x, y = (np.array(a).transpose((0, 2, 1)).reshape((-1, 2)) for a in (x, y)) coll = LineCollection(np.stack((x, y), axis=2), zorder=zorder) if isinstance(color, str): @@ -97,28 +83,6 @@ class Result: def __init__(self, source: SourceReading, readings: dict[str, dict]): self._source = source self._components = MappingProxyType(readings) - # self._choppers = {} - # for name, chopper in choppers.items(): - # self._choppers[name] = ChopperReading( - # distance=chopper["distance"], - # name=chopper["name"], - # frequency=chopper["frequency"], - # open=chopper["open"], - # close=chopper["close"], - # phase=chopper["phase"], - # open_times=chopper["open_times"], - # close_times=chopper["close_times"], - # data=chopper["data"], - # ) - - # self._detectors = {} - # for name, det in detectors.items(): - # self._detectors[name] = DetectorReading( - # distance=det["distance"], name=det["name"], data=det["data"] - # ) - - # self._choppers = MappingProxyType(self._choppers) - # self._detectors = MappingProxyType(self._detectors) @property def choppers(self) -> MappingProxyType[str, ChopperReading]: @@ -151,8 +115,6 @@ def __iter__(self): return iter(self._components) def __getitem__(self, name: str) -> ComponentReading: - # if name not in self: - # raise KeyError(f"No component with name {name} was found.") return self._components[name] def plot( @@ -232,7 +194,6 @@ def plot( replace=False, ) x, y, c = _get_rays(components, pulse=i, inds=inds) - print(x.shape, y.shape, c.shape) _add_rays( ax=ax, x=x, diff --git a/src/tof/utils.py b/src/tof/utils.py index 640d121..a46c42d 100644 --- a/src/tof/utils.py +++ b/src/tof/utils.py @@ -69,6 +69,34 @@ def energy_to_speed(x: sc.Variable, unit="m/s") -> sc.Variable: return sc.sqrt(x / (0.5 * const.m_n)).to(unit=unit) +def wavelength_to_energy(x: sc.Variable, unit="meV") -> sc.Variable: + """ + Convert neutron wavelengths to energies. + + Parameters + ---------- + x: + Input wavelengths. + unit: + The unit of the output energies. + """ + return speed_to_energy(wavelength_to_speed(x)).to(unit=unit) + + +def energy_to_wavelength(x: sc.Variable, unit="angstrom") -> sc.Variable: + """ + Convert neutron energies to wavelengths. + + Parameters + ---------- + x: + Input energies. + unit: + The unit of the output wavelengths. + """ + return speed_to_wavelength(energy_to_speed(x)).to(unit=unit) + + def one_mask( masks: MappingProxyType[str, sc.Variable], unit: str | None = None ) -> sc.Variable: From 87fefc43a882a2c7ac8d4e22832c8619ed680ce0 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 00:06:48 +0100 Subject: [PATCH 14/22] fix plotting blocked rays --- src/tof/result.py | 42 ++++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/src/tof/result.py b/src/tof/result.py index 76df04d..e99a714 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -25,7 +25,6 @@ def _get_rays( y = [] c = [] data = components[0].data["pulse", pulse] - # TODO optimize: should not index multiple times xstart = data.coords["toa"].values[inds] ystart = np.full_like(xstart, components[0].distance.value) color = data.coords["wavelength"].values[inds] @@ -38,7 +37,11 @@ def _get_rays( xstart, ystart = xend, yend color = comp.data["pulse", pulse].coords["wavelength"].values[inds] - return np.array(x), np.array(y), np.array(c) + return ( + np.array(x).transpose((0, 2, 1)), + np.array(y).transpose((0, 2, 1)), + np.array(c), + ) def _add_rays( @@ -53,7 +56,7 @@ def _add_rays( cax: plt.Axes | None = None, zorder: int = 1, ): - x, y = (np.array(a).transpose((0, 2, 1)).reshape((-1, 2)) for a in (x, y)) + x, y = (a.reshape((-1, 2)) for a in (x, y)) coll = LineCollection(np.stack((x, y), axis=2), zorder=zorder) if isinstance(color, str): coll.set_color(color) @@ -206,21 +209,24 @@ def plot( cax=cax, ) - # # Plot blocked rays - # inds = rng.choice( - # ids[blocked], size=min(blocked_rays, nblocked), replace=False - # ) - # x, y = _get_rays(components, pulse=i, inds=inds) - # blocked_by_others = np.stack( - # [ - # comp.data["pulse", i].masks["blocked_by_others"].values[inds] - # for comp in components - # ], - # axis=1, - # ) - # x[blocked_by_others] = np.nan - # y[blocked_by_others] = np.nan - # _add_rays(ax=ax, x=x, y=y, color="lightgray", zorder=-1) + # Plot blocked rays + inds = rng.choice( + ids[blocked], size=min(blocked_rays, nblocked), replace=False + ) + x, y, _ = _get_rays(components, pulse=i, inds=inds) + blocked_by_others = np.stack( + [ + comp.data["pulse", i].masks["blocked_by_others"].values[inds] + for comp in components[1:] + ], + axis=1, + ).T + line_selection = np.broadcast_to( + blocked_by_others.reshape((*blocked_by_others.shape, 1)), x.shape + ) + x[line_selection] = np.nan + y[line_selection] = np.nan + _add_rays(ax=ax, x=x, y=y, color="lightgray", zorder=-1) # Plot pulse self.source.plot_on_time_distance_diagram(ax, pulse=i) From a1997845643cb65326d043a181cd0584b455f4b1 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 10:25:08 +0100 Subject: [PATCH 15/22] start fixing tests --- src/tof/model.py | 78 +++++++++++++++------- src/tof/result.py | 40 ++++++------ src/tof/source.py | 12 ++++ src/tof/utils.py | 18 +++++ tests/model_test.py | 155 ++++++++++++++++++++------------------------ 5 files changed, 176 insertions(+), 127 deletions(-) diff --git a/src/tof/model.py b/src/tof/model.py index 1bf4895..a169144 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -5,6 +5,7 @@ import warnings from itertools import chain +from types import MappingProxyType import scipp as sc @@ -13,6 +14,7 @@ from .detector import Detector from .result import Result from .source import Source +from .utils import extract_component_group ComponentType = Chopper | Detector @@ -33,17 +35,25 @@ def make_beamline(instrument: dict) -> dict[str, list[Chopper] | list[Detector]] type, see the documentation of the :class:`Chopper` and :class:`Detector` classes for details. """ - components = [] + beamline = {"components": []} mapping = {"chopper": Chopper, "detector": Detector} for name, comp in instrument.items(): if comp["type"] == "source": + if "source" in beamline: + raise ValueError( + "Only one source is allowed, but multiple were found in the" + "instrument parameters." + ) + beamline["source"] = Source.from_json(params=comp) continue if comp["type"] not in mapping: raise ValueError( f"Unknown component type: {comp['type']} for component {name}. " ) - components.append(mapping[comp["type"]].from_json(name=name, params=comp)) - return components + beamline["components"].append( + mapping[comp["type"]].from_json(name=name, params=comp) + ) + return beamline class Model: @@ -73,10 +83,38 @@ def __init__( detectors: list[Detector] | tuple[Detector, ...] | None = None, ): self.source = source - self.components = {} + self._components = {} for comp in chain((choppers or ()), (detectors or ()), (components or ())): self.add(comp) + @property + def components(self) -> dict[str, Component]: + """ + A dictionary of the components in the instrument. + """ + return self._components + + @property + def choppers(self) -> MappingProxyType[str, Chopper]: + """ + A dictionary of the choppers in the instrument. + """ + return extract_component_group(self._components, "chopper") + + @property + def detectors(self) -> MappingProxyType[str, Detector]: + """ + A dictionary of the detectors in the instrument. + """ + return extract_component_group(self._components, "detector") + + @property + def samples(self) -> MappingProxyType[str, Component]: + """ + A dictionary of the samples in the instrument. + """ + return extract_component_group(self._components, "sample") + @classmethod def from_json(cls, filename: str) -> Model: """ @@ -95,20 +133,7 @@ def from_json(cls, filename: str) -> Model: with open(filename) as f: instrument = json.load(f) - beamline = make_beamline(instrument) - source = None - for item in instrument.values(): - if item.get("type") == "source": - if "facility" not in item: - raise ValueError( - "Currently, only sources from facilities are supported when " - "loading from JSON." - ) - source_args = item.copy() - del source_args["type"] - source = Source(**source_args) - break - return cls(source=source, **beamline) + return cls(**make_beamline(instrument)) def as_json(self) -> dict: """ @@ -128,7 +153,7 @@ def as_json(self) -> dict: ) else: instrument_dict['source'] = self.source.as_json() - for comp in self.components.values(): + for comp in self._components.values(): instrument_dict[comp.name] = comp.as_json() return instrument_dict @@ -161,14 +186,19 @@ def add(self, component: Component) -> None: component: A chopper or detector. """ + if not isinstance(component, Component): + raise TypeError( + "Component must be an instance of Component or derived class, " + f"but got {type(component)}." + ) # Note that the name "source" is reserved for the source. - if component.name in (*self.components, "source"): + if component.name in (*self._components, "source"): raise KeyError( f"Component with name {component.name} already exists. " "If you wish to replace/update an existing component, use " "``model.components['name'] = new_component``." ) - self.components[component.name] = component + self._components[component.name] = component def remove(self, name: str): """ @@ -179,7 +209,7 @@ def remove(self, name: str): name: The name of the component to remove. """ - del self.components[name] + del self._components[name] def run(self) -> Result: """ @@ -190,7 +220,7 @@ def run(self) -> Result: "No source has been defined for this model. Please add a source using " "`model.source = Source(...)` before running the simulation." ) - components = sorted(self.components.values(), key=lambda c: c.distance.value) + components = sorted(self._components.values(), key=lambda c: c.distance.value) if len(components) == 0: raise ValueError("Cannot run model: no components have been defined.") @@ -245,7 +275,7 @@ def run(self) -> Result: def __repr__(self) -> str: out = f"Model:\n Source: {self.source}\n" groups = {} - for comp in self.components.values(): + for comp in self._components.values(): if comp.kind not in groups: groups[comp.kind] = [] groups[comp.kind].append(comp) diff --git a/src/tof/result.py b/src/tof/result.py index e99a714..4b8914d 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -15,7 +15,7 @@ from .component import ComponentReading from .detector import DetectorReading from .source import SourceReading -from .utils import Plot, one_mask +from .utils import Plot, extract_component_group, one_mask def _get_rays( @@ -89,25 +89,24 @@ def __init__(self, source: SourceReading, readings: dict[str, dict]): @property def choppers(self) -> MappingProxyType[str, ChopperReading]: - """The choppers in the model.""" - return MappingProxyType( - { - key: comp - for key, comp in self._components.items() - if comp.kind == "chopper" - } - ) + """ + A dictionary of the choppers in the instrument. + """ + return extract_component_group(self._components, "chopper") @property def detectors(self) -> MappingProxyType[str, DetectorReading]: - """The detectors in the model.""" - return MappingProxyType( - { - key: comp - for key, comp in self._components.items() - if comp.kind == "detector" - } - ) + """ + A dictionary of the detectors in the instrument. + """ + return extract_component_group(self._components, "detector") + + @property + def samples(self) -> MappingProxyType[str, ComponentReading]: + """ + A dictionary of the samples in the instrument. + """ + return extract_component_group(self._components, "sample") @property def source(self) -> SourceReading: @@ -284,11 +283,12 @@ def to_nxevent_data(self, key: str | None = None) -> sc.DataArray: start = sc.datetime("2024-01-01T12:00:00.000000") period = sc.reciprocal(self.source.frequency) - keys = list(self._detectors.keys()) if key is None else [key] + detectors = self.detectors + keys = list(detectors.keys()) if key is None else [key] event_data = [] for name in keys: - raw_data = self._detectors[name].data.flatten(to="event") + raw_data = detectors[name].data.flatten(to="event") events = ( raw_data[~raw_data.masks["blocked_by_others"]] .copy() @@ -307,7 +307,7 @@ def to_nxevent_data(self, key: str | None = None) -> sc.DataArray: "toa" ) % period.to(unit=dt.unit) out = ( - event_data.drop_coords(["tof", "speed", "birth_time", "wavelength"]) + event_data.drop_coords(["speed", "birth_time", "wavelength"]) .group("distance") .rename_dims(distance="detector_number") ) diff --git a/src/tof/source.py b/src/tof/source.py index 5dd0599..413a1e5 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2025 Scipp contributors (https://github.com/scipp) +from __future__ import annotations import warnings from dataclasses import dataclass @@ -492,6 +493,17 @@ def __repr__(self) -> str: f" distance={self.distance:c}" ) + @classmethod + def from_json(cls, params: dict) -> Source: + if params.get("facility") is None: + raise ValueError( + "Currently, only sources from facilities are supported when " + "loading from JSON." + ) + source_args = params.copy() + del source_args["type"] + return cls(**source_args) + def as_json(self) -> dict: """ Return the source as a JSON-serializable dictionary. diff --git a/src/tof/utils.py b/src/tof/utils.py index a46c42d..822fa4e 100644 --- a/src/tof/utils.py +++ b/src/tof/utils.py @@ -130,6 +130,24 @@ def var_to_dict(var: sc.Variable) -> dict: } +def extract_component_group( + components: dict | MappingProxyType, kind: str +) -> MappingProxyType: + """ + Extract a group of components of a given kind from a dictionary of components. + + Parameters + ---------- + components: + The components to extract from. + kind: + The kind of components to extract. + """ + return MappingProxyType( + {name: comp for name, comp in components.items() if kind in comp.kind} + ) + + @dataclass class Plot: ax: plt.Axes diff --git a/tests/model_test.py b/tests/model_test.py index 11e9005..3a91f7d 100644 --- a/tests/model_test.py +++ b/tests/model_test.py @@ -17,7 +17,8 @@ ms = sc.Unit('ms') -def test_one_chopper_one_opening(make_chopper, make_source): +@pytest.mark.parametrize("use_components", [True, False]) +def test_one_chopper_one_opening(make_chopper, make_source, use_components): # Make a chopper open from 10-20 ms. Assume zero phase. topen = 10.0 * ms tclose = 20.0 * ms @@ -40,7 +41,11 @@ def test_one_chopper_one_opening(make_chopper, make_source): distance=chopper.distance, ) - model = tof.Model(source=source, choppers=[chopper], detectors=[detector]) + if use_components: + args = {"components": [chopper, detector]} + else: + args = {"choppers": [chopper], "detectors": [detector]} + model = tof.Model(source=source, **args) res = model.run() toa = res.choppers['chopper'].toa.data @@ -71,7 +76,8 @@ def test_one_chopper_one_opening(make_chopper, make_source): ) -def test_two_choppers_one_opening(make_chopper, make_source): +@pytest.mark.parametrize("use_components", [True, False]) +def test_two_choppers_one_opening(make_chopper, make_source, use_components): # Make a first chopper open from 5-16 ms. Assume zero phase. topen = 5.0 * ms tclose = 16.0 * ms @@ -105,9 +111,11 @@ def test_two_choppers_one_opening(make_chopper, make_source): distance=chopper1.distance, ) - model = tof.Model( - source=source, choppers=[chopper1, chopper2], detectors=[detector] - ) + if use_components: + args = {"components": [chopper1, chopper2, detector]} + else: + args = {"choppers": [chopper1, chopper2], "detectors": [detector]} + model = tof.Model(source=source, **args) res = model.run() ch1_toas = res.choppers['chopper1'].toa.data @@ -158,7 +166,8 @@ def test_two_choppers_one_opening(make_chopper, make_source): ) -def test_two_choppers_one_and_two_openings(make_chopper, make_source): +@pytest.mark.parametrize("use_components", [True, False]) +def test_two_choppers_one_and_two_openings(make_chopper, make_source, use_components): topen = 5.0 * ms tclose = 16.0 * ms chopper1 = make_chopper( @@ -201,9 +210,11 @@ def test_two_choppers_one_and_two_openings(make_chopper, make_source): distance=chopper1.distance, ) - model = tof.Model( - source=source, choppers=[chopper1, chopper2], detectors=[detector] - ) + if use_components: + args = {"components": [chopper1, chopper2, detector]} + else: + args = {"choppers": [chopper1, chopper2], "detectors": [detector]} + model = tof.Model(source=source, **args) res = model.run() assert res.choppers['chopper1'].toa.data.sum().value == 5 @@ -216,7 +227,8 @@ def test_two_choppers_one_and_two_openings(make_chopper, make_source): ) -def test_neutron_conservation(make_chopper): +@pytest.mark.parametrize("use_components", [True, False]) +def test_neutron_conservation(make_chopper, use_components): N = 100_000 source = tof.Source(facility='ess', neutrons=N) @@ -238,9 +250,11 @@ def test_neutron_conservation(make_chopper): ) detector = tof.Detector(distance=20 * meter, name='detector') - model = tof.Model( - source=source, choppers=[chopper1, chopper2], detectors=[detector] - ) + if use_components: + args = {"components": [chopper1, chopper2, detector]} + else: + args = {"choppers": [chopper1, chopper2], "detectors": [detector]} + model = tof.Model(source=source, **args) res = model.run() ch1 = res.choppers['chopper1'].toa.data @@ -258,40 +272,40 @@ def test_neutron_conservation(make_chopper): assert det.sum().value + det.masks['blocked_by_others'].sum().value == N -def test_neutron_time_of_flight(make_chopper): - N = 10_000 - source = tof.Source(facility='ess', neutrons=N) - - chopper1 = make_chopper( - topen=[5.0 * ms], - tclose=[16.0 * ms], - f=10.0 * Hz, - phase=0.0 * deg, - distance=10 * meter, - name='chopper1', - ) - chopper2 = make_chopper( - topen=[9.0 * ms, 15.0 * ms], - tclose=[15.0 * ms, 20.0 * ms], - f=15.0 * Hz, - phase=0.0 * deg, - distance=15 * meter, - name='chopper2', - ) - - detector = tof.Detector(distance=20 * meter, name='detector') - model = tof.Model( - source=source, choppers=[chopper1, chopper2], detectors=[detector] - ) - res = model.run() - - ch1 = res.choppers['chopper1'].data - ch2 = res.choppers['chopper2'].data - det = res.detectors['detector'].data - - assert sc.allclose(ch1.coords['tof'], ch1.coords['toa'] - ch1.coords['birth_time']) - assert sc.allclose(ch2.coords['tof'], ch2.coords['toa'] - ch2.coords['birth_time']) - assert sc.allclose(det.coords['tof'], det.coords['toa'] - det.coords['birth_time']) +# def test_neutron_time_of_flight(make_chopper): +# N = 10_000 +# source = tof.Source(facility='ess', neutrons=N) + +# chopper1 = make_chopper( +# topen=[5.0 * ms], +# tclose=[16.0 * ms], +# f=10.0 * Hz, +# phase=0.0 * deg, +# distance=10 * meter, +# name='chopper1', +# ) +# chopper2 = make_chopper( +# topen=[9.0 * ms, 15.0 * ms], +# tclose=[15.0 * ms, 20.0 * ms], +# f=15.0 * Hz, +# phase=0.0 * deg, +# distance=15 * meter, +# name='chopper2', +# ) + +# detector = tof.Detector(distance=20 * meter, name='detector') +# model = tof.Model( +# source=source, choppers=[chopper1, chopper2], detectors=[detector] +# ) +# res = model.run() + +# ch1 = res.choppers['chopper1'].data +# ch2 = res.choppers['chopper2'].data +# det = res.detectors['detector'].data + +# assert sc.allclose(ch1.coords['tof'], ch1.coords['toa'] - ch1.coords['birth_time']) +# assert sc.allclose(ch2.coords['tof'], ch2.coords['toa'] - ch2.coords['birth_time']) +# assert sc.allclose(det.coords['tof'], det.coords['toa'] - det.coords['birth_time']) def test_source_not_at_origin(make_chopper): @@ -405,50 +419,25 @@ def test_remove(dummy_chopper, dummy_detector, dummy_source): chopper = dummy_chopper detector = dummy_detector model = tof.Model(source=dummy_source, choppers=[chopper], detectors=[detector]) - del model['dummy_chopper'] - assert 'dummy_chopper' not in model - assert 'dummy_detector' in model - del model['dummy_detector'] - assert 'dummy_detector' not in model - - -def test_getitem(dummy_chopper, dummy_detector, dummy_source): - chopper = dummy_chopper - detector = dummy_detector - model = tof.Model(source=dummy_source, choppers=[chopper], detectors=[detector]) - assert model['dummy_chopper'] is chopper - assert model['dummy_detector'] is detector - with pytest.raises(KeyError, match='No component with name foo'): - model['foo'] + model.remove('dummy_chopper') + assert 'dummy_chopper' not in model.components + assert 'dummy_detector' in model.components + model.remove('dummy_detector') + assert 'dummy_detector' not in model.components def test_bad_input_type_raises(dummy_chopper, dummy_detector, dummy_source): chopper = dummy_chopper detector = dummy_detector - with pytest.raises( - TypeError, match='Beamline components: expected Chopper instance' - ): + err = "Component must be an instance of Component or derived class" + with pytest.raises(TypeError, match=err): _ = tof.Model(source=dummy_source, choppers='bad chopper') - with pytest.raises( - TypeError, match='Beamline components: expected Detector instance' - ): + with pytest.raises(TypeError, match=err): _ = tof.Model(source=dummy_source, choppers=[chopper], detectors='abc') - with pytest.raises( - TypeError, match='Beamline components: expected Chopper instance' - ): + with pytest.raises(TypeError, match=err): _ = tof.Model(source=dummy_source, choppers=[chopper, 'bad chopper']) - with pytest.raises( - TypeError, match='Beamline components: expected Detector instance' - ): + with pytest.raises(TypeError, match=err): _ = tof.Model(source=dummy_source, detectors=(1234, detector)) - with pytest.raises( - TypeError, match='Beamline components: expected Chopper instance' - ): - _ = tof.Model(source=dummy_source, choppers=[detector]) - with pytest.raises( - TypeError, match='Beamline components: expected Detector instance' - ): - _ = tof.Model(source=dummy_source, detectors=[chopper]) def test_model_repr_does_not_raise(make_chopper): From acb156b79b09d99a71c716b310efc6e3327fe015 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 10:37:53 +0100 Subject: [PATCH 16/22] fix remaining tests --- src/tof/result.py | 3 ++- tests/result_test.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/tof/result.py b/src/tof/result.py index 4b8914d..cdb1bf1 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -181,10 +181,11 @@ def plot( wmin, wmax = wavelengths.min(), wavelengths.max() rng = np.random.default_rng(seed) + # Make ids for neutrons per pulse, instead of using their id coord + ids = np.arange(self.source.neutrons) for i in range(self._source.data.sizes["pulse"]): component_data = furthest_component.data["pulse", i] - ids = component_data.coords["id"].values # Plot visible rays blocked = one_mask(component_data.masks).values diff --git a/tests/result_test.py b/tests/result_test.py index f72872a..7b9f8b9 100644 --- a/tests/result_test.py +++ b/tests/result_test.py @@ -343,7 +343,7 @@ def test_plot_reading_pulse_skipping_does_not_raise(): distance=10 * meter, name='skip', ) - model.choppers['skip'] = skip + model.components['skip'] = skip res = model.run() res.choppers['chopper'].toa.plot() res.choppers['chopper'].wavelength.plot() @@ -361,7 +361,7 @@ def test_plot_reading_nothing_to_plot_raises(): distance=10 * meter, name='skip', ) - model.choppers['skip'] = skip + model.components['skip'] = skip res = model.run() with pytest.raises(RuntimeError, match="Nothing to plot."): res.detectors['detector'].toa.plot() From d10610fae12fcf046f0cf3c9de2a4131558b094f Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 10:38:39 +0100 Subject: [PATCH 17/22] remove commented test --- tests/model_test.py | 36 ------------------------------------ 1 file changed, 36 deletions(-) diff --git a/tests/model_test.py b/tests/model_test.py index 3a91f7d..15c8733 100644 --- a/tests/model_test.py +++ b/tests/model_test.py @@ -272,42 +272,6 @@ def test_neutron_conservation(make_chopper, use_components): assert det.sum().value + det.masks['blocked_by_others'].sum().value == N -# def test_neutron_time_of_flight(make_chopper): -# N = 10_000 -# source = tof.Source(facility='ess', neutrons=N) - -# chopper1 = make_chopper( -# topen=[5.0 * ms], -# tclose=[16.0 * ms], -# f=10.0 * Hz, -# phase=0.0 * deg, -# distance=10 * meter, -# name='chopper1', -# ) -# chopper2 = make_chopper( -# topen=[9.0 * ms, 15.0 * ms], -# tclose=[15.0 * ms, 20.0 * ms], -# f=15.0 * Hz, -# phase=0.0 * deg, -# distance=15 * meter, -# name='chopper2', -# ) - -# detector = tof.Detector(distance=20 * meter, name='detector') -# model = tof.Model( -# source=source, choppers=[chopper1, chopper2], detectors=[detector] -# ) -# res = model.run() - -# ch1 = res.choppers['chopper1'].data -# ch2 = res.choppers['chopper2'].data -# det = res.detectors['detector'].data - -# assert sc.allclose(ch1.coords['tof'], ch1.coords['toa'] - ch1.coords['birth_time']) -# assert sc.allclose(ch2.coords['tof'], ch2.coords['toa'] - ch2.coords['birth_time']) -# assert sc.allclose(det.coords['tof'], det.coords['toa'] - det.coords['birth_time']) - - def test_source_not_at_origin(make_chopper): N = 100_000 source1 = tof.Source(facility='ess', neutrons=N, seed=123) From 0e1e99d6531096cd4d11e7dcd0cb8e23b50b4606 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 14:29:04 +0100 Subject: [PATCH 18/22] fix dashboard --- docs/api-reference/index.md | 5 ++++- src/tof/dashboard.py | 37 +++++++++++++++++++------------------ src/tof/source.py | 16 ++++++++++++++++ 3 files changed, 39 insertions(+), 19 deletions(-) diff --git a/docs/api-reference/index.md b/docs/api-reference/index.md index 5458e71..38d33a4 100644 --- a/docs/api-reference/index.md +++ b/docs/api-reference/index.md @@ -12,15 +12,17 @@ Chopper ChopperReading + Component ComponentReading Dashboard Detector DetectorReading + InelasticSample Model ReadingField Result Source - SourceParameters + SourceReading ``` ## Top-level functions @@ -39,5 +41,6 @@ :template: module-template.rst :recursive: + facilities utils ``` diff --git a/src/tof/dashboard.py b/src/tof/dashboard.py index d9fbae2..28e4f21 100644 --- a/src/tof/dashboard.py +++ b/src/tof/dashboard.py @@ -259,24 +259,25 @@ def populate_from_instrument(self, change): for det in self.detectors_container.children: self.remove_detector(None, uid=det._uid) params = INSTRUMENT_LIBRARY[change["new"]] - for ch in params["choppers"]: - self.add_chopper(None) - chop = self.choppers_container.children[-1] - chop.frequency_widget.value = ch.frequency.to(unit='Hz').value - chop.open_widget.value = ", ".join( - str(x) for x in ch.open.to(unit='deg').values - ) - chop.close_widget.value = ", ".join( - str(x) for x in ch.close.to(unit='deg').values - ) - chop.phase_widget.value = ch.phase.to(unit='deg').value - chop.distance_widget.value = ch.distance.to(unit='m').value - chop.name_widget.value = ch.name - for d in params["detectors"]: - self.add_detector(None) - det = self.detectors_container.children[-1] - det.distance_widget.value = d.distance.to(unit='m').value - det.name_widget.value = d.name + for comp in params["components"]: + if comp.kind == "chopper": + self.add_chopper(None) + chop = self.choppers_container.children[-1] + chop.frequency_widget.value = comp.frequency.to(unit='Hz').value + chop.open_widget.value = ", ".join( + str(x) for x in comp.open.to(unit='deg').values + ) + chop.close_widget.value = ", ".join( + str(x) for x in comp.close.to(unit='deg').values + ) + chop.phase_widget.value = comp.phase.to(unit='deg').value + chop.distance_widget.value = comp.distance.to(unit='m').value + chop.name_widget.value = comp.name + elif comp.kind == "detector": + self.add_detector(None) + det = self.detectors_container.children[-1] + det.distance_widget.value = comp.distance.to(unit='m').value + det.name_widget.value = comp.name self.run(None) self.continuous_update.value = cont_update_value diff --git a/src/tof/source.py b/src/tof/source.py index 413a1e5..bae0c39 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -495,6 +495,22 @@ def __repr__(self) -> str: @classmethod def from_json(cls, params: dict) -> Source: + """ + Create a source from a JSON-serializable dictionary. + Currently, only sources from facilities are supported when loading from JSON. + + The dictionary should have the following format: + + .. code-block:: json + + { + "type": "source", + "facility": "ess", + "neutrons": 1000000, + "pulses": 1, + "seed": 42 + } + """ if params.get("facility") is None: raise ValueError( "Currently, only sources from facilities are supported when " From 0d59007364d3301c32596035552f3044a6c2a136 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 15:03:18 +0100 Subject: [PATCH 19/22] add tests for inelastic sample --- src/tof/inelastic.py | 8 +- tests/inelastic_test.py | 230 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 236 insertions(+), 2 deletions(-) create mode 100644 tests/inelastic_test.py diff --git a/src/tof/inelastic.py b/src/tof/inelastic.py index 5cf2102..59be940 100644 --- a/src/tof/inelastic.py +++ b/src/tof/inelastic.py @@ -94,7 +94,9 @@ def __init__( self.energies = delta_e.coords[dim] # TODO: check for bin edges self._noise_scale = ( - 0.5 * (self.energies.max() - self.energies.min()).value / (len(delta_e) - 1) + 0.5 + * (self.energies.max() - self.energies.min()).value + / (max(len(delta_e), 2) - 1) ) self.kind = "inelastic_sample" self._rng = np.random.default_rng(seed) @@ -136,7 +138,9 @@ def as_json(self) -> dict: 'type': 'inelastic_sample', 'distance': var_to_dict(self.distance), 'name': self.name, - 'delta_e': var_to_dict(self.delta_e), + 'energies': var_to_dict(self.energies), + 'probabilities': var_to_dict(self.probabilities), + 'seed': self._rng.bit_generator.state['state']['key'][0], } def as_readonly(self, neutrons: sc.DataArray) -> InelasticSampleReading: diff --git a/tests/inelastic_test.py b/tests/inelastic_test.py new file mode 100644 index 0000000..aa0aea2 --- /dev/null +++ b/tests/inelastic_test.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2023 Scipp contributors (https://github.com/scipp) + +import json +import os +import tempfile + +import numpy as np +import pytest +import scipp as sc + +import tof + +Hz = sc.Unit('Hz') +deg = sc.Unit('deg') +meter = sc.Unit('m') +ms = sc.Unit('ms') + + +def test_inelastic_sample_flat_distribution(): + sample = tof.InelasticSample( + distance=28.0 * meter, + name="sample", + delta_e=sc.DataArray( + data=sc.ones(sizes={'e': 100}), + coords={'e': sc.linspace('e', -0.2, 0.2, 100, unit='meV')}, + ), + ) + + choppers = [ + tof.Chopper( + frequency=70.0 * Hz, + open=sc.array(dims=['cutout'], values=[0.0], unit='deg'), + close=sc.array(dims=['cutout'], values=[1.0], unit='deg'), + phase=0.0 * deg, + distance=20.0 * meter, + name="fastchopper", + ), + ] + + detectors = [ + tof.Detector(distance=26.0 * meter, name='monitor'), + tof.Detector(distance=32.0 * meter, name='detector'), + ] + + source = tof.Source(facility='ess', neutrons=500_000, seed=77) + + model = tof.Model(source=source, components=choppers + detectors + [sample]) + model_no_sample = tof.Model(source=source, components=choppers + detectors) + + res = model.run() + res_no_sample = model_no_sample.run() + + assert sc.identical( + res_no_sample['monitor'].data.coords['wavelength'], + res_no_sample['detector'].data.coords['wavelength'], + ) + assert not sc.identical( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + assert not sc.allclose( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + + +def test_inelastic_sample_doube_peaked_distribution(): + delta_e = sc.DataArray( + data=sc.zeros(sizes={'e': 100}), + coords={'e': sc.linspace('e', -0.2, 0.2, 100, unit='meV')}, + ) + delta_e.values[[0, -1]] = 1.0 + sample = tof.InelasticSample(distance=28.0 * meter, name="sample", delta_e=delta_e) + + choppers = [ + tof.Chopper( + frequency=70.0 * Hz, + open=sc.array(dims=['cutout'], values=[0.0], unit='deg'), + close=sc.array(dims=['cutout'], values=[1.0], unit='deg'), + phase=0.0 * deg, + distance=20.0 * meter, + name="fastchopper", + ), + ] + + detectors = [ + tof.Detector(distance=26.0 * meter, name='monitor'), + tof.Detector(distance=32.0 * meter, name='detector'), + ] + + source = tof.Source(facility='ess', neutrons=500_000, seed=78) + + model = tof.Model(source=source, components=choppers + detectors + [sample]) + model_no_sample = tof.Model(source=source, components=choppers + detectors) + + res = model.run() + res_no_sample = model_no_sample.run() + + assert sc.identical( + res_no_sample['monitor'].data.coords['wavelength'], + res_no_sample['detector'].data.coords['wavelength'], + ) + assert not sc.identical( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + assert not sc.allclose( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + + +def test_inelastic_sample_normal_distribution(): + x = sc.linspace('e', -0.2, 0.2, 100, unit='meV') + sig = sc.scalar(0.03, unit='meV') + y = 1.0 / (np.sqrt(2.0 * np.pi) * sig) * sc.exp(-((x / sig) ** 2) / 2) + y.unit = "" + + sample = tof.InelasticSample( + distance=28.0 * meter, + name="sample", + delta_e=sc.DataArray(data=y, coords={'e': x}), + ) + + choppers = [ + tof.Chopper( + frequency=70.0 * Hz, + open=sc.array(dims=['cutout'], values=[0.0], unit='deg'), + close=sc.array(dims=['cutout'], values=[1.0], unit='deg'), + phase=0.0 * deg, + distance=20.0 * meter, + name="fastchopper", + ), + ] + + detectors = [ + tof.Detector(distance=26.0 * meter, name='monitor'), + tof.Detector(distance=32.0 * meter, name='detector'), + ] + + source = tof.Source(facility='ess', neutrons=500_000, seed=78) + + model = tof.Model(source=source, components=choppers + detectors + [sample]) + model_no_sample = tof.Model(source=source, components=choppers + detectors) + + res = model.run() + res_no_sample = model_no_sample.run() + + assert sc.identical( + res_no_sample['monitor'].data.coords['wavelength'], + res_no_sample['detector'].data.coords['wavelength'], + ) + assert not sc.identical( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + assert not sc.allclose( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + + +def test_inelastic_sample_that_has_zero_delta_e(): + sample = tof.InelasticSample( + distance=28.0 * meter, + name="sample", + delta_e=sc.DataArray( + data=sc.array(dims=['e'], values=[1.0], unit=''), + coords={'e': sc.array(dims=['e'], values=[0.0], unit='meV')}, + ), + ) + + choppers = [ + tof.Chopper( + frequency=70.0 * Hz, + open=sc.array(dims=['cutout'], values=[0.0], unit='deg'), + close=sc.array(dims=['cutout'], values=[1.0], unit='deg'), + phase=0.0 * deg, + distance=20.0 * meter, + name="fastchopper", + ), + ] + + detectors = [ + tof.Detector(distance=26.0 * meter, name='monitor'), + tof.Detector(distance=32.0 * meter, name='detector'), + ] + + source = tof.Source(facility='ess', neutrons=500_000, seed=78) + + model = tof.Model(source=source, components=choppers + detectors + [sample]) + model_no_sample = tof.Model(source=source, components=choppers + detectors) + + res = model.run() + res_no_sample = model_no_sample.run() + + assert sc.identical( + res_no_sample['monitor'].data.coords['wavelength'], + res_no_sample['detector'].data.coords['wavelength'], + ) + assert sc.allclose( + res['monitor'].data.coords['wavelength'], + res['detector'].data.coords['wavelength'], + ) + + +def test_inelastic_sample_as_json(): + # delta_e = sc.DataArray( + # data=sc.zeros(sizes={'e': 10}), + # coords={'e': sc.linspace('e', -0.2, 0.2, 10, unit='meV')}, + # ) + # delta_e.values[[0, -1]] = 1.0 + # sample = tof.InelasticSample(distance=28.0 * meter, name="sample", delta_e=delta_e) + + # json_dict = sample.as_json() + # assert json_dict['type'] == 'inelastic_sample' + # assert json_dict['name'] == 'sample' + # assert json_dict['distance']['value'] == 28.0 + # assert json_dict['distance']['unit'] == 'm' + # assert json_dict['delta_e']['values'][0] == 1.0 + # assert json_dict['delta_e']['values'][-1] == 1.0 + # assert json_dict['delta_e']['unit'] == 'meV' + + # sample_from_json = tof.InelasticSample.from_json( + # name=json_dict['name'], params=json_dict + # ) + # assert sc.identical(sample.distance, sample_from_json.distance) + # assert sc.identical(sample.delta_e, sample_from_json.delta_e) + return From cfde0c40d25440224b8dd5b47d21e2b899e870be Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 16:03:48 +0100 Subject: [PATCH 20/22] finish sample tests --- src/tof/chopper.py | 17 ++++------- src/tof/detector.py | 6 ++-- src/tof/inelastic.py | 20 ++++++------- src/tof/utils.py | 21 +++++++++++++ tests/inelastic_test.py | 66 +++++++++++++++++++++++++++-------------- 5 files changed, 82 insertions(+), 48 deletions(-) diff --git a/src/tof/chopper.py b/src/tof/chopper.py index 2a50484..9f367cc 100644 --- a/src/tof/chopper.py +++ b/src/tof/chopper.py @@ -10,7 +10,7 @@ import scipp as sc from .component import Component, ComponentReading -from .utils import two_pi, var_to_dict +from .utils import two_pi, var_from_dict, var_to_dict if TYPE_CHECKING: try: @@ -29,13 +29,7 @@ class Direction(Enum): def _array_or_none(container: dict, key: str) -> sc.Variable | None: - return ( - sc.array( - dims=["cutout"], values=container[key]["value"], unit=container[key]["unit"] - ) - if key in container - else None - ) + return var_from_dict(container[key], dim="cutout") if key in container else None @dataclass(frozen=True) @@ -294,15 +288,14 @@ def from_json(cls, name: str, params: dict) -> Chopper: f"'{params['direction']}' for component {name}." ) return cls( - frequency=params["frequency"]["value"] - * sc.Unit(params["frequency"]["unit"]), + frequency=var_from_dict(params["frequency"]), direction=_dir, open=_array_or_none(params, "open"), close=_array_or_none(params, "close"), centers=_array_or_none(params, "centers"), widths=_array_or_none(params, "widths"), - phase=params["phase"]["value"] * sc.Unit(params["phase"]["unit"]), - distance=params["distance"]["value"] * sc.Unit(params["distance"]["unit"]), + phase=var_from_dict(params["phase"]), + distance=var_from_dict(params["distance"]), name=name, ) diff --git a/src/tof/detector.py b/src/tof/detector.py index bf9b136..57027a8 100644 --- a/src/tof/detector.py +++ b/src/tof/detector.py @@ -7,7 +7,7 @@ import scipp as sc from .component import Component, ComponentReading -from .utils import var_to_dict +from .utils import var_from_dict, var_to_dict @dataclass(frozen=True) @@ -86,9 +86,7 @@ def from_json(cls, name: str, params: dict) -> Detector: Create a detector from a JSON-serializable dictionary. """ return cls( - distance=sc.scalar( - params["distance"]["value"], unit=params["distance"]["unit"] - ), + distance=var_from_dict(params["distance"]), name=name, ) diff --git a/src/tof/inelastic.py b/src/tof/inelastic.py index 59be940..3e412fd 100644 --- a/src/tof/inelastic.py +++ b/src/tof/inelastic.py @@ -11,6 +11,7 @@ from .component import Component, ComponentReading from .utils import ( energy_to_wavelength, + var_from_dict, var_to_dict, wavelength_to_energy, wavelength_to_speed, @@ -88,8 +89,7 @@ def __init__( self.name = name if delta_e.ndim != 1: raise ValueError("delta_e must be a 1D array.") - self.probabilities = delta_e.values - self.probabilities = self.probabilities / self.probabilities.sum() + self.probabilities = delta_e.data / delta_e.data.sum() dim = delta_e.dim self.energies = delta_e.coords[dim] # TODO: check for bin edges @@ -99,7 +99,8 @@ def __init__( / (max(len(delta_e), 2) - 1) ) self.kind = "inelastic_sample" - self._rng = np.random.default_rng(seed) + self.seed = seed + self._rng = np.random.default_rng(self.seed) def __repr__(self) -> str: return f"InelasticSample(name={self.name}, distance={self.distance:c})" @@ -119,12 +120,11 @@ def from_json(cls, name: str, params: dict) -> InelasticSample: Create an inelastic sample from a JSON-serializable dictionary. """ return cls( - distance=sc.scalar( - params["distance"]["value"], unit=params["distance"]["unit"] - ), + distance=var_from_dict(params["distance"]), name=name, - delta_e=sc.scalar( - params["delta_e"]["value"], unit=params["delta_e"]["unit"] + delta_e=sc.DataArray( + data=var_from_dict(params['probabilities'], dim='e'), + coords={'e': var_from_dict(params['energies'], dim='e')}, ), seed=params.get("seed"), ) @@ -140,7 +140,7 @@ def as_json(self) -> dict: 'name': self.name, 'energies': var_to_dict(self.energies), 'probabilities': var_to_dict(self.probabilities), - 'seed': self._rng.bit_generator.state['state']['key'][0], + 'seed': self.seed, } def as_readonly(self, neutrons: sc.DataArray) -> InelasticSampleReading: @@ -165,7 +165,7 @@ def apply( w_initial = neutrons.coords["wavelength"] n = neutrons.shape - inds = self._rng.choice(len(self.energies), size=n, p=self.probabilities) + inds = self._rng.choice(len(self.energies), size=n, p=self.probabilities.values) de = sc.array( dims=w_initial.dims, values=self.energies.values[inds] diff --git a/src/tof/utils.py b/src/tof/utils.py index 822fa4e..83356d5 100644 --- a/src/tof/utils.py +++ b/src/tof/utils.py @@ -6,6 +6,7 @@ from types import MappingProxyType import matplotlib.pyplot as plt +import numpy as np import scipp as sc import scipp.constants as const @@ -130,6 +131,26 @@ def var_to_dict(var: sc.Variable) -> dict: } +def var_from_dict(data: dict, dim: str | None = None) -> sc.Variable: + """ + Convert a dictionary with 'value' and 'unit' keys to a scipp Variable. + + Parameters + ---------- + data: + The dictionary to convert. + dim: + The dimension of the output variable (non-scalar data only). + """ + values = np.asarray(data["value"]) + unit = data['unit'] + if values.shape: + if dim is None: + raise ValueError("Missing dimension to construct variable from json.") + return sc.array(dims=[dim], values=values, unit=unit) + return sc.scalar(values, unit=unit) + + def extract_component_group( components: dict | MappingProxyType, kind: str ) -> MappingProxyType: diff --git a/tests/inelastic_test.py b/tests/inelastic_test.py index aa0aea2..b0cd662 100644 --- a/tests/inelastic_test.py +++ b/tests/inelastic_test.py @@ -206,25 +206,47 @@ def test_inelastic_sample_that_has_zero_delta_e(): def test_inelastic_sample_as_json(): - # delta_e = sc.DataArray( - # data=sc.zeros(sizes={'e': 10}), - # coords={'e': sc.linspace('e', -0.2, 0.2, 10, unit='meV')}, - # ) - # delta_e.values[[0, -1]] = 1.0 - # sample = tof.InelasticSample(distance=28.0 * meter, name="sample", delta_e=delta_e) - - # json_dict = sample.as_json() - # assert json_dict['type'] == 'inelastic_sample' - # assert json_dict['name'] == 'sample' - # assert json_dict['distance']['value'] == 28.0 - # assert json_dict['distance']['unit'] == 'm' - # assert json_dict['delta_e']['values'][0] == 1.0 - # assert json_dict['delta_e']['values'][-1] == 1.0 - # assert json_dict['delta_e']['unit'] == 'meV' - - # sample_from_json = tof.InelasticSample.from_json( - # name=json_dict['name'], params=json_dict - # ) - # assert sc.identical(sample.distance, sample_from_json.distance) - # assert sc.identical(sample.delta_e, sample_from_json.delta_e) - return + p = np.array([0.4, 0.45, 0.7, 1.0, 0.7, 0.45, 0.4]) + e = np.array([-0.5, -0.25, -0.1, 0.0, 0.1, 0.25, 0.5]) + + delta_e = sc.DataArray( + data=sc.array(dims=['e'], values=p), + coords={'e': sc.array(dims=['e'], values=e, unit='meV')}, + ) + sample = tof.InelasticSample( + distance=28.0 * meter, name="sample1", delta_e=delta_e, seed=66 + ) + + json_dict = sample.as_json() + assert json_dict['type'] == 'inelastic_sample' + assert json_dict['name'] == 'sample1' + assert json_dict['distance']['value'] == 28.0 + assert json_dict['distance']['unit'] == 'm' + assert np.array_equal(json_dict['probabilities']['value'], p / p.sum()) + assert json_dict['probabilities']['unit'] == 'dimensionless' + assert np.array_equal(json_dict['energies']['value'], e) + assert json_dict['energies']['unit'] == 'meV' + assert json_dict['seed'] == 66 + + +def test_inelastic_sample_from_json(): + p = np.array([0.4, 0.7, 1.0, 0.7, 0.4]) + e = np.array([-0.5, -0.25, 0.0, 0.25, 0.5]) + json_dict = { + 'type': 'inelastic_sample', + 'distance': {'value': 28.0, 'unit': 'm'}, + 'name': 'sample1', + 'energies': {'value': e, 'unit': 'meV'}, + 'probabilities': {'value': p, 'unit': ''}, + 'seed': 78, + } + sample = tof.InelasticSample.from_json(name=json_dict['name'], params=json_dict) + + assert sample.distance.value == 28.0 + assert sample.distance.unit == 'm' + assert sample.name == 'sample1' + assert np.array_equal(sample.energies.values, e) + assert sample.energies.unit == 'meV' + assert np.array_equal(sample.probabilities.values, p / p.sum()) + assert sample.probabilities.unit == 'dimensionless' + assert sample.seed == 78 From 13ab32c308c0cf87203672262923b32c00813f55 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 16:05:44 +0100 Subject: [PATCH 21/22] static analysis --- tests/inelastic_test.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/inelastic_test.py b/tests/inelastic_test.py index b0cd662..f39ad17 100644 --- a/tests/inelastic_test.py +++ b/tests/inelastic_test.py @@ -1,12 +1,7 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2023 Scipp contributors (https://github.com/scipp) -import json -import os -import tempfile - import numpy as np -import pytest import scipp as sc import tof @@ -14,7 +9,6 @@ Hz = sc.Unit('Hz') deg = sc.Unit('deg') meter = sc.Unit('m') -ms = sc.Unit('ms') def test_inelastic_sample_flat_distribution(): From 7201e0c31e35ab5af87ff873768a59c4d9d4ad72 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Wed, 25 Feb 2026 16:15:12 +0100 Subject: [PATCH 22/22] small style change --- src/tof/source.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/tof/source.py b/src/tof/source.py index bae0c39..fc0489f 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -516,9 +516,7 @@ def from_json(cls, params: dict) -> Source: "Currently, only sources from facilities are supported when " "loading from JSON." ) - source_args = params.copy() - del source_args["type"] - return cls(**source_args) + return cls(**{k: v for k, v in params.items() if k != "type"}) def as_json(self) -> dict: """