From 34bf819c23129d469d97ac19a13bea1acbfbf8c9 Mon Sep 17 00:00:00 2001 From: Philipp Schmidt Date: Fri, 25 Apr 2025 14:56:06 +0200 Subject: [PATCH] Add fel_shutter/ppl_shutter arguments to PumpProbePulses --- src/extra/components/pulses.py | 39 +++++++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/src/extra/components/pulses.py b/src/extra/components/pulses.py index 0bacb958..717f0424 100644 --- a/src/extra/components/pulses.py +++ b/src/extra/components/pulses.py @@ -1258,6 +1258,10 @@ class PumpProbePulses(XrayPulses, OpticalLaserPulses): extrapolated from past trains if missing, true by default. An exception is raised if disabled and the PPL anchoring method is missing the minimum number of required FEL pulses. + fel_shutter (KeyData, optional): Scalar key evaluated as a + shutter state blocking FEL pulses in a given train. + ppl_shutter (KeyData, optional): Scalar key evaluated as a + shutter state blocking PPL pulses in a given train. """ # This class inherits from two classes which both have the same @@ -1277,7 +1281,8 @@ class PumpProbePulses(XrayPulses, OpticalLaserPulses): def __init__(self, data, source=None, instrument=None, *, bunch_table_position=None, bunch_table_offset=None, - pulse_offset=None, extrapolate=True): + pulse_offset=None, extrapolate=True, + fel_shutter=None, ppl_shutter=None): self._bunch_table_position = None self._bunch_table_offset = None self._pulse_offset = None # Allowed to be float! @@ -1325,6 +1330,20 @@ def __init__(self, data, source=None, instrument=None, *, self._sase = sase + if isinstance(fel_shutter, KeyData): + self._fel_shutter = fel_shutter + elif fel_shutter is not None: + raise TypeError('expected KeyData object for fel_shutter') + else: + self._fel_shutter = None + + if isinstance(ppl_shutter, KeyData): + self._ppl_shutter = ppl_shutter + elif ppl_shutter is not None: + raise TypeError('expected KeyData object for ppl_shutter') + else: + self._ppl_shutter = None + def __repr__(self): if self._bunch_table_position is not None: offset_str = f'@{self._bunch_table_position}b' @@ -1385,6 +1404,18 @@ def _get_pulse_ids(self): train_ids = self._key.train_id_coordinates() prev_fel_pids = None + if self._fel_shutter is not None: + is_closed = self._fel_shutter.xarray().astype(bool) + trains_without_fel = set(is_closed.trainId[is_closed].data) + else: + trains_without_fel = set() + + if self._ppl_shutter is not None: + is_closed = self._ppl_shutter.xarray().astype(bool) + trains_without_ppl = set(is_closed.trainId[is_closed].data) + else: + trains_without_ppl = set() + for train_id, (fel_pids, ppl_pids) in zip(train_ids, iter_pulse_ids): try: ppl_pids += self._get_ppl_offset(fel_pids) @@ -1413,6 +1444,12 @@ def _get_pulse_ids(self): else: prev_fel_pids = fel_pids + if train_id in trains_without_fel: + fel_pids = np.zeros(0) + + if train_id in trains_without_ppl: + ppl_pids = np.zeros(0) + pids = np.union1d(fel_pids, ppl_pids) pids_by_train.append(pids)