Add jpp reader

km3io/reader.py

@@ -142,3 +142,69 @@ class AanetReader:
                 key.decode('utf8') for key in tracks_tree.keys()
             ]
         return self._tracks_keys
+
+
+class JppReader:
+    """Reader for Jpp ROOT files"""
+    def __init__(self, filename):
+        self.fobj = uproot.open(filename)
+        self._events = None
+
+    @property
+    def events(self):
+        if self._events is None:
+            tree = self.fobj["KM3NET_EVENT"]
+
+            headers = tree["KM3NETDAQ::JDAQEventHeader"].array(
+                uproot.interpret(tree["KM3NETDAQ::JDAQEventHeader"],
+                                 cntvers=True))
+            snapshot_hits = tree["snapshotHits"].array(
+                uproot.asjagged(uproot.astable(
+                    uproot.asdtype([("dom_id", "i4"), ("channel_id", "u1"),
+                                    ("time", "u4"), ("tot", "u1")])),
+                                skipbytes=10))
+            triggered_hits = tree["triggeredHits"].array(
+                uproot.asjagged(uproot.astable(
+                    uproot.asdtype([("dom_id", "i4"), ("channel_id", "u1"),
+                                    ("time", "u4"), ("tot", "u1"),
+                                    (" cnt", "u4"), (" vers", "u2"),
+                                    ("trigger_mask", "u8")])),
+                                skipbytes=10))
+            self._events = JppEvents(headers, snapshot_hits, triggered_hits)
+        return self._events
+
+
+class JppEvents:
+    """A simple wrapper for Jpp events"""
+    def __init__(self, headers, snapshot_hits, triggered_hits):
+        self.headers = headers
+        self.snapshot_hits = snapshot_hits
+        self.triggered_hits = triggered_hits
+
+    def __getitem__(self, item):
+        return JppEvent(self.headers[item], self.snapshot_hits[item],
+                        self.triggered_hits[item])
+
+    def __len__(self):
+        return len(self.headers)
+
+    def __str__(self):
+        return "Number of events: {}".format(len(self.headers))
+
+    def __repr__(self):
+        return str(self)
+
+
+class JppEvent:
+    """A wrapper for a Jpp event"""
+    def __init__(self, header, snapshot_hits, triggered_hits):
+        self.header = header
+        self.snapshot_hits = snapshot_hits
+        self.triggered_hits = triggered_hits
+
+    def __str__(self):
+        return "Jpp event with {} snapshot hits and {} triggered hits".format(
+            len(self.snapshot_hits), len(self.triggered_hits))
+
+    def __repr__(self):
+        return str(self)

tests/test_jpp.py

+import os
+import unittest
+
+from km3io import JppReader
+
+SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "samples")
+
+
+class TestJppEventsSnapshotHits(unittest.TestCase):
+    def setUp(self):
+        self.events = JppReader(os.path.join(SAMPLES_DIR,
+                                             "jpp_v12.0.0.root")).events
+        self.lengths = {0: 96, 1: 124, -1: 78}
+        self.total_item_count = 298
+
+    def test_reading_snapshot_hits(self):
+        hits = self.events.snapshot_hits
+
+        for event_id, length in self.lengths.items():
+            assert length == len(hits[event_id].dom_id)
+            assert length == len(hits[event_id].channel_id)
+            assert length == len(hits[event_id].time)
+
+    def test_total_item_counts(self):
+        hits = self.events.snapshot_hits
+
+        assert self.total_item_count == sum(hits.dom_id.count())
+        assert self.total_item_count == sum(hits.channel_id.count())
+        assert self.total_item_count == sum(hits.time.count())
+
+    def test_data_values(self):
+        hits = self.events.snapshot_hits
+
+        self.assertListEqual([806451572, 806451572, 806455814],
+                             list(hits.dom_id[0][:3]))
+        self.assertListEqual([10, 13, 0], list(hits.channel_id[0][:3]))
+        self.assertListEqual([1593234433, 1559680001, 3371422721],
+                             list(hits.time[0][:3]))
+
+    def test_channel_ids_have_valid_values(self):
+        hits = self.events.snapshot_hits
+
+        # channel IDs are always between [0, 30]
+        assert all(c >= 0 for c in hits.channel_id.min())
+        assert all(c < 31 for c in hits.channel_id.max())
+
+
+class TestJppEventsTriggeredHits(unittest.TestCase):
+    def setUp(self):
+        self.events = JppReader(os.path.join(SAMPLES_DIR,
+                                             "jpp_v12.0.0.root")).events
+        self.lengths = {0: 18, 1: 53, -1: 9}
+        self.total_item_count = 80
+
+    def test_data_lengths(self):
+        hits = self.events.triggered_hits
+
+        for event_id, length in self.lengths.items():
+            assert length == len(hits[event_id].dom_id)
+            assert length == len(hits[event_id].channel_id)
+            assert length == len(hits[event_id].time)
+            assert length == len(hits[event_id].trigger_mask)
+
+    def test_total_item_counts(self):
+        hits = self.events.triggered_hits
+
+        assert self.total_item_count == sum(hits.dom_id.count())
+        assert self.total_item_count == sum(hits.channel_id.count())
+        assert self.total_item_count == sum(hits.time.count())
+
+    def test_data_values(self):
+        hits = self.events.triggered_hits
+
+        self.assertListEqual([806451572, 806451572, 808432835],
+                             list(hits.dom_id[0][:3]))
+        self.assertListEqual([10, 13, 1], list(hits.channel_id[0][:3]))
+        self.assertListEqual([1593234433, 1559680001, 1978979329],
+                             list(hits.time[0][:3]))
+        self.assertListEqual([16, 16, 4], list(hits.trigger_mask[0][:3]))
+
+    def test_channel_ids_have_valid_values(self):
+        hits = self.events.triggered_hits
+
+        # channel IDs are always between [0, 30]
+        assert all(c >= 0 for c in hits.channel_id.min())
+        assert all(c < 31 for c in hits.channel_id.max())