Merge pull request #119 from scipp/result-data

Add data property on results that returns a DataGroup

Author: nvaytet

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(
@@ -131,12 +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 = _input_to_dict(choppers, kind=Chopper)
-        self.detectors = _input_to_dict(detectors, kind=Detector)
+        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)
 
     @classmethod
     def from_json(cls, filename: str) -> Model:
@@ -212,31 +196,33 @@ 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):
+        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(
                 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``."
             )
-        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):
         """

src/tof/result.py

@@ -344,3 +344,18 @@ 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:
+        """
+        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)

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(

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])