Add XtdfPulses

Author: philsmt

docs/components/pulse-patterns.md

@@ -6,4 +6,6 @@
 
 ::: extra.components.MachinePulses
 
+::: extra.components.XtdfPulses
+
 ::: extra.components.DldPulses

src/extra/components/__init__.py

@@ -1,7 +1,8 @@
 
 from .scantool import Scantool  # noqa
 from .pulses import XrayPulses, OpticalLaserPulses, MachinePulses, \
-    PumpProbePulses, DldPulses  # noqa
+    PumpProbePulses, XtdfPulses, AgipdPulses, LpdPulses, DsscPulses, \
+    DldPulses  # noqa
 from .scan import Scan  # noqa
 from .xgm import XGM  # noqa
 from .dld import DelayLineDetector  # noqa

src/extra/components/pulses.py

@@ -1508,6 +1508,94 @@ def pulse_mask(self, labelled=True, field=None):
             raise ValueError(f"{field=!r} parameter was not 'fel'/'ppl'/None")
 
 
+class XtdfPulses(PulsePattern):
+    """An interface to pulses from XTDF sources.
+
+    The custom fast MHz detectors used at European XFEL send their data
+    via the XFEL Train Data Format (XTDF), which contains information
+    about the pulse pattern recorded by these detectors.
+
+    This component exposes this information via the pulse pattern
+    interface. As XTDF detectors generally enumerate pulses
+    independently, it also allows to correlate its pulse ID with those
+    obtained from other global sources, e.g. the machine information
+    from [XrayPulses][extra.components.Xraypulses].
+
+    Args:
+        detector (SourceData): Instrument source of an XTDF source.
+        ppt_anchor (int or PulsePattern, optional): Constant offset to
+            add to XTDF pulse IDs, or PulsePattern-based component to
+            correlate every train individually. No correlation is done
+            by default.
+        first_lit_cell (int, optional): Which cell coincides with the
+            first pulse from passed ppt_anchor, ignored unless a
+            PulsePattern component is passed.
+        include_cells (bool, optional): Whether to include cell IDs
+            in the pulse index generated by this component, True by
+            default.
+    """
+
+    def __init__(self, xtdf_source, ppt_anchor=None, first_lit_cell=0,
+                include_cells=True):
+        super().__init__(xtdf_source, xtdf_source['image.pulseId'])
+
+        self._ppt_anchor = ppt_anchor
+        self._first_lit_cell = first_lit_cell
+        self._cell_key = xtdf_source['image.cellId'] if include_cells else None
+
+    def __repr__(self):
+        return "<{} using {}>".format(type(self).__name__, self._source.source)
+
+    def _get_train_ids(self):
+        return self._source.train_ids
+
+    def _get_pulse_ids(self):
+        pulse_ids = self._key.ndarray().astype(np.int32).ravel()
+
+        if isinstance(self._ppt_anchor, int):
+            pulse_ids += self._ppt_anchor
+
+        elif isinstance(self._ppt_anchor, PulsePattern):
+            first, last = self._key.train_index_bounds(False)
+
+            # TODO: Must be aligned!
+            for i, (train_id, anchor_pids) in enumerate(
+                self._ppt_anchor.pulse_ids().groupby('trainId')
+            ):
+                pulse_ids[first[i]:last[i]] += pids.iloc[0] - pulse_ids[
+                    first[i] + self._first_lit_cell]
+
+        elif self._ppt_anchor is not None:
+            raise TypeError(type(self._ppt_anchor))
+
+        index_levels = {
+            'trainId': self._key.train_id_coordinates(),
+            'pulseIndex': np.concatenate([
+                np.arange(count, dtype=np.int32) for count
+                in self._key.data_counts(labelled=False)])}
+
+        if self._cell_key is not None:
+            index_levels['cellId'] = self._cell_key.ndarray().ravel()
+
+        import pandas as pd
+        index = pd.MultiIndex.from_arrays(
+            list(index_levels.values()), names=list(index_levels.keys()))
+
+        return pd.Series(data=pulse_ids, index=index, dtype=np.int32)
+
+
+class AgipdPulses(XtdfPulses):
+    pass
+
+
+class LpdPulses(XtdfPulses):
+    pass
+
+
+class DsscPulses(XtdfPulses):
+    pass
+
+
 class DldPulses(PulsePattern):
     """An interface to pulses from DLD reconstruction.