Fix subarray event source to check for telescope source exhaustion

Author: maxnoe

src/ctapipe_io_zfits/conftest.py

@@ -101,6 +101,16 @@ def get_module_and_pixel_id_map(n_modules, n_pixels_module, missing_modules=None
         },
         id="no_tel_ids_with_data",
     ),
+    pytest.param(
+        {
+            "all_chunks": True,
+            "obs_start": Time("2025-02-04T20:45:31"),
+            "sb_creator_id": 2,
+            "sb_id": 126,
+            "obs_id": 127,
+        },
+        id="all_chunks",
+    ),
 ]
 
 

src/ctapipe_io_zfits/dl0.py

@@ -308,6 +308,10 @@ def _generator(self):
         for count, subarray_trigger in enumerate(
             self._subarray_trigger_file.SubarrayEvents
         ):
+            # stop processing when all telescope event sources are exhausted
+            if all(f.exhausted for f in self._telescope_files.values()):
+                return
+
             array_event = ArrayEventContainer(
                 count=count,
                 index=EventIndexContainer(

src/ctapipe_io_zfits/multifile.py

@@ -223,6 +223,7 @@ def __init__(self, path, *args, **kwargs):
         self._events_headers = {}
         self.camera_config = None
         self.data_stream = None
+        self._exhausted_sources = set()
 
         for data_source in self.data_sources:
             self._load_next_chunk(data_source)
@@ -236,6 +237,11 @@ def _valid_filename_convention(self, proposal):
             raise TraitError(msg)
         return value
 
+    @property
+    def exhausted(self):
+        """Whether all data sources / chunks available have been consumed."""
+        return len(self._exhausted_sources) == len(self.data_sources)
+
     @property
     def n_open_files(self):
         """Number of currently open files."""
@@ -326,6 +332,8 @@ def __next__(self):  # noqa: D105
                 try:
                     self._load_next_chunk(data_source)
                 except FileNotFoundError:
-                    pass
+                    self._exhausted_sources.add(data_source)
+            else:
+                self._exhausted_sources.add(data_source)
 
         return event

src/ctapipe_io_zfits/tests/test_dl0.py

@@ -5,6 +5,7 @@
 from ctapipe.instrument import SubarrayDescription
 from ctapipe.io import EventSource, TableLoader
 from ctapipe.tools.process import ProcessorTool
+from traitlets.config import Config
 
 
 def test_is_compatible(dummy_dl0):
@@ -27,8 +28,10 @@ def test_subarray(dummy_dl0):
 
 def test_subarray_events(dummy_dl0):
     time = dummy_dl0["obs_start"]
+    all_chunks = dummy_dl0.get("all_chunks", False)
+    config = Config({"MultiFiles": {"all_chunks": all_chunks}})
 
-    with EventSource(dummy_dl0["trigger_path"]) as source:
+    with EventSource(dummy_dl0["trigger_path"], config=config) as source:
         n_read = 0
         for i, array_event in enumerate(source):
             assert array_event.count == i
@@ -46,7 +49,7 @@ def test_subarray_events(dummy_dl0):
             n_read += 1
             time = time + 0.001 * u.s
 
-        assert n_read == 100
+        assert n_read == (100 if all_chunks else 40)
 
 
 def test_process(dummy_dl0, tmp_path):
@@ -130,7 +133,8 @@ def test_telescope_event_source_missing_ids(dummy_tel_file_no_ids):
     assert ProtozfitsDL0TelescopeEventSource.is_compatible(first_ff_file)
     assert ProtozfitsDL0TelescopeEventSource.is_compatible(first_ped_file)
 
-    with EventSource(first_ff_file) as source:
+    config = Config({"MultiFiles": {"all_chunks": True}})
+    with EventSource(first_ff_file, config=config) as source:
         assert isinstance(source, ProtozfitsDL0TelescopeEventSource)
 
         n_read = 0
@@ -142,7 +146,7 @@ def test_telescope_event_source_missing_ids(dummy_tel_file_no_ids):
 
         assert n_read == 50
 
-    with EventSource(first_ped_file) as source:
+    with EventSource(first_ped_file, config=config) as source:
         assert isinstance(source, ProtozfitsDL0TelescopeEventSource)
 
         n_read = 0

src/ctapipe_io_zfits/tests/test_multifiles.py

@@ -23,6 +23,8 @@ def test_multifiles(all_chunks, dummy_tel_file, dl0_base):
         else:
             assert expected_event_id == 40
 
+        assert mf.exhausted
+
     recorded_inputs = Provenance().current_activity.provenance["input"]
     # five chunks per stream
     if all_chunks: