Repository: European-XFEL/EXtra-data

Issue:
split_trains method for xtdf detector classes
Something I noticed while working on code for FXE - it would be useful to be able to split the `LPD1M` object into chunks, like we can with `DataCollection` and `KeyData`. In particular, a `frames_per_part` option would be useful when you want to limit how much data a chunk contains without needing to know the number of frames per train. My idea would still round it to a whole number of trains in each chunk, though.

In a similar vein, `KeyData.split_trains` could have an option like `bytes_per_part`. :thinking: 

Discussion:

Patch:
diff --git a/.github/dependabot/constraints.txt b/.github/dependabot/constraints.txt
index 20ca231..d5be32a 100644
--- a/.github/dependabot/constraints.txt
+++ b/.github/dependabot/constraints.txt
@@ -8,7 +8,7 @@ msgpack==1.0.3
 msgpack-numpy==0.4.7.1
 numpy<1.23.0
 pandas<1.5.0
-Pillow==9.0.1
+Pillow==9.0.0
 psutil==5.9.0
 pyparsing==3.0.7
 python-dateutil==2.8.2
diff --git a/docs/agipd_lpd_data.rst b/docs/agipd_lpd_data.rst
index 30485a2..8deca41 100644
--- a/docs/agipd_lpd_data.rst
+++ b/docs/agipd_lpd_data.rst
@@ -35,6 +35,10 @@ DSSC and JUNGFRAU, pulling together the separate modules into a single array.
 
    .. automethod:: trains
 
+   .. automethod:: select_trains
+
+   .. automethod:: split_trains
+
    .. automethod:: write_frames
 
    .. automethod:: write_virtual_cxi
@@ -51,6 +55,12 @@ DSSC and JUNGFRAU, pulling together the separate modules into a single array.
 
    .. automethod:: trains
 
+   .. automethod:: select_trains
+
+   .. automethod:: split_trains
+
+   .. automethod:: write_virtual_cxi
+
 .. autofunction:: identify_multimod_detectors
 
 If you get data for a train from the main :class:`DataCollection` interface,
diff --git a/extra_data/components.py b/extra_data/components.py
index 664d170..c51d917 100644
--- a/extra_data/components.py
+++ b/extra_data/components.py
@@ -1,14 +1,17 @@
 """Interfaces to data from specific instruments
 """
 import logging
+import re
+from copy import copy
+from warnings import warn
+
 import numpy as np
 import pandas as pd
-import re
 import xarray
 
 from .exceptions import SourceNameError
 from .reader import DataCollection, by_id, by_index
-from .read_machinery import roi_shape, DataChunk
+from .read_machinery import DataChunk, roi_shape, split_trains
 from .writer import FileWriter
 from .write_cxi import XtdfCXIWriter, JUNGFRAUCXIWriter
 
@@ -94,6 +97,7 @@ class MultimodDetectorBase:
     _source_re = re.compile(r'(?P<detname>.+)/DET/(\d+)CH')
     # Override in subclass
     _main_data_key = ''  # Key to use for checking data counts match
+    _frames_per_entry = 1  # Override if separate pulse dimension in files
     module_shape = (0, 0)
     n_modules = 0
 
@@ -135,17 +139,13 @@ class MultimodDetectorBase:
         self.modno_to_source = {m: s for (s, m) in source_to_modno.items()}
         assert len(self.modno_to_source) == len(self.source_to_modno)
 
-        train_id_arr = np.asarray(self.data.train_ids)
-        split_indices = np.where(np.diff(train_id_arr) != 1)[0] + 1
-        self.train_id_chunks = np.split(train_id_arr, split_indices)
-        self.frame_counts = frame_counts[train_id_arr]
+        self.frame_counts = frame_counts[self.data.train_ids]
 
         self.train_ids_perframe = np.repeat(
             self.frame_counts.index.values, self.frame_counts.values.astype(np.intp)
         )
-
-        # Cumulative sum gives the end of each train, subtract to get start
-        self.train_id_to_ix = self.frame_counts.cumsum() - self.frame_counts
+        # If we add extra instance attributes, check whether they should be
+        # updated in .select_trains() below.
 
     @classmethod
     def _find_detector_name(cls, data):
@@ -252,18 +252,123 @@ class MultimodDetectorBase:
     def train_ids(self):
         return self.data.train_ids
 
+    @property
+    def train_id_chunks(self):
+        # Used to be used internally. Kept temporarily in case anyone else used it.
+        warn(
+            "detector.train_id_chunks is likely to be removed in the future. "
+            "Please contact da-support@xfel.eu if you're using it",
+            stacklevel=2
+        )
+        train_id_arr = np.asarray(self.data.train_ids)
+        split_indices = np.where(np.diff(train_id_arr) != 1)[0] + 1
+        return np.split(train_id_arr, split_indices)
+
+    @property
+    def train_id_to_ix(self):
+        # Used to be used internally. Kept temporarily in case anyone else used it.
+        warn(
+            "detector.train_id_to_ix is likely to be removed in the future. "
+            "Please contact da-support@xfel.eu if you're using it",
+            stacklevel=2
+        )
+        # Cumulative sum gives the end of each train, subtract to get start
+        return self.frame_counts.cumsum() - self.frame_counts
+
     @property
     def frames_per_train(self):
         counts = set(self.frame_counts.unique()) - {0}
         if len(counts) > 1:
             raise ValueError(f"Varying number of frames per train: {counts}")
-        return counts.pop()
+        return counts.pop() * self._frames_per_entry
 
     def __repr__(self):
         return "<{}: Data interface for detector {!r} with {} modules>".format(
             type(self).__name__, self.detector_name, len(self.source_to_modno),
         )
 
+    def select_trains(self, trains):
+        """Select a subset of trains from this data as a new object.
+
+        Slice trains by position within this data::
+
+            sel = det.select_trains(np.s_[:5])
+
+        Or select trains by train ID, with a slice or a list::
+
+            from extra_data import by_id
+            sel1 = det.select_trains(by_id[142844490 : 142844495])
+            sel2 = det.select_trains(by_id[[142844490, 142844493, 142844494]])
+        """
+        # Using a copy to bypass the source & train checks in __init__
+        res = copy(self)
+        res.data = self.data.select_trains(trains)
+        res.frame_counts = self.frame_counts[res.data.train_ids]
+        res.train_ids_perframe = np.repeat(
+            res.frame_counts.index.values, res.frame_counts.values.astype(np.intp)
+        )
+        return res
+
+    def split_trains(self, parts=None, trains_per_part=None, frames_per_part=None):
+        """Split this data into chunks with a fraction of the trains each.
+
+        At least one of *parts*, *trains_per_part* or *frames_per_part* must be
+        specified. You can pass any combination of these.
+
+        Parameters
+        ----------
+
+        parts: int
+            How many parts to split the data into. If trains_per_part is also
+            specified, this is a minimum, and it may make more parts.
+            It may also make fewer if there are fewer trains in the data.
+        trains_per_part: int
+            A maximum number of trains in each part. Parts will often have
+            fewer trains than this.
+        frames_per_part: int
+            A target number of frames in each part. Each chunk should have up
+            to this many frames, but chunks always contain complete trains,
+            so if this is less than one train, you may get single train chunks
+            with more frames. When ``frames_per_part`` is used, the final
+            chunk may be much smaller than the others.
+        """
+        if {parts, trains_per_part, frames_per_part} == {None}:
+            raise ValueError(
+                "One of parts, trains_per_part, frames_per_part must be specified"
+            )
+        if frames_per_part is None:
+            for s in split_trains(len(self.train_ids), parts, trains_per_part):
+                yield self.select_trains(s)
+        else:
+            # frames_per_part was specified. We don't assume that the number
+            # of frames per train is constant, so we'll iterate over trains
+            # and cut off each chunk when we reach the relevant number.
+            if not self.train_ids:
+                return  # No data to split
+
+            if trains_per_part is None:
+                trains_per_part = np.inf
+            if parts:
+                trains_per_part = min(trains_per_part, len(self.train_ids) // parts)
+
+            chunk_start = 0
+            ntrains = 1
+            nentries = self.frame_counts.iloc[0]
+
+            for frame_ct in self.frame_counts.iloc[1:]:
+                ntrains += 1
+                nentries += frame_ct
+                if (ntrains > trains_per_part) or (nentries * self._frames_per_entry > frames_per_part):
+                    # We've got a full chunk
+                    chunk_end = chunk_start + ntrains - 1
+                    yield self.select_trains(np.s_[chunk_start:chunk_end])
+                    chunk_start = chunk_end
+                    ntrains = 1
+                    nentries = frame_ct
+
+            # There will always be at least the last train left to yield
+            yield self.select_trains(np.s_[chunk_start:])
+
     @staticmethod
     def _concat(arrays, index, fill_value, astype):
         dtype = arrays[0].dtype if astype is None else np.dtype(astype)
@@ -1336,6 +1441,10 @@ class JUNGFRAU(MultimodDetectorBase):
                 data, self.detector_name
             ))
 
+        # In burst mode, JUNGFRAU can have 16 frames per train
+        src = next(iter(self.source_to_modno))
+        self._frames_per_entry = self.data[src, self._main_data_key].entry_shape[0]
+
     @staticmethod
     def _label_dims(arr):
         # Label dimensions to match the AGIPD/DSSC/LPD data access
diff --git a/extra_data/reader.py b/extra_data/reader.py
index 8035bff..29ccfb8 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -850,18 +850,15 @@ class DataCollection:
     def select_trains(self, train_range):
         """Select a subset of trains from this data.
 
-        Choose a slice of trains by train ID::
+        Slice trains by position within this data::
 
-            from extra_data import by_id
-            sel = run.select_trains(by_id[142844490:142844495])
-
-        Or select a list of trains::
-
-            sel = run.select_trains(by_id[[142844490, 142844493, 142844494]])
+            sel = run.select_trains(np.s_[:5])
 
-        Or select trains by index within this collection::
+        Or select trains by train ID, with a slice or a list::
 
-            sel = run.select_trains(np.s_[:5])
+            from extra_data import by_id
+            sel1 = run.select_trains(by_id[142844490 : 142844495])
+            sel2 = run.select_trains(by_id[[142844490, 142844493, 142844494]])
 
         Returns a new :class:`DataCollection` object for the selected trains.
 
diff --git a/extra_data/write_cxi.py b/extra_data/write_cxi.py
index fb8caa1..c73c14f 100644
--- a/extra_data/write_cxi.py
+++ b/extra_data/write_cxi.py
@@ -46,9 +46,6 @@ class VirtualCXIWriterBase:
             frame_counts.index.values, frame_counts.values.astype(np.intp)
         )
 
-        # Cumulative sum gives the end of each train, subtract to get start
-        self.train_id_to_ix = frame_counts.cumsum() - frame_counts
-
         # For AGIPD, DSSC & LPD detectors modules are numbered from 0.
         # Overridden for JUNGFRAU to number from 1.
         self.modulenos = list(range(self.nmodules))


Test patch:
diff --git a/extra_data/tests/test_components.py b/extra_data/tests/test_components.py
index 34f413b..b737a11 100644
--- a/extra_data/tests/test_components.py
+++ b/extra_data/tests/test_components.py
@@ -265,6 +265,60 @@ def test_get_dask_array_jungfrau(mock_jungfrau_run):
     np.testing.assert_array_equal(arr.coords['train'], np.arange(10000, 10100))
 
 
+def test_select_trains(mock_fxe_raw_run):
+    run = RunDirectory(mock_fxe_raw_run)
+    det = LPD1M(run.select_trains(np.s_[:20]))
+    assert len(det.train_ids) == 20
+    det = det.select_trains(np.s_[:2])
+    assert len(det.train_ids) == 2
+
+    arr = det.get_array('image.data')
+    assert arr.shape == (16, 2, 128, 256, 256)
+
+
+def test_split_trains(mock_fxe_raw_run):
+    run = RunDirectory(mock_fxe_raw_run)
+    det = LPD1M(run.select_trains(np.s_[:20]))
+    assert len(det.train_ids) == 20
+
+    parts = list(det.split_trains(parts=4))
+    assert len(parts) == 4
+    assert [len(p.train_ids) for p in parts] == [5, 5, 5, 5]
+    assert sum([p.train_ids for p in parts], []) == det.train_ids
+
+    arr = parts[-1].get_array('image.data', pulses=np.s_[:1])
+    assert arr.shape == (16, 5, 1, 256, 256)
+
+    # Split by a number of frames
+    parts = list(det.split_trains(frames_per_part=256))
+    assert [len(p.train_ids) for p in parts] == [2] * 10
+
+    # frames_per_part less than one train (128 frames in this example)
+    parts = list(det.split_trains(frames_per_part=100))
+    assert [len(p.train_ids) for p in parts] == [1] * 20
+
+    # trains_per_part cuts off before frames_per_part
+    parts = list(det.split_trains(trains_per_part=3, frames_per_part=1024))
+    assert [len(p.train_ids) for p in parts] == ([3] * 6) + [2]
+
+    # parts cuts off before frames_per_part
+    parts = list(det.split_trains(parts=6, frames_per_part=1024))
+    assert [len(p.train_ids) for p in parts] == ([3] * 6) + [2]
+
+    # frames_per_part > all frames in selection
+    parts = list(det.split_trains(frames_per_part=3000))
+    assert [len(p.train_ids) for p in parts] == [20]
+
+
+def test_split_trains_jungfrau(mock_jungfrau_run):
+    run = RunDirectory(mock_jungfrau_run)
+    jf = JUNGFRAU(run.select_trains(np.s_[:20]))
+    assert jf.frames_per_train == 16
+
+    parts = list(jf.split_trains(frames_per_part=64))
+    assert [len(p.train_ids) for p in parts] == [4] * 5
+
+
 def test_iterate(mock_fxe_raw_run):
     run = RunDirectory(mock_fxe_raw_run)
     det = LPD1M(run.select_trains(by_index[:2]))