Merge pull request #210 from cta-observatory/fix_iter_mc_events

Fix iter mc events reading of telescope data

eventio/__init__.py

@@ -6,7 +6,7 @@
 from .histograms import Histograms
 
 
-__version__ = '1.4.0'
+__version__ = '1.4.1'
 
 __all__ = [
     'EventIOFile',

eventio/base.py

@@ -69,6 +69,7 @@ def __init__(self, path, zcat=True):
         self.path = path
         self.read_process = None
         self.zstd = False
+        self.next = None
 
         if not is_eventio(path):
             raise ValueError('File {} is not an eventio file'.format(path))
@@ -116,6 +117,11 @@ def __iter__(self):
         return self
 
     def __next__(self):
+        if self.next is not None:
+            o = self.next
+            self.next = None
+            return o
+
         self.seek(self._next_header_pos)
         read_sync_marker(self)
         header = read_header(
@@ -130,6 +136,11 @@ def __next__(self):
             filehandle=self._filehandle,
         )
 
+    def peek(self):
+        if self.next is None:
+            self.next = next(self)
+        return self.next
+
     def seek(self, position, whence=0):
         return self._filehandle.seek(position, whence)
 

eventio/iact/objects.py

@@ -161,7 +161,7 @@ class TelescopeData(EventIOObject):
     eventio_type = 1204
 
     def __str__(self):
-        return '{}[{}](reuse_id={})'.format(
+        return '{}[{}](event_id={})'.format(
             self.__class__.__name__,
             self.header.type,
             self.header.id,
@@ -314,11 +314,17 @@ def parse(self):
         data = self.read()
 
         pe.update(PhotoElectrons.parse_1208(
-            data, pe['n_pixels'], pe['non_empty'], self.header.version, flags, pe['n_pe']
+            data, pe['n_pixels'], pe['non_empty'],
+            self.header.version, flags, pe['n_pe']
         ))
 
         return pe
 
+    def __str__(self):
+        return super().__str__() + '(array_id={}, telescope_id={})'.format(
+            self.array_id, self.telescope_id,
+        )
+
 
 class EventEnd(EventIOObject):
     eventio_type = 1209

eventio/simtel/simtelfile.py

@@ -89,7 +89,10 @@ def __init__(self, path, allowed_telescopes=None, skip_calibration=False, zcat=T
         self.camera_monitorings = defaultdict(dict)
         self.laser_calibrations = defaultdict(dict)
         self.current_mc_shower = None
+        self.current_mc_shower_id = None
         self.current_mc_event = None
+        self.current_mc_event_id = None
+        self.current_telescope_data_event_id = None
         self.current_photoelectron_sum = None
         self.current_photoelectrons = {}
         self.current_photons = {}
@@ -139,6 +142,7 @@ def next_low_level(self):
 
         elif isinstance(o, MCShower):
             self.current_mc_shower = o.parse()
+            self.current_mc_shower_id = o.header.id
 
         elif isinstance(o, ArrayEvent):
             self.current_array_event = parse_array_event(
@@ -147,7 +151,8 @@ def next_low_level(self):
             )
 
         elif isinstance(o, iact.TelescopeData):
-            photons, emitter, photoelectrons = parse_telescope_data(o)
+            event_id, photons, emitter, photoelectrons = parse_telescope_data(o)
+            self.current_telescope_data_event_id = event_id
             self.current_photons = photons
             self.current_emitter = emitter
             self.current_photoelectrons = photoelectrons
@@ -206,14 +211,19 @@ def iter_mc_events(self):
 
     def try_build_mc_event(self):
         if self.current_mc_event:
+
             event_data = {
                 'event_id': self.current_mc_event_id,
                 'mc_shower': self.current_mc_shower,
                 'mc_event': self.current_mc_event,
-                'photons': self.current_photons,
-                'emitter': self.current_emitter,
-                'photoelectrons': self.current_photoelectrons,
             }
+            # if next object is TelescopeData, it belongs to this event
+            if isinstance(self.peek(), iact.TelescopeData):
+                self.next_low_level()
+                event_data['photons'] = self.current_photons
+                event_data['emitter'] = self.current_emitter
+                event_data['photoelectrons'] = self.current_photoelectrons
+
             self.current_mc_event = None
             return event_data
         self.next_low_level()
@@ -375,7 +385,7 @@ def parse_telescope_data(telescope_data):
             if e is not None:
                 emitter[o.telescope_id] = e
 
-    return photons, emitter, photo_electrons
+    return telescope_data.header.id, photons, emitter, photo_electrons
 
 
 def parse_telescope_event(telescope_event):

tests/simtel/test_simtelfile.py

@@ -129,13 +129,20 @@ def test_iterate_complete_file_zst():
 def test_iterate_mc_events():
     expected = 200
     with SimTelFile(prod4_path) as f:
+        for counter, event in enumerate(f.iter_mc_events(), start=1):
+            assert 'event_id' in event
+            assert 'mc_shower' in event
+            assert 'mc_event' in event
+
+    assert counter == expected
+
+    with SimTelFile('tests/resources/lst_with_photons.simtel.zst') as f:
         for counter, event in enumerate(f.iter_mc_events(), start=1):
             assert event.keys() == {
                 'event_id',
                 'mc_shower', 'mc_event',
-                'photoelectrons', 'photons', 'emitter',
+                'photons', 'photoelectrons', 'emitter'
             }
-    assert counter == expected
 
 
 def test_allowed_tels():
@@ -200,6 +207,14 @@ def test_new_prod4():
         assert i == 10
 
 
+def test_correct_event_ids_iter_mc_events():
+
+    with SimTelFile('tests/resources/lst_with_photons.simtel.zst') as f:
+        for e in f:
+            assert f.current_mc_event_id == f.current_telescope_data_event_id
+            assert f.current_mc_shower_id == f.current_mc_event_id // 100
+
+
 def test_photons():
     from eventio.iact.objects import Photons
 

tests/test_iterator.py

@@ -16,3 +16,13 @@ def test_iterator():
 
     for o in it:
         assert o.header.content_address > first.header.content_address
+
+
+def test_peek():
+    f = EventIOFile(testfile)
+
+    o = f.peek()
+    assert o is f.peek()  # make sure peek does not advance
+    assert o is next(f)   # make sure peek gives us the next object
+    assert o is not f.peek()  # assure we get the next
+    assert f.peek() is next(f)