Refactor offline

This refactoring includes some slicing and performance improvement attempts.

The slicing is still experimental and has some flaws:

• fancy indexing over thousands of elements might be horrible slow, eg f.events.tracks[idx_list, 0] where idx_list is a huge list of indices
• trivial problem of iterating: for track in f.evens.tracks does not stop after the last element
• meaningless indexing: f.events[5][23][420] is simply f.events[420] since the index/slice is not chained but overwritten when propagated. We need custom checks whether it makes sense
  • events are flat lists, so [x][y] does not make sense, it should throw an error
  • f.events[5:23][4] should effectively be the same as f.events[9] due to chaining of slice/index

km3io/offline.py

@@ -30,14 +30,14 @@ EVENTS_MAP = BranchMapper(name="events",
                           flat=True)
 
 SUBBRANCH_MAPS = [
-    BranchMapper(
-        name="tracks",
-        key="trks",
-        extra={},
-        exclude=EXCLUDE_KEYS + ['trks.usr_data', 'trks.usr', 'trks.fUniqueID', 'trks.fBits'],
-        update={},
-        attrparser=_nested_mapper,
-        flat=False),
+    BranchMapper(name="tracks",
+                 key="trks",
+                 extra={},
+                 exclude=EXCLUDE_KEYS +
+                 ['trks.usr_data', 'trks.usr', 'trks.fUniqueID', 'trks.fBits'],
+                 update={},
+                 attrparser=_nested_mapper,
+                 flat=False),
     BranchMapper(name="mc_tracks",
                  key="mc_trks",
                  extra={},
@@ -72,6 +72,111 @@ SUBBRANCH_MAPS = [
 ]
 
 
+class OfflineBranch(Branch):
+    @cached_property
+    def usr(self):
+        return Usr(self._mapper, self._branch, index=self._index)
+
+
+class Usr:
+    """Helper class to access AAObject `usr` stuff"""
+    def __init__(self, mapper, branch, index=None):
+        self._mapper = mapper
+        self._name = mapper.name
+        self._index = index
+        self._branch = branch
+        self._usr_names = []
+        self._usr_idx_lookup = {}
+
+        self._usr_key = 'usr' if mapper.flat else mapper.key + '.usr'
+
+        self._initialise()
+
+    def _initialise(self):
+        try:
+            self._branch[self._usr_key]
+            # This will raise a KeyError in old aanet files
+            # which has a different strucuter and key (usr_data)
+            # We do not support those (yet)
+        except (KeyError, IndexError):
+            print("The `usr` fields could not be parsed for the '{}' branch.".
+                  format(self._name))
+            return
+
+        if self._mapper.flat:
+            self._initialise_flat()
+
+    def _initialise_flat(self):
+        # Here, we assume that every event has the same names in the same order
+        # to massively increase the performance. This needs triple check if
+        # it's always the case.
+        self._usr_names = [
+            n.decode("utf-8")
+            for n in self._branch[self._usr_key + '_names'].lazyarray()[0]
+        ]
+        self._usr_idx_lookup = {
+            name: index
+            for index, name in enumerate(self._usr_names)
+        }
+
+        data = self._branch[self._usr_key].lazyarray()
+
+        if self._index is not None:
+            data = data[self._index]
+
+        self._usr_data = data
+
+        for name in self._usr_names:
+            setattr(self, name, self[name])
+
+    # def _initialise_nested(self):
+    #     self._usr_names = [
+    #         n.decode("utf-8") for n in self.branch['usr_names'].lazyarray(
+    #             # TODO this will be fixed soon in uproot,
+    #             # see https://github.com/scikit-hep/uproot/issues/465
+    #             uproot.asgenobj(
+    #                 uproot.SimpleArray(uproot.STLVector(uproot.STLString())),
+    #                 self.branch['usr_names']._context, 6),
+    #             basketcache=BASKET_CACHE)[0]
+    #     ]
+
+    def __getitem__(self, item):
+        if self._mapper.flat:
+            return self.__getitem_flat__(item)
+        return self.__getitem_nested__(item)
+
+    def __getitem_flat__(self, item):
+        if self._index is not None:
+            return self._usr_data[self._index][:, self._usr_idx_lookup[item]]
+        else:
+            return self._usr_data[:, self._usr_idx_lookup[item]]
+
+    def __getitem_nested__(self, item):
+        data = self._branch[self._usr_key + '_names'].lazyarray(
+            # TODO this will be fixed soon in uproot,
+            # see https://github.com/scikit-hep/uproot/issues/465
+            uproot.asgenobj(
+                uproot.SimpleArray(uproot.STLVector(uproot.STLString())),
+                self._branch[self._usr_key + '_names']._context, 6),
+            basketcache=BASKET_CACHE)
+        if self._index is None:
+            return data
+        else:
+            return data[self._index]
+
+    def keys(self):
+        return self._usr_names
+
+    def __str__(self):
+        entries = []
+        for name in self.keys():
+            entries.append("{}: {}".format(name, self[name]))
+        return '\n'.join(entries)
+
+    def __repr__(self):
+        return "<{}[{}]>".format(self.__class__.__name__, self._name)
+
+
 class OfflineReader:
     """reader for offline ROOT files"""
     def __init__(self, file_path=None):
@@ -90,7 +195,7 @@ class OfflineReader:
     @cached_property
     def events(self):
         """The `E` branch, containing all offline events."""
-        return Branch(self._tree,
+        return OfflineBranch(self._tree,
                       mapper=EVENTS_MAP,
                       subbranchmaps=SUBBRANCH_MAPS)
 
@@ -107,8 +212,6 @@ class OfflineReader:
             warnings.warn("Your file header has an unsupported format")
 
 
-
-
 class Header:
     """The header"""
     def __init__(self, header):
@@ -150,5 +253,3 @@ class Header:
             else:
                 lines.append("  {}: {}".format(key, value))
         return "\n".join(lines)
-
-