diff --git a/eventio/__init__.py b/eventio/__init__.py index 4388daf..75e5f9b 100644 --- a/eventio/__init__.py +++ b/eventio/__init__.py @@ -6,7 +6,7 @@ from .histograms import Histograms -__version__ = '1.4.0' +__version__ = '1.4.1' __all__ = [ 'EventIOFile', diff --git a/eventio/base.py b/eventio/base.py index 82a68fc..8f79faa 100644 --- a/eventio/base.py +++ b/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) diff --git a/eventio/iact/objects.py b/eventio/iact/objects.py index 5676fb4..1445d4a 100644 --- a/eventio/iact/objects.py +++ b/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 diff --git a/eventio/simtel/simtelfile.py b/eventio/simtel/simtelfile.py index 5380d4b..dbe6fd8 100644 --- a/eventio/simtel/simtelfile.py +++ b/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): diff --git a/tests/resources/lst_muons_with_true_photons.simtel.zst b/tests/resources/lst_muons_with_true_photons.simtel.zst new file mode 100644 index 0000000..60ba0f9 Binary files /dev/null and b/tests/resources/lst_muons_with_true_photons.simtel.zst differ diff --git a/tests/simtel/test_simtelfile.py b/tests/simtel/test_simtelfile.py index 06c8d97..2d383f9 100644 --- a/tests/simtel/test_simtelfile.py +++ b/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 diff --git a/tests/test_iterator.py b/tests/test_iterator.py index 97cfda1..8758437 100644 --- a/tests/test_iterator.py +++ b/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)