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

@@ -2,11 +2,15 @@ from collections import namedtuple
 import uproot
 import warnings
 from .definitions import mc_header
-from .tools import Branch, BranchMapper, cached_property, _to_num
+from .tools import Branch, BranchMapper, cached_property, _to_num, _unfold_indices
 
 MAIN_TREE_NAME = "E"
 EXCLUDE_KEYS = ["AAObject", "t", "fBits", "fUniqueID"]
 
+# 110 MB based on the size of the largest basket found so far in km3net
+BASKET_CACHE_SIZE = 110 * 1024**2
+BASKET_CACHE = uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)
+
 
 def _nested_mapper(key):
     """Maps a key in the ROOT file to another key (e.g. trks.pos.x -> pos_x)"""
@@ -30,14 +34,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 +76,110 @@ SUBBRANCH_MAPS = [
 ]
 
 
+class OfflineBranch(Branch):
+    @cached_property
+    def usr(self):
+        return Usr(self._mapper, self._branch, index_chain=self._index_chain)
+
+
+class Usr:
+    """Helper class to access AAObject `usr` stuff"""
+    def __init__(self, mapper, branch, index_chain=None):
+        self._mapper = mapper
+        self._name = mapper.name
+        self._index_chain = [] if index_chain is None else index_chain
+        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_chain:
+            data = _unfold_indices(data, self._index_chain)
+
+        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_chain:
+            return _unfold_indices(
+                self._usr_data, self._index_chain)[:,
+                                                   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)
+        return _unfold_indices(data, self._index_chain)
+
+    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,9 +198,9 @@ class OfflineReader:
     @cached_property
     def events(self):
         """The `E` branch, containing all offline events."""
-        return Branch(self._tree,
-                      mapper=EVENTS_MAP,
-                      subbranchmaps=SUBBRANCH_MAPS)
+        return OfflineBranch(self._tree,
+                             mapper=EVENTS_MAP,
+                             subbranchmaps=SUBBRANCH_MAPS)
 
     @cached_property
     def header(self):
@@ -107,8 +215,6 @@ class OfflineReader:
             warnings.warn("Your file header has an unsupported format")
 
 
-
-
 class Header:
     """The header"""
     def __init__(self, header):
@@ -150,5 +256,3 @@ class Header:
             else:
                 lines.append("  {}: {}".format(key, value))
         return "\n".join(lines)
-
-