WIP: Resolve "uproot4 integration"

Closes #58 (closed)

km3io/offline.py

@@ -3,7 +3,7 @@ import uproot4 as uproot
 import warnings
 
 from .definitions import mc_header
-from .tools import cached_property
+from .tools import cached_property, to_num
 
 
 class OfflineReader:
@@ -13,22 +13,22 @@ class OfflineReader:
     item_name = "OfflineEvent"
     skip_keys = ["t", "AAObject"]
     aliases = {
-        "t_s": "t.fSec",
+        "t_sec": "t.fSec",
         "t_ns": "t.fNanoSec",
-        "usr": "AAObject/usr",
-        "usr_names": "AAObject/usr_names",
     }
     special_branches = {
         "hits": {
+            "id": "hits.id",
             "channel_id": "hits.channel_id",
             "dom_id": "hits.dom_id",
-            "time": "hits.t",
+            "t": "hits.t",
             "tot": "hits.tot",
-            "triggered": "hits.trig",  # non-zero if the hit is a triggered hit
+            "trig": "hits.trig",  # non-zero if the hit is a triggered hit
         },
         "mc_hits": {
+            "id": "mc_hits.id",
             "pmt_id": "mc_hits.pmt_id",
-            "time": "mc_hits.t",  # hit time (MC truth)
+            "t": "mc_hits.t",  # hit time (MC truth)
             "a": "mc_hits.a",  # hit amplitude (in p.e.)
             "origin": "mc_hits.origin",  # track id of the track that created this hit
             "pure_t": "mc_hits.pure_t",  # photon time before pmt simultion
@@ -90,6 +90,14 @@ class OfflineReader:
         self._uuid = self._fobj._file.uuid
         self._iterator_index = 0
         self._keys = None
+        self._grouped_counts = {}  # TODO: e.g. {"events": [3, 66, 34]}
+
+        if "E/Evt/AAObject/usr" in self._fobj:
+            if ak.count(f["E/Evt/AAObject/usr"].array()) > 0:
+                self.aliases.update({
+                    "usr": "AAObject/usr",
+                    "usr_names": "AAObject/usr_names",
+                })
 
         self._initialise_keys()
 
@@ -104,10 +112,14 @@ class OfflineReader:
         )
 
     def _initialise_keys(self):
+        skip_keys = set(self.skip_keys)
         toplevel_keys = set(k.split("/")[0] for k in self._fobj[self.event_path].keys())
-        keys = (toplevel_keys - set(self.skip_keys)).union(
+        keys = (toplevel_keys - skip_keys).union(
             list(self.aliases.keys()) + list(self.special_aliases)
         )
+        for key in list(self.special_branches) + list(self.special_aliases):
+            keys.add("n_" + key)
+        # self._grouped_branches = {k for k in toplevel_keys - skip_keys if isinstance(self._fobj[self.event_path][k].interpretation, uproot.AsGrouped)}
         self._keys = keys
 
     def keys(self):
@@ -124,6 +136,7 @@ class OfflineReader:
 
     def __getattr__(self, attr):
         attr = self._keyfor(attr)
+        # if attr in self.keys() or (attr.startswith("n_") and self._keyfor(attr.split("n_")[1]) in self._grouped_branches):
         if attr in self.keys():
             return self.__getitem__(attr)
         raise AttributeError(
@@ -131,6 +144,10 @@ class OfflineReader:
         )
 
     def __getitem__(self, key):
+        if key.startswith("n_"):  # group counts, for e.g. n_events, n_hits etc.
+            key = self._keyfor(key.split("n_")[1])
+            return self._fobj[self.event_path][key].array(uproot.AsDtype(">i4"))
+
         key = self._keyfor(key)
         branch = self._fobj[self.event_path]
         # These are special branches which are nested, like hits/trks/mc_trks
@@ -149,11 +166,13 @@ class OfflineReader:
 
     def _event_generator(self):
         events = self._fobj[self.event_path]
-        keys = list(
+        group_count_keys = set(k for k in self.keys() if k.startswith("n_"))
+        keys = set(list(
             set(self.keys())
             - set(self.special_branches.keys())
             - set(self.special_aliases)
-        ) + list(self.aliases.keys())
+            - group_count_keys
+        ) + list(self.aliases.keys()))
         events_it = events.iterate(keys, aliases=self.aliases, step_size=self.step_size)
         specials = []
         special_keys = (
@@ -167,14 +186,20 @@ class OfflineReader:
                     step_size=self.step_size,
                 )
             )
+        group_counts = {}
+        for key in group_count_keys:
+            group_counts[key] = iter(self[key])
         for event_set, *special_sets in zip(events_it, *specials):
             for _event, *special_items in zip(event_set, *special_sets):
-                data = {
-                    **{k: _event[k] for k in keys},
-                    **{k: i for (k, i) in zip(special_keys, special_items)},
-                }
+                data = {}
+                for k in keys:
+                    data[k] = _event[k]
+                for (k, i) in zip(special_keys, special_items):
+                    data[k] = i
                 for tokey, fromkey in self.special_aliases.items():
                     data[tokey] = data[fromkey]
+                for key in group_counts:
+                    data[key] = next(group_counts[key])
                 yield self._event_ctor(**data)
 
     def __next__(self):