Add a PPU component

src/extra/components/__init__.py

@@ -1,4 +1,5 @@
 
 from .scantool import Scantool  # noqa
+from .ppu import PPU
 from .pulses import XrayPulses, OpticalLaserPulses, DldPulses  # noqa
 from .scan import Scan

src/extra/components/ppu.py

@@ -41,19 +41,19 @@ def _find_ppu(run: DataCollection, device: str = None):
         raise KeyError("Could not find a PPU device in this data")
     elif len(available_ppus) == 1:
         return run[available_ppus[0]]
-    elif len(available_ppus) > 1:
+    else:  # len(available_ppus) > 1
         if device:
             # And unique substrings of available PPU
             matches = [name for name in available_ppus if device.upper() in name]
             if len(matches) == 1:
                 return run[matches[0]]
             elif len(matches) == 0:
-                KeyError(
-                    f"Couldn't identify an PPU from '{device}'; please pass a valid device name, alias, or unique substring"
+                raise KeyError(
+                    f"Couldn't identify a PPU from '{device}'; please pass a valid device name, alias, or unique substring"
                 )
             else:
-                KeyError(
-                    f"Multiple XGMs found matching '{device}', please be more specific: {matches}"
+                raise KeyError(
+                    f"Multiple PPUs found matching '{device}', please be more specific: {matches}"
                 )
         raise KeyError(f"Multiple PPU devices found in that data: {available_ppus}")
 
@@ -116,7 +116,7 @@ def train_ids(
         for seq, train_id in enumerate(start_train_ids):
             n_trains = self.device["trainTrigger.numberOfTrains"]
             n_trains = n_trains.select_trains(by_id[[train_id]]).ndarray()[0]
-            train_ids.extend(list(range(train_id, train_id + n_trains)))
+            train_ids.extend(np.arange(train_id, train_id + n_trains).tolist())
             sequences.extend([seq] * n_trains)
         # drop train ids missing from the run
         train_ids = sorted(set(train_ids).intersection(self.device.train_ids))

tests/conftest.py

@@ -2,14 +2,26 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
+import h5py
+import numpy as np
 import pytest
-
 from extra_data import RunDirectory
 from extra_data.tests.mockdata import write_file
-from extra_data.tests.mockdata.xgm import XGM
+from extra_data.tests.mockdata.base import DeviceBase
 from extra_data.tests.mockdata.motor import Motor
+from extra_data.tests.mockdata.xgm import XGM
+
+from .mockdata.timeserver import PulsePatternDecoder, Timeserver
 
-from .mockdata.timeserver import Timeserver, PulsePatternDecoder
+
+class PPU(DeviceBase):
+    control_keys = [
+        ('trainTrigger.numberOfTrains', 'i4', ()),
+        ('trainTrigger.sequenceStart', 'i4', ()),
+    ]
+    extra_run_values = [
+        ('classId', None, 'PulsePickerTrainTrigger'),
+    ]
 
 
 @pytest.fixture(scope='session')
@@ -24,3 +36,27 @@ def mock_spb_aux_run():
     with TemporaryDirectory() as td:
         write_file(Path(td) / 'RAW-R0001-DA01-S00000.h5', sources, 100)
         yield RunDirectory(td)
+
+
+@pytest.fixture(scope='session')
+def ppu_run():
+    sources = [
+        PPU('HED_XTD6_PPU/MDL/PPU_TRIGGER'),
+        PPU('HED_DIPOLE_PPU/MDL/PPU_TRIGGER'),
+        Timeserver('HED_RR_SYS/TSYS/TIMESERVER'),
+    ]
+
+    with TemporaryDirectory() as td:
+        fpath = Path(td) / 'RAW-R0001-DA01-S00000.h5'
+        write_file(fpath, sources, 100, firsttrain=10000, format_version='1.3')
+
+        with h5py.File(fpath, 'r+') as f:
+            f['/CONTROL/HED_XTD6_PPU/MDL/PPU_TRIGGER/trainTrigger/numberOfTrains'] = np.array([10] * 100, dtype=np.int64)
+            f['/CONTROL/HED_XTD6_PPU/MDL/PPU_TRIGGER/trainTrigger/sequenceStart'] = np.repeat([9000, 10080], 50)
+            f['/CONTROL/HED_DIPOLE_PPU/MDL/PPU_TRIGGER/trainTrigger/numberOfTrains'] = np.array([1] * 100, dtype=np.int64)
+            f['/CONTROL/HED_DIPOLE_PPU/MDL/PPU_TRIGGER/trainTrigger/sequenceStart'] = np.repeat([9985, 10015, 10045, 10075], 25)
+
+        aliases = {'ppu-hed': 'HED_XTD6_PPU/MDL/PPU_TRIGGER',
+                   'ppu-dipole': 'HED_DIPOLE_PPU/MDL/PPU_TRIGGER'}
+        run = RunDirectory(td)
+        yield run.with_aliases(aliases)

tests/test_components_ppu.py

@@ -0,0 +1,88 @@
+import pandas as pd
+import pytest
+
+from extra_data.reader import DataCollection
+from extra.components import PPU
+from extra.components.ppu import _find_ppu
+
+
+def test_find_ppu(ppu_run):
+    source = _find_ppu(ppu_run, ppu_run['HED_DIPOLE_PPU/MDL/PPU_TRIGGER'])
+    assert source.source == 'HED_DIPOLE_PPU/MDL/PPU_TRIGGER'
+
+    source = _find_ppu(ppu_run, ppu_run['HED_DIPOLE_PPU/MDL/PPU_TRIGGER', 'trainTrigger.sequenceStart'])
+    assert source.source == 'HED_DIPOLE_PPU/MDL/PPU_TRIGGER'
+
+    source = _find_ppu(ppu_run, 'HED_DIPOLE_PPU/MDL/PPU_TRIGGER')
+    assert source.source == 'HED_DIPOLE_PPU/MDL/PPU_TRIGGER'
+
+    source = _find_ppu(ppu_run, 'ppu-hed')
+    assert source.source == 'HED_XTD6_PPU/MDL/PPU_TRIGGER'
+
+    source = _find_ppu(ppu_run, 'XTD6')
+    assert source.source == 'HED_XTD6_PPU/MDL/PPU_TRIGGER'
+
+    source = _find_ppu(ppu_run.select('HED_XTD6_PPU*'))
+    assert source.source == 'HED_XTD6_PPU/MDL/PPU_TRIGGER'
+
+    # fails with multiple PPUs
+    with pytest.raises(KeyError) as excinfo:
+        _find_ppu(ppu_run)
+    assert 'Multiple PPU' in str(excinfo.value)
+
+    # fails with invalid device type
+    with pytest.raises(KeyError) as excinfo:
+        _find_ppu(ppu_run, 1)
+    assert 'not int' in str(excinfo.value)
+
+    # fails with 0 PPUs
+    with pytest.raises(KeyError) as excinfo:
+        _find_ppu(ppu_run.select('*TIMESERVER'))
+    assert 'Could not find a PPU' in str(excinfo.value)
+
+    # too many match
+    with pytest.raises(KeyError) as excinfo:
+        _find_ppu(ppu_run, 'PPU')
+    assert 'Multiple PPUs found matching' in str(excinfo.value)
+
+    # no match
+    with pytest.raises(KeyError) as excinfo:
+        _find_ppu(ppu_run, 'PPU2')
+    assert 'Couldn\'t identify a PPU' in str(excinfo.value)
+
+
+def test_train_ids(ppu_run):
+    # single trigger sequence
+    ppu = PPU(ppu_run, 'ppu-hed')
+    train_ids = ppu.train_ids()
+    assert isinstance(train_ids, list)
+    assert len(train_ids) == 10
+    train_ids = ppu.train_ids(labelled=True)
+    assert isinstance(train_ids, pd.Series)
+    assert train_ids.size == 10  # 10 trains in total
+    assert train_ids.index.unique().size == 1  # single trigger sequence
+
+    # multiple trigger sequences
+    ppu = PPU(ppu_run, 'ppu-dipole')
+    train_ids = ppu.train_ids()
+    assert isinstance(train_ids, list)
+    assert len(train_ids) == 3
+    train_ids = ppu.train_ids(labelled=True)
+    assert isinstance(train_ids, pd.Series)
+    assert train_ids.index.unique().size == 3  # 3 trigger sequence
+    assert train_ids.size == 3  # 1 train per sequence
+
+
+def test_trains(ppu_run):
+    ppu = PPU(ppu_run, 'ppu-dipole')
+    reduced_run = ppu.trains()
+    assert isinstance(reduced_run, DataCollection)
+    assert len(reduced_run.train_ids) == 3
+    assert reduced_run.train_ids == [10015, 10045, 10075]
+
+    # split per sequence
+    reduced_run = ppu.trains(split_sequence=True)
+    assert isinstance(reduced_run, list)
+    assert len(reduced_run) == 3
+    assert len(reduced_run[0].train_ids) == 1
+    assert reduced_run[0].train_ids == [10015]