diff --git a/src/ctapipe_io_zfits/dl0.py b/src/ctapipe_io_zfits/dl0.py index 9cbab80..605dfb5 100644 --- a/src/ctapipe_io_zfits/dl0.py +++ b/src/ctapipe_io_zfits/dl0.py @@ -236,8 +236,11 @@ def scheduling_blocks(self) -> Dict[int, SchedulingBlockContainer]: return self._scheduling_blocks def _generator(self): - for subarray_trigger in self._subarray_trigger_file.SubarrayEvents: + for count, subarray_trigger in enumerate( + self._subarray_trigger_file.SubarrayEvents + ): array_event = ArrayEventContainer( + count=count, index=EventIndexContainer( obs_id=subarray_trigger.obs_id, event_id=subarray_trigger.event_id ), @@ -337,7 +340,7 @@ def observation_blocks(self) -> Dict[int, ObservationBlockContainer]: def scheduling_blocks(self) -> Dict[int, SchedulingBlockContainer]: return self._scheduling_blocks - def _fill_event(self, zfits_event) -> ArrayEventContainer: + def _fill_event(self, count, zfits_event) -> ArrayEventContainer: tel_id = self.tel_id # until ctapipe allows telescope event sources # we have to fill an arrayevent with just one telescope here @@ -345,6 +348,7 @@ def _fill_event(self, zfits_event) -> ArrayEventContainer: zfits_event.event_time_s, zfits_event.event_time_qns ) array_event = ArrayEventContainer( + count=count, index=EventIndexContainer( obs_id=self.obs_id, event_id=zfits_event.event_id, @@ -364,5 +368,5 @@ def _fill_event(self, zfits_event) -> ArrayEventContainer: return array_event def _generator(self): - for event in self._multi_file: - yield self._fill_event(event) + for count, zfits_event in enumerate(self._multi_file): + yield self._fill_event(count, zfits_event) diff --git a/src/ctapipe_io_zfits/tests/test_dl0.py b/src/ctapipe_io_zfits/tests/test_dl0.py index c5f37e6..3f420c6 100644 --- a/src/ctapipe_io_zfits/tests/test_dl0.py +++ b/src/ctapipe_io_zfits/tests/test_dl0.py @@ -31,7 +31,8 @@ def test_subarray_events(dummy_dl0): with EventSource(dummy_dl0) as source: n_read = 0 - for array_event in source: + for i, array_event in enumerate(source): + assert array_event.count == i assert array_event.index.obs_id == 456 assert array_event.index.event_id == n_read + 1 dt = np.abs(array_event.trigger.time - time).to(u.ns) @@ -74,7 +75,8 @@ def test_telescope_event_source(dummy_tel_file): with EventSource(dummy_tel_file) as source: assert isinstance(source, ProtozfitsDL0TelescopeEventSource) - for event in source: + for i, event in enumerate(source): + assert event.count == i assert event.dl0.tel.keys() == {1}