diff --git a/extra_data/components.py b/extra_data/components.py index 4d3d56e9..81c59560 100644 --- a/extra_data/components.py +++ b/extra_data/components.py @@ -425,6 +425,17 @@ def __repr__(self): return (f"<{det}: Data interface for detector {self.detector_name!r} " f"- {rp} data with {len(self.source_to_modno)} modules>") + def _with_train_sel_data(self, data_sel): + # Used in select_trains & split_trains below + # Using a copy to bypass the source & train checks in __init__ + res = copy(self) + res.data = data_sel + 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 select_trains(self, trains): """Select a subset of trains from this data as a new object. @@ -438,14 +449,7 @@ def select_trains(self, trains): 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 + return self._with_train_sel_data(self.data.select_trains(trains)) 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. @@ -474,9 +478,21 @@ def split_trains(self, parts=None, trains_per_part=None, frames_per_part=None): raise ValueError( "One of parts, trains_per_part, frames_per_part must be specified" ) + if frames_per_part is not None: + # If possible, convert frames_per_part to trains_per_part, which + # uses the more efficient code path. + try: + fpt = self.frames_per_train + except ValueError: + pass + else: + ntpp = max(1, frames_per_part // fpt) + trains_per_part = min(ntpp, trains_per_part or ntpp) + frames_per_part = None + if frames_per_part is None: - for s in split_trains(len(self.train_ids), parts, trains_per_part): - yield self.select_trains(s) + for part in self.data.split_trains(parts=parts, trains_per_part=trains_per_part): + yield self._with_train_sel_data(part) 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 diff --git a/extra_data/tests/test_components.py b/extra_data/tests/test_components.py index caa4b8d0..7f362168 100644 --- a/extra_data/tests/test_components.py +++ b/extra_data/tests/test_components.py @@ -479,11 +479,12 @@ def test_split_trains(mock_fxe_raw_run): # 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] + assert {len(p.train_ids) for p in parts} == {2, 3} # 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] + assert len(parts) == 6 + assert {len(p.train_ids) for p in parts} == {3, 4} # frames_per_part > all frames in selection parts = list(det.split_trains(frames_per_part=3000))