diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1b51d34bc53910b15b7ffa670fe6227a2ab797de..ec80e3d63f8e5a542c4efe4ac03cfc172ce1047b 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,9 +1,11 @@ Unreleased changes ------------------ * ``km3net-dataformat`` updated to 2.0.0 +* mother ID and status are now read out for MC tracks Version 0 --------- + 0.20.0 / 2021-02-18 ~~~~~~~~~~~~~~~~~~~ * The fields ``.tdc``, ``.pos_{xyz}`` and ``.dir_{xyz}`` in ``.hits`` are diff --git a/km3io/offline.py b/km3io/offline.py index f1d2bd9dc1ecc7f5edb80994a343e94a843034d9..14c6318551ecfdaac2b51c3bffc78c68a60369cf 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -79,8 +79,8 @@ class OfflineReader(EventReader): "E": "mc_trks.E", "t": "mc_trks.t", "len": "mc_trks.len", - # "status": "mc_trks.status", # TODO: check this - # "mother_id": "mc_trks.mother_id", # TODO: check this + "status": "mc_trks.status", + "mother_id": "mc_trks.mother_id", "pdgid": "mc_trks.type", "hit_ids": "mc_trks.hit_ids", "usr": "mc_trks.usr", # TODO: trouble with uproot4 diff --git a/tests/test_offline.py b/tests/test_offline.py index 5762bf3174be4f2d26994bb0b26c58884f51f1fb..8adda39facd746a143ce7f0f69f38bb1c4b00b94 100644 --- a/tests/test_offline.py +++ b/tests/test_offline.py @@ -18,6 +18,9 @@ OFFLINE_MC_TRACK_USR = OfflineReader( ) ) OFFLINE_NUMUCC = OfflineReader(data_path("offline/numucc.root")) # with mc data +OFFLINE_MC_TRACK = OfflineReader( + data_path("gseagen/gseagen_v7.0.0_numuCC_diffuse.aa.root") +) class TestOfflineReader(unittest.TestCase): @@ -412,7 +415,11 @@ class TestOfflineTracks(unittest.TestCase): self.f = OFFLINE_FILE self.tracks = OFFLINE_FILE.events.tracks self.tracks_numucc = OFFLINE_NUMUCC + self.mc_tracks = OFFLINE_MC_TRACK.mc_tracks + self.mc_tracks_old = OFFLINE_MC_TRACK_USR.mc_tracks self.n_events = 10 + self.status = [100, 5, 11, 15, 1, 12, 12, 12, 12, 12] + self.mother_id = [-1, -1, 1, 1, 0, 2, 5, 5, 6, 8] def test_fields(self): for field in [ @@ -433,6 +440,18 @@ class TestOfflineTracks(unittest.TestCase): ]: getattr(self.tracks, field) + def test_status(self): + assert np.allclose(self.status, self.mc_tracks[1].status[:10].tolist()) + + def test_mother_id(self): + assert np.allclose(self.mother_id, self.mc_tracks[1].mother_id[:10].tolist()) + + def test_attribute_error_raised_for_older_files(self): + with self.assertRaises(AttributeError): + self.mc_tracks_old[1].mother_id + with self.assertRaises(AttributeError): + self.mc_tracks_old[1].status + def test_item_selection(self): self.assertListEqual( list(self.tracks[0].dir_z[:2]), [-0.872885221293917, -0.872885221293917]