MultiFile: Return stream from which event was loaded

Author: maxnoe

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

@@ -146,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
 
@@ -235,4 +236,4 @@ def __next__(self):
                 except FileNotFoundError:
                     pass
 
-        return event
+        return stream, event

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,25 @@ 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