From ba8d0b47ae2c7dfc1e7707f5ab211bc8d5dd2cb2 Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Mon, 2 Oct 2023 17:34:37 +0200 Subject: [PATCH 1/7] first draft of a ppu component --- src/extra/components/ppu.py | 151 ++++++++++++++++++++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 src/extra/components/ppu.py diff --git a/src/extra/components/ppu.py b/src/extra/components/ppu.py new file mode 100644 index 00000000..a6ad32e1 --- /dev/null +++ b/src/extra/components/ppu.py @@ -0,0 +1,151 @@ +import logging +from typing import List, Union + +import numpy as np +import pandas as pd +from extra_data import by_id +from extra_data.keydata import KeyData +from extra_data.reader import DataCollection +from extra_data.sourcedata import SourceData + +log = logging.getLogger(__name__) + + +def _find_ppu(run: DataCollection, device: str = None): + """Helper function to find a PPU device.""" + + # fast path, we don't validate if the type or name match + if isinstance(device, SourceData): + return device + elif isinstance(device, KeyData): + return run[device.source] + elif isinstance(device, str): + if device in run.control_sources: + return run[device] + elif device in run.alias: + return _find_ppu(run, run.alias[device]) + # else search substring for match + elif device is not None: + raise KeyError(f"ppu must be a SourceData or str, not {type(device).__name__}") + + # Then we list all PPU device in the run + available_ppus = [ + source + for source in run.control_sources + if run[source].device_class in PPU._DEVICE_CLASSES + ] + if len(available_ppus) == 0: + available_ppus = [s for s in run.control_sources if "MDL/PPU" in s] + + if len(available_ppus) == 0: + 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: + 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" + ) + else: + KeyError( + f"Multiple XGMs found matching '{device}', please be more specific: {matches}" + ) + raise KeyError(f"Multiple PPU devices found in that data: {available_ppus}") + + +class PPU: + """Interface to a PPU (Pulse Picker Unit). + + Despite its name, the PPU picks trains, not pulses. + """ + + _DEVICE_CLASSES = ["PulsePickerTrainTrigger", "PulsePickerTrainTriggerCopy"] + + def __init__( + self, data: DataCollection, ppu: Union[KeyData, SourceData, str] = None + ): + """ + + Args: + data (DataCollection): + ppu (Union[KeyData, SourceData, str], optional): + Specify a Pulse Picker Unit device to use, necessary if a run + contains more than one PPU. This can be any of: + - The device name of the source. + - A `SourceData` or [KeyData][extra_data.KeyData] of the + control source (e.g. `HED_XTD6_PPU/MDL/PPU_TRIGGER`) of an + XGM. + - The alias name of either a `SourceData` or + [KeyData][extra_data.KeyData] belonging to a PPU. + - A unique (case-insensitive) substring of a PPU source name. + + Raises: + KeyError: If we can't identify a unique PPU device from the + arguments. + """ + self.data = data + self.device = _find_ppu(data, ppu) + + def train_ids( + self, offset: int = 0, labelled: bool = False + ) -> Union[List[int], pd.Series]: + """All train IDs picked by the PPU. + + Args: + offset (int, optional): + offset to add to the selected trains. Defaults to 0. + labelled (bool, optional): + Returns a Pandas Series if set to True, where this index represents the + trigger sequence a train ID is part of. Defaults to False. + + Returns: + Union[List[int], pd.Series]: Train IDs picked by the PPU. + """ + seq_start = self.device["trainTrigger.sequenceStart"].ndarray() + # The trains picked are the unique values of trainTrigger.sequenceStart + # minus the first (previous trigger before this run). + start_train_ids = np.unique(seq_start)[1:] + offset + + train_ids = [] + sequences = [] + 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))) + sequences.extend([seq] * n_trains) + # drop train ids missing from the run + train_ids = sorted(set(train_ids).intersection(self.device.train_ids)) + + log.info( + f"PPU device {self.device.source} triggered for {len(train_ids)} train(s) across {len(sequences)} sequence(s)." + ) + + if labelled: + train_ids = pd.Series(train_ids, index=sequences) + return train_ids + + def trains( + self, split_sequence: bool = False, offset: int = 0 + ) -> Union[DataCollection, List[DataCollection]]: + """Returns a subset of the data only with Trains selected by the PPU. + + Args: + split_sequence (bool, optional): Split data per PPU trigger sequence. Defaults to False. + offset (int, optional): offset to apply to train IDs to be selected. Defaults to 0. + + Returns: + Union[DataCollection, List[DataCollection]]: + DataCollection(s) containing only trains triggered by the PPU + """ + train_ids = self.train_ids(labelled=True, offset=offset) + if split_sequence: + return [ + self.data.select_trains(by_id[seq.values]) + for _, seq in train_ids.groupby(train_ids.index) + ] + return self.data.select_trains(by_id[train_ids.values]) From 50399e4dd372c6c36e8e5d9dae72d42e7fe6e0bc Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Sat, 7 Oct 2023 12:54:50 +0200 Subject: [PATCH 2/7] add unit tests --- src/extra/components/__init__.py | 1 + src/extra/components/ppu.py | 12 ++--- tests/conftest.py | 42 +++++++++++++-- tests/test_components_ppu.py | 88 ++++++++++++++++++++++++++++++++ 4 files changed, 134 insertions(+), 9 deletions(-) create mode 100644 tests/test_components_ppu.py diff --git a/src/extra/components/__init__.py b/src/extra/components/__init__.py index eea3caf1..845c394b 100644 --- a/src/extra/components/__init__.py +++ b/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 diff --git a/src/extra/components/ppu.py b/src/extra/components/ppu.py index a6ad32e1..89280b28 100644 --- a/src/extra/components/ppu.py +++ b/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)) diff --git a/tests/conftest.py b/tests/conftest.py index dafb4b6c..7df10cde 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/test_components_ppu.py b/tests/test_components_ppu.py new file mode 100644 index 00000000..f0832afd --- /dev/null +++ b/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] From 38359f20a6a05c048c7f3ef64c5ccd3a5051d49c Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Mon, 9 Oct 2023 11:24:53 +0200 Subject: [PATCH 3/7] docs --- docs/components/index.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/components/index.md b/docs/components/index.md index d726bbe2..db2f44ae 100644 --- a/docs/components/index.md +++ b/docs/components/index.md @@ -3,10 +3,12 @@ This module contains classes that abstract various Karabo devices to make access easier. -- [Scans](scans.md) - - [Scantool][extra.components.Scantool] - - [Scan][extra.components.Scan] + - [Pulse patterns](pulse-patterns.md) - [XrayPulses][extra.components.XrayPulses] - [OpticalLaserPulses][extra.components.OpticalLaserPulses] - [DldPulses][extra.components.DldPulses] +- [Pulse Picker Unit][extra.components.PPU] +- [Scans](scans.md) + - [Scantool][extra.components.Scantool] + - [Scan][extra.components.Scan] From f107703d1c3b018688310f4644788a170459db29 Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Mon, 9 Oct 2023 11:52:16 +0200 Subject: [PATCH 4/7] docs --- docs/components/index.md | 2 +- docs/components/pulse-picker-unit.md | 1 + mkdocs.yml | 2 -- src/extra/components/ppu.py | 16 ++++++++++++---- tests/test_components_ppu.py | 2 -- 5 files changed, 14 insertions(+), 9 deletions(-) create mode 100644 docs/components/pulse-picker-unit.md diff --git a/docs/components/index.md b/docs/components/index.md index db2f44ae..bec657b0 100644 --- a/docs/components/index.md +++ b/docs/components/index.md @@ -8,7 +8,7 @@ easier. - [XrayPulses][extra.components.XrayPulses] - [OpticalLaserPulses][extra.components.OpticalLaserPulses] - [DldPulses][extra.components.DldPulses] -- [Pulse Picker Unit][extra.components.PPU] +- [Pulse Picker Unit](pulse-picker-unit.md) - [Scans](scans.md) - [Scantool][extra.components.Scantool] - [Scan][extra.components.Scan] diff --git a/docs/components/pulse-picker-unit.md b/docs/components/pulse-picker-unit.md new file mode 100644 index 00000000..f978dcc6 --- /dev/null +++ b/docs/components/pulse-picker-unit.md @@ -0,0 +1 @@ +::: extra.components.PPU diff --git a/mkdocs.yml b/mkdocs.yml index 666c58e4..72d02285 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -22,8 +22,6 @@ nav: - detector-geometry.md - Components: - components/index.md - - components/scans.md - - components/pulse-patterns.md - karabo-bridge.md - utilities.md - changelog.md diff --git a/src/extra/components/ppu.py b/src/extra/components/ppu.py index 89280b28..395aec5c 100644 --- a/src/extra/components/ppu.py +++ b/src/extra/components/ppu.py @@ -59,9 +59,17 @@ def _find_ppu(run: DataCollection, device: str = None): class PPU: - """Interface to a PPU (Pulse Picker Unit). + """Interface to a Pulse Picker Unit (PPU). - Despite its name, the PPU picks trains, not pulses. + Despite its name, the PPU selects a bunch train from within the 10Hz + structure and block the remainder of the beam. + + Technical description: + A motor-driven absorber rotor is rotated into the beam axis in order to + block the XFEL beam when triggered. The rotor is contained within a UHV + chamber. In terms of temporal structure, the beam pipe is blocked by an + absorbing rotor for up to 9/10ths of a second or vice versa, + synchronized to the facility clock/trigger. """ _DEVICE_CLASSES = ["PulsePickerTrainTrigger", "PulsePickerTrainTriggerCopy"] @@ -78,8 +86,8 @@ def __init__( contains more than one PPU. This can be any of: - The device name of the source. - A `SourceData` or [KeyData][extra_data.KeyData] of the - control source (e.g. `HED_XTD6_PPU/MDL/PPU_TRIGGER`) of an - XGM. + control source (e.g. `HED_XTD6_PPU/MDL/PPU_TRIGGER`) of a + PPU. - The alias name of either a `SourceData` or [KeyData][extra_data.KeyData] belonging to a PPU. - A unique (case-insensitive) substring of a PPU source name. diff --git a/tests/test_components_ppu.py b/tests/test_components_ppu.py index f0832afd..056c636d 100644 --- a/tests/test_components_ppu.py +++ b/tests/test_components_ppu.py @@ -77,12 +77,10 @@ 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] From 4aea70edb4b5848b0b814118318edbf355977721 Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Tue, 10 Oct 2023 16:08:57 +0200 Subject: [PATCH 5/7] do not drop train IDs missing from the run data, this might still be helpful information to have ; the .trains() method will still work and return the correct information --- src/extra/components/ppu.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/extra/components/ppu.py b/src/extra/components/ppu.py index 395aec5c..705f38bd 100644 --- a/src/extra/components/ppu.py +++ b/src/extra/components/ppu.py @@ -126,8 +126,6 @@ def train_ids( n_trains = n_trains.select_trains(by_id[[train_id]]).ndarray()[0] 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)) log.info( f"PPU device {self.device.source} triggered for {len(train_ids)} train(s) across {len(sequences)} sequence(s)." From 44e92b5c0d69952acd9affa8b6cd8e4eb1bfadcd Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Sat, 14 Oct 2023 10:32:14 +0200 Subject: [PATCH 6/7] support for the dipole ppu-like device --- src/extra/components/ppu.py | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/src/extra/components/ppu.py b/src/extra/components/ppu.py index 705f38bd..e85276ee 100644 --- a/src/extra/components/ppu.py +++ b/src/extra/components/ppu.py @@ -1,4 +1,5 @@ import logging +from functools import lru_cache from typing import List, Union import numpy as np @@ -36,6 +37,7 @@ def _find_ppu(run: DataCollection, device: str = None): ] if len(available_ppus) == 0: available_ppus = [s for s in run.control_sources if "MDL/PPU" in s] + available_ppus += [s for s in run.control_sources if "MDL/DIPOLE_PPU" in s] if len(available_ppus) == 0: raise KeyError("Could not find a PPU device in this data") @@ -72,7 +74,11 @@ class PPU: synchronized to the facility clock/trigger. """ - _DEVICE_CLASSES = ["PulsePickerTrainTrigger", "PulsePickerTrainTriggerCopy"] + _DEVICE_CLASSES = [ + "PulsePickerTrainTrigger", # PPU + "PulsePickerTrainTriggerCopy", + "StandardTrigger", # DIPOLE PPU + ] def __init__( self, data: DataCollection, ppu: Union[KeyData, SourceData, str] = None @@ -99,6 +105,20 @@ def __init__( self.data = data self.device = _find_ppu(data, ppu) + @lru_cache() + def number_of_trains(self, train_id: int) -> int: + """Number of trains picked for the sequence starting at train_id. + + Args: + train_id (int): train ID of the sequence start. + """ + + # The Dipole PPU-like device does not allow to pick multiple trains + if "trainTrigger.numberOfTrains" not in self.device.keys(): + return 1 + n_trains = self.device["trainTrigger.numberOfTrains"] + return int(n_trains.select_trains(by_id[[train_id]]).ndarray()[0]) + def train_ids( self, offset: int = 0, labelled: bool = False ) -> Union[List[int], pd.Series]: @@ -122,10 +142,9 @@ def train_ids( train_ids = [] sequences = [] 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(np.arange(train_id, train_id + n_trains).tolist()) - sequences.extend([seq] * n_trains) + span = self.number_of_trains(train_id) + train_ids.extend(np.arange(train_id, train_id + span).tolist()) + sequences.extend([seq] * span) log.info( f"PPU device {self.device.source} triggered for {len(train_ids)} train(s) across {len(sequences)} sequence(s)." From b2af097c34cb368f7714af389d408fc1d49baa0f Mon Sep 17 00:00:00 2001 From: Thomas Michelat Date: Sat, 14 Oct 2023 10:38:00 +0200 Subject: [PATCH 7/7] allow filtering trains on other datas --- src/extra/components/ppu.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/extra/components/ppu.py b/src/extra/components/ppu.py index e85276ee..3e78baf8 100644 --- a/src/extra/components/ppu.py +++ b/src/extra/components/ppu.py @@ -155,11 +155,16 @@ def train_ids( return train_ids def trains( - self, split_sequence: bool = False, offset: int = 0 + self, + data: Union[DataCollection, SourceData, KeyData] = None, + *, + split_sequence: bool = False, + offset: int = 0, ) -> Union[DataCollection, List[DataCollection]]: """Returns a subset of the data only with Trains selected by the PPU. Args: + data: Data to filter. If set to None (defaut) use the data used at initialization. split_sequence (bool, optional): Split data per PPU trigger sequence. Defaults to False. offset (int, optional): offset to apply to train IDs to be selected. Defaults to 0. @@ -167,10 +172,12 @@ def trains( Union[DataCollection, List[DataCollection]]: DataCollection(s) containing only trains triggered by the PPU """ + data = data or self.data + train_ids = self.train_ids(labelled=True, offset=offset) if split_sequence: return [ - self.data.select_trains(by_id[seq.values]) + data.select_trains(by_id[seq.values]) for _, seq in train_ids.groupby(train_ids.index) ] - return self.data.select_trains(by_id[train_ids.values]) + return data.select_trains(by_id[train_ids.values])