From 4ae4bbd262ad363f5e0fed7d4e4c890389cf401f Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Fri, 30 Jan 2026 22:30:31 +0100 Subject: [PATCH 1/3] add data property on results that returns a datagroup --- src/tof/result.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/tof/result.py b/src/tof/result.py index 266812a..6589830 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -344,3 +344,12 @@ def to_nxevent_data(self, key: str | None = None) -> sc.DataArray: ) out.coords["Ltotal"] = out.coords.pop("distance") return out + + @property + def data(self) -> sc.DataGroup: + return sc.DataGroup( + { + key: value.data + for key, value in chain(self.choppers.items(), self.detectors.items()) + } + ) From b8a605f39185c261bb49d5ab37ad4fddd7f912a7 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Mon, 2 Feb 2026 20:52:56 +0100 Subject: [PATCH 2/3] order components by distance in .data output and add a couple of fixes --- src/tof/model.py | 37 +++++++++---------------------------- src/tof/result.py | 16 +++++++++++----- src/tof/source.py | 9 ++++++++- 3 files changed, 28 insertions(+), 34 deletions(-) diff --git a/src/tof/model.py b/src/tof/model.py index 31732ce..9dfa6d1 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -16,30 +16,6 @@ ComponentType = Chopper | Detector -def _input_to_dict( - obj: None | list[ComponentType] | tuple[ComponentType, ...] | ComponentType, - kind: type, -): - if isinstance(obj, list | tuple): - out = {} - for item in obj: - new = _input_to_dict(item, kind=kind) - for key in new.keys(): - if key in out: - raise ValueError(f"More than one component named '{key}' found.") - out.update(new) - return out - elif isinstance(obj, kind): - return {obj.name: obj} - elif obj is None: - return {} - else: - raise TypeError( - "Invalid input type. Must be a Chopper or a Detector, " - "or a list/tuple of Choppers or Detectors." - ) - - def _array_or_none(container: dict, key: str) -> sc.Variable | None: return ( sc.array( @@ -134,9 +110,11 @@ def __init__( choppers: Chopper | list[Chopper] | tuple[Chopper, ...] | None = None, detectors: Detector | list[Detector] | tuple[Detector, ...] | None = None, ): - self.choppers = _input_to_dict(choppers, kind=Chopper) - self.detectors = _input_to_dict(detectors, kind=Detector) + self.choppers = {} + self.detectors = {} self.source = source + for c in chain(choppers or (), detectors or ()): + self.add(c) @classmethod def from_json(cls, filename: str) -> Model: @@ -212,17 +190,20 @@ def to_json(self, filename: str): with open(filename, 'w') as f: json.dump(self.as_json(), f, indent=2) - def add(self, component): + def add(self, component: Chopper | Detector): """ Add a component to the instrument. Component names must be unique across choppers and detectors. + The name "source" is reserved for the source, and can thus not be used for other + components. Parameters ---------- component: A chopper or detector. """ - if component.name in chain(self.choppers, self.detectors): + # Note that the name "source" is reserved for the source. + if component.name in chain(self.choppers, self.detectors, ("source",)): raise KeyError( f"Component with name {component.name} already exists. " "If you wish to replace/update an existing component, use " diff --git a/src/tof/result.py b/src/tof/result.py index 6589830..0e3ad2b 100644 --- a/src/tof/result.py +++ b/src/tof/result.py @@ -347,9 +347,15 @@ def to_nxevent_data(self, key: str | None = None) -> sc.DataArray: @property def data(self) -> sc.DataGroup: - return sc.DataGroup( - { - key: value.data - for key, value in chain(self.choppers.items(), self.detectors.items()) - } + """ + Get the data for the source, choppers, and detectors, as a DataGroup. + The components are sorted by distance. + """ + out = {"source": self.source.data} + components = sorted( + chain(self.choppers.values(), self.detectors.values()), + key=lambda c: c.distance.value, ) + for comp in components: + out[comp.name] = comp.data + return sc.DataGroup(out) diff --git a/src/tof/source.py b/src/tof/source.py index 46f33ea..504d4c3 100644 --- a/src/tof/source.py +++ b/src/tof/source.py @@ -290,7 +290,14 @@ def data(self) -> sc.DataArray: """ The data array containing the neutrons in the pulse. """ - return self._data + return self._data.assign_coords( + { + "distance": self._distance, + "eto": self._data.coords["birth_time"] + % (1.0 / self._frequency).to(unit=TIME_UNIT, copy=False), + "toa": self._data.coords["birth_time"], + } + ) @classmethod def from_neutrons( From 3badc1d69dba81ed2b6c82bb62a72ae7c30af8f0 Mon Sep 17 00:00:00 2001 From: Neil Vaytet Date: Tue, 3 Feb 2026 10:33:40 +0100 Subject: [PATCH 3/3] fix tests --- src/tof/model.py | 29 +++++++++++++++++------------ tests/model_test.py | 36 ++++++++++++++++++++---------------- 2 files changed, 37 insertions(+), 28 deletions(-) diff --git a/src/tof/model.py b/src/tof/model.py index 9dfa6d1..36c241f 100644 --- a/src/tof/model.py +++ b/src/tof/model.py @@ -107,14 +107,20 @@ class Model: def __init__( self, source: Source | None = None, - choppers: Chopper | list[Chopper] | tuple[Chopper, ...] | None = None, - detectors: Detector | list[Detector] | tuple[Detector, ...] | None = None, + choppers: list[Chopper] | tuple[Chopper, ...] | None = None, + detectors: list[Detector] | tuple[Detector, ...] | None = None, ): self.choppers = {} self.detectors = {} self.source = source - for c in chain(choppers or (), detectors or ()): - self.add(c) + 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) @classmethod def from_json(cls, filename: str) -> Model: @@ -202,6 +208,11 @@ 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." + ) # Note that the name "source" is reserved for the source. if component.name in chain(self.choppers, self.detectors, ("source",)): raise KeyError( @@ -210,14 +221,8 @@ def add(self, component: Chopper | Detector): "``model.choppers['name'] = new_chopper`` or " "``model.detectors['name'] = new_detector``." ) - if isinstance(component, Chopper): - self.choppers[component.name] = component - elif isinstance(component, Detector): - self.detectors[component.name] = component - else: - raise TypeError( - f"Cannot add component of type {type(component)} to the model." - ) + container = self.choppers if isinstance(component, Chopper) else self.detectors + container[component.name] = component def remove(self, name: str): """ diff --git a/tests/model_test.py b/tests/model_test.py index 86659eb..11e9005 100644 --- a/tests/model_test.py +++ b/tests/model_test.py @@ -382,11 +382,11 @@ def test_create_model_with_duplicate_component_names_raises( chopper = dummy_chopper detector = dummy_detector with pytest.raises( - ValueError, match="More than one component named 'dummy_chopper' found" + KeyError, match="Component with name dummy_chopper already exists" ): tof.Model(source=dummy_source, choppers=[chopper, chopper]) with pytest.raises( - ValueError, match="More than one component named 'dummy_detector' found" + KeyError, match="Component with name dummy_detector already exists" ): tof.Model(source=dummy_source, detectors=[detector, detector]) @@ -422,28 +422,32 @@ def test_getitem(dummy_chopper, dummy_detector, dummy_source): model['foo'] -def test_input_can_be_single_component(dummy_chopper, dummy_detector, dummy_source): - chopper = dummy_chopper - detector = dummy_detector - model = tof.Model(source=dummy_source, choppers=chopper, detectors=detector) - assert 'dummy_chopper' in model.choppers - assert 'dummy_detector' in model.detectors - - def test_bad_input_type_raises(dummy_chopper, dummy_detector, dummy_source): chopper = dummy_chopper detector = dummy_detector - with pytest.raises(TypeError, match='Invalid input type'): + with pytest.raises( + TypeError, match='Beamline components: expected Chopper instance' + ): _ = tof.Model(source=dummy_source, choppers='bad chopper') - with pytest.raises(TypeError, match='Invalid input type'): + with pytest.raises( + TypeError, match='Beamline components: expected Detector instance' + ): _ = tof.Model(source=dummy_source, choppers=[chopper], detectors='abc') - with pytest.raises(TypeError, match='Invalid input type'): + with pytest.raises( + TypeError, match='Beamline components: expected Chopper instance' + ): _ = tof.Model(source=dummy_source, choppers=[chopper, 'bad chopper']) - with pytest.raises(TypeError, match='Invalid input type'): + with pytest.raises( + TypeError, match='Beamline components: expected Detector instance' + ): _ = tof.Model(source=dummy_source, detectors=(1234, detector)) - with pytest.raises(TypeError, match='Invalid input type'): + with pytest.raises( + TypeError, match='Beamline components: expected Chopper instance' + ): _ = tof.Model(source=dummy_source, choppers=[detector]) - with pytest.raises(TypeError, match='Invalid input type'): + with pytest.raises( + TypeError, match='Beamline components: expected Detector instance' + ): _ = tof.Model(source=dummy_source, detectors=[chopper])