Merge pull request #210 from cta-observatory/multi_file_stream

Return stream and allow pure_protobuf in `MultiFile`.

environment.yml

@@ -6,7 +6,7 @@ dependencies:
   - astropy=5.2
   - python=3.11 # nail the python version, so conda does not try upgrading / dowgrading
   - ctapipe=0.19
-  - protozfits=2.2
+  - protozfits=2.4
   - eventio
   - corsikaio
   - zeromq

setup.cfg

@@ -32,7 +32,7 @@ zip_safe = False
 install_requires=
     astropy~=5.2
     ctapipe >=0.19.0,<0.21.0a0
-    protozfits~=2.2
+    protozfits~=2.4
     numpy>=1.20
 
 [options.package_data]

src/ctapipe_io_lst/__init__.py

@@ -612,7 +612,7 @@ def _generator(self):
         mon = self.initialize_mon_container()
 
         # loop on events
-        for count, zfits_event in enumerate(self.multi_file):
+        for count, (_, zfits_event) in enumerate(self.multi_file):
             # Skip "empty" events that occur at the end of some runs
             if zfits_event.event_id == 0:
                 self.log.warning('Event with event_id=0 found, skipping')

src/ctapipe_io_lst/multifiles.py

@@ -60,6 +60,14 @@ class MultiFiles(Component):
         )
     ).tag(config=True)
 
+    pure_protobuf = Bool(
+        default_value=False,
+        help=(
+            "By default, protozfits converts protobuf message to namedtuples of numpy arrays."
+            "If this option is true, the protobuf Message object will be returned instead."
+        ),
+    ).tag(config=True)
+
     last_subrun = Integer(
         default_value=None,
         allow_none=True,
@@ -138,6 +146,7 @@ def _load_next_subrun(self, stream):
 
         if stream is None:
             path = self.path
+            stream = self.file_info.stream
         else:
             self.current_subrun[stream] += 1
 
@@ -160,7 +169,7 @@ def _load_next_subrun(self, stream):
             self._files.pop(stream).close()
 
         Provenance().add_input_file(str(path), "R0")
-        file_ = File(str(path))
+        file_ = File(str(path), pure_protobuf=self.pure_protobuf)
         self._files[stream] = file_
         self.log.info("Opened file %s", path)
         self._events_tables[stream] = file_.Events
@@ -227,4 +236,4 @@ def __next__(self):
                 except FileNotFoundError:
                     pass
 
-        return event
+        return stream, event

src/ctapipe_io_lst/tests/test_cta_r1.py

@@ -8,8 +8,8 @@
 import astropy.units as u
 
 import protozfits
-from protozfits.CTA_R1_pb2 import CameraConfiguration, Event, TelescopeDataStream
-from protozfits.Debug_R1_pb2 import DebugEvent, DebugCameraConfiguration
+from protozfits.R1v1_pb2 import CameraConfiguration, Event, TelescopeDataStream
+from protozfits.R1v1_debug_pb2 import DebugEvent, DebugCameraConfiguration
 from protozfits.CoreMessages_pb2 import AnyArray
 from traitlets.config import Config
 from ctapipe_io_lst import LSTEventSource

src/ctapipe_io_lst/tests/test_multifile.py

@@ -1,6 +1,5 @@
-import pytest
-from pathlib import Path
 import os
+from pathlib import Path
 
 test_data = Path(os.getenv('LSTCHAIN_TEST_DATA', 'test_data')).absolute()
 test_r0_dir = test_data / 'multifile_test'
@@ -16,9 +15,10 @@ def test_multifile_streams():
         assert multi_files.dvr_applied is False
 
         event_count = 0
-        for event in multi_files:
+        for stream, event in multi_files:
             event_count += 1
             assert event.event_id == event_count
+            assert stream in (1, 2, 3, 4)
 
         assert event_count == 40
 
@@ -32,9 +32,10 @@ def test_multifile_all_subruns():
         assert multi_files.n_open_files == 4
 
         event_count = 0
-        for event in multi_files:
+        for stream, event in multi_files:
             event_count += 1
             assert event.event_id == event_count
+            assert stream in (1, 2, 3, 4)
 
         assert event_count == 200
 
@@ -48,9 +49,10 @@ def test_multifile_last_subrun():
         assert multi_files.n_open_files == 4
 
         event_count = 80
-        for event in multi_files:
+        for stream, event in multi_files:
             event_count += 1
             assert event.event_id == event_count
+            assert stream in (1, 2, 3, 4)
 
         assert event_count == 200
 
@@ -59,9 +61,10 @@ def test_multifile_last_subrun():
         assert multi_files.n_open_files == 4
 
         event_count = 80
-        for event in multi_files:
+        for stream, event in multi_files:
             event_count += 1
             assert event.event_id == event_count
+            assert stream in (1, 2, 3, 4)
 
         assert event_count == 160
 
@@ -71,21 +74,47 @@ def test_multifile_single():
 
     path = test_r0_dir / 'LST-1.3.Run00001.0002.fits.fz'
 
+    # only load multiple streams if stream 1 is passed
     with MultiFiles(path, all_streams=True, all_subruns=True) as multi_files:
         assert multi_files.n_open_files == 1
 
         event_count = 79
-        for event in multi_files:
+        for stream, event in multi_files:
             event_count += 4
             assert event.event_id == event_count
+            assert stream == 3
         assert event_count == 119
 
+    # explicitly turn multiple streams off
     path = test_r0_dir / 'LST-1.1.Run00001.0000.fits.fz'
     with MultiFiles(path, all_streams=False, all_subruns=False) as multi_files:
         assert multi_files.n_open_files == 1
 
         event_count = -3
-        for event in multi_files:
+        for stream, event in multi_files:
             event_count += 4
             assert event.event_id == event_count
+            assert stream == 1
         assert event_count == 37
+
+
+def test_multifile_pure_protobuf():
+    from protozfits import get_class_from_PBFHEAD
+    from ctapipe_io_lst.multifiles import MultiFiles
+
+    path = test_r0_dir / 'LST-1.1.Run00001.0000.fits.fz'
+
+    with MultiFiles(path, pure_protobuf=True) as multi_files:
+        assert multi_files.n_open_files == 4
+        assert multi_files.dvr_applied is False
+
+        cls = get_class_from_PBFHEAD("ProtoR1.CameraEvent")
+
+        event_count = 0
+        for stream, event in multi_files:
+            assert isinstance(event, cls)
+            event_count += 1
+            assert event.event_id == event_count
+            assert stream in (1, 2, 3, 4)
+
+        assert event_count == 40