Add support for passing a union of runs to Scantool

docs/changelog.md

@@ -24,6 +24,8 @@ Added:
 - [Grating2DCalibration][extra.recipes.Grating2DCalibration] to calibrate data from a 2D grating detector (!284).
 - Exposed detector data components from `extra_data` in `extra.components`
   (AGIPD1M, AGIPD500K, DSSC1M, JUNGFRAU, LPD1M) (!177).
+- Added support for passing a union of runs to
+  [Scantool][extra.components.Scantool] (!327).
 
 Changed:
 - [Timepix3.spatial_bins()] is now a static method.

src/extra/components/scantool.py

@@ -1,4 +1,5 @@
 import ast
+import logging
 from warnings import warn
 
 from extra_data import SourceData
@@ -43,14 +44,34 @@ def __init__(self, run, src=None):
             else:
                 raise RuntimeError(f"Found multiple possible scantools, please pass one explicitly with the `src` argument: {', '.join(possible_devices)}")
 
-        values = run.get_run_values(src)
+        self._source_name = src
+        self._source = run[src]
+
+        # If the run is a union then we have to use the CONTROL values rather
+        # than RUN values.
+        run_values = run.get_run_values(src) if run.is_single_run else { }
+        if not run.is_single_run:
+            logging.warning("The passed DataCollection represents multiple runs, "
+                            "but this component will only take the Scantool settings from the first train.")
+
+        def get_value(key, raise_on_missing=True, is_str=False):
+            if key in run_values:
+                return run_values[key]
+            elif key in self._source:
+                arr = self._source[key][0].ndarray().squeeze()
+                return arr.item().decode() if is_str else arr
+            elif not raise_on_missing:
+                return None
+
+            raise KeyError(f"Could not find key '{key}' in either the RUN or CONTROL section")
 
         def get_first_value(keys):
             for key in keys:
-                if key in values:
-                    return values[key]
+                x = get_value(key, raise_on_missing=False)
+                if x is not None:
+                    return x
 
-            raise KeyError(f"Could not find any of these RUN section keys: {', '.join(keys)}")
+            raise KeyError(f"Could not find any of these keys in the RUN or CONTROL section: {', '.join(keys)}")
 
         # These are a list of possible property names for different versions of
         # the scantool. So far we've only seen the names being different, the
@@ -60,10 +81,8 @@ def get_first_value(keys):
         active_motors_keys = ["deviceEnv.activeMotors.value", "activeMotors.value"]
 
         # Get scan metadata and list of motors
-        self._source_name = src
-        self._source = run[src]
         self._active = self.source["isMoving"].ndarray().any()
-        self._scan_type = values["scanEnv.scanType.value"]
+        self._scan_type = get_value("scanEnv.scanType.value", is_str=True)
         self._motors = [x.decode() for x in get_first_value(active_motors_keys) if len(x) > 0]
 
         # The acquisition time vector gives the length of each step, unless the
@@ -73,14 +92,14 @@ def get_first_value(keys):
         # - https://git.xfel.eu/karaboDevices/Karabacon/-/blob/bd22d4a69bf7a401856f49920789ef42fda14ad2/src/karabacon/enums.py#L67
         self._acquisition_time = get_first_value(acquisition_time_keys)
         if _isinstance_no_import(self._acquisition_time, "numpy", "ndarray"):
-            if "Continuous" in values["deviceEnv.acquisitionMode.value"]:
+            if "Continuous" in get_value("deviceEnv.acquisitionMode.value", is_str=True):
                 self._acquisition_time = self._acquisition_time[0]
 
         # The deviceEnv.activeMotors property stores the motor aliases,
         # but we can try to get the actual device names from the
         # actualConfiguration property.
         self._motor_devices = None
-        motors_line = [x for x in values["actualConfiguration.value"].split("---") if "Motors:" in x]
+        motors_line = [x for x in get_value("actualConfiguration.value", is_str=True).split("---") if "Motors:" in x]
         device_names_warning = "Couldn't extract the Karabo device names for the active motors."
         if len(motors_line) == 1:
             try:
@@ -95,11 +114,11 @@ def get_first_value(keys):
         # Get the number of steps and start/stop positions for each motor
         n_motors = len(self.motors)
         self._steps = dict(zip(self.motors,
-                               values["scanEnv.steps.value"][:n_motors]))
+                               get_value("scanEnv.steps.value")[:n_motors]))
         self._start_positions = dict(zip(self.motors,
-                                         values["scanEnv.startPoints.value"][:n_motors]))
+                                         get_value("scanEnv.startPoints.value")[:n_motors]))
         self._stop_positions = dict(zip(self.motors,
-                                        values["scanEnv.stopPoints.value"][:n_motors]))
+                                        get_value("scanEnv.stopPoints.value")[:n_motors]))
 
     @property
     def source_name(self) -> str:

tests/conftest.py

@@ -18,6 +18,7 @@
 from .mockdata.timepix import Timepix3Receiver, Timepix3Centroids
 from .mockdata.timeserver import PulsePatternDecoder, Timeserver
 from .mockdata.xgm import XGM, XGMD, XGMReduced
+from .mockdata.karabacon import Karabacon
 
 
 @pytest.fixture(scope='session')
@@ -156,3 +157,13 @@ def mock_timepix_exceeded_buffer_run(mock_sqs_timepix_directory):
             size_dset[np.argmax(size_dset)] += tpx_root['data/x'].shape[1]
 
         yield RunDirectory(td).deselect('SQS_EXTRA*')
+
+@pytest.fixture(scope="function")
+def mock_scantool_run():
+    sources = [
+        Karabacon("FXE_DAQ_SCAN/MDL/KARABACON"),
+    ]
+
+    with TemporaryDirectory() as td:
+        write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 100)
+        yield RunDirectory(td)

tests/mockdata/karabacon.py

@@ -0,0 +1,42 @@
+import h5py
+import numpy as np
+from extra_data.tests.mockdata.base import DeviceBase
+
+str_dtype = h5py.special_dtype(vlen=str)
+
+class Karabacon(DeviceBase):
+    control_keys = [
+        ("isMoving", "uint8", ()),
+        ("deviceEnv/acquisitionMode", str_dtype, (1,)),
+        ("deviceEnv/acquisitionTimes", "f8", (10_000,)),
+        ("deviceEnv/activeMotors", str_dtype, (1000,)),
+        ("scanEnv/scanType", str_dtype, (1,)),
+        ("scanEnv/steps", "int32", (6,)),
+        ("scanEnv/startPoints", "f8", (6,)),
+        ("scanEnv/stopPoints", "f8", (6,)),
+        ("actualConfiguration", str_dtype, (1,))
+    ]
+
+    def write_control(self, f):
+        super().write_control(f)
+
+        # Values taken from p7948, run 154
+        mock_values = {
+            "deviceEnv/acquisitionMode": np.array([b'Continuous Averaged'], dtype=object),
+            "deviceEnv/acquisitionTimes": np.insert(np.ones(9999), 0, 30),
+            "deviceEnv/activeMotors": np.array([b"ppl_odl"] + [b""] * 999, dtype=object),
+            "scanEnv/scanType": np.array([b'dscan'], dtype=object),
+            "scanEnv/steps": np.array([20,  1,  1,  1,  1,  1], dtype=np.int32),
+            "scanEnv/startPoints": np.array([-15., 1., 1., 1., 1., 1.]),
+            "scanEnv/stopPoints": np.array([25., 1., 1., 1., 1., 1.]),
+            "actualConfiguration": np.array([b"--- Motors: ['FXE_AUXT_LIC/DOOCS/PPODL:default']--- Data Sources: ['FXE_EXP_ONC/METRO/USER_XAS:output0.schema.data.value', 'FXE_EXP_ONC/METRO/USER_XAS:output1.schema.data.value', 'FXE_EXP_ONC/METRO/USER_XAS:output2.schema.data.value', 'FXE_EXP_ONC/METRO/USER_XAS:output6.schema.data.value', 'FXE_EXP_ONC/METRO/USER_XAS:output7.schema.data.value', 'FXE_EXP_ONC/METRO/USER_XAS:output8.schema.data.value']--- Triggers: []"],
+                                                  dtype=object)
+        }
+
+
+        for key, value in mock_values.items():
+            ds_key = f"CONTROL/{self.device_id}/{key}/value"
+            ds = f[ds_key]
+            ntrains = ds.shape[0]
+            for i in range(ntrains):
+                ds[i, :] = value

tests/test_components_scantool.py

@@ -1,11 +1,12 @@
+import logging
 from unittest.mock import MagicMock
 
 from extra.components import Scantool
 
 import pytest
 import numpy as np
 
-def test_scantool():
+def test_scantool(mock_scantool_run, mock_spb_aux_run, caplog):
     mock_name = "CRISPY/KARABACON"
     mock_source = MagicMock()
 
@@ -77,6 +78,18 @@ def test_scantool():
     scantool = Scantool(mock_run)
     assert scantool.acquisition_time == 20
 
+    # Test creating a Scantool from a union of two runs with only control data,
+    # see mockdata/karabacon.py for the settings it was created with.
+    with caplog.at_level(logging.WARNING):
+        scantool = Scantool(mock_scantool_run.union(mock_spb_aux_run))
+        assert len(caplog.records) == 1
+
+    assert scantool.scan_type == "dscan"
+    assert scantool.motors == ["ppl_odl"]
+    assert scantool.acquisition_time == 30
+    assert scantool.start_positions == { "ppl_odl": -15 }
+    assert scantool.stop_positions == { "ppl_odl": 25 }
+
     # Smoke tests
     scantool.info()
     scantool.format()