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

@@ -12,7 +12,7 @@ EXCLUDE_KEYS = set(["AAObject", "t", "fBits", "fUniqueID"])
 
 BranchMapper = namedtuple(
     "BranchMapper",
-    ['name', 'key', 'extra', 'exclude', 'update', 'attrparser'])
+    ['name', 'key', 'extra', 'exclude', 'update', 'attrparser', 'flat'])
 
 
 def _nested_mapper(key):
@@ -33,7 +33,8 @@ EVENTS_MAP = BranchMapper(name="events",
                               'n_tracks': 'trks',
                               'n_mc_tracks': 'mc_trks'
                           },
-                          attrparser=lambda a: a)
+                          attrparser=lambda a: a,
+                          flat=True)
 
 SUBBRANCH_MAPS = [
     BranchMapper(
@@ -42,7 +43,8 @@ SUBBRANCH_MAPS = [
         extra={},
         exclude=['trks.usr_data', 'trks.usr', 'trks.fUniqueID', 'trks.fBits'],
         update={},
-        attrparser=_nested_mapper),
+        attrparser=_nested_mapper,
+        flat=False),
     BranchMapper(name="mc_tracks",
                  key="mc_trks",
                  extra={},
@@ -51,7 +53,8 @@ SUBBRANCH_MAPS = [
                      'mc_trks.fitinf', 'mc_trks.fUniqueID', 'mc_trks.fBits'
                  ],
                  update={},
-                 attrparser=_nested_mapper),
+                 attrparser=_nested_mapper,
+                 flat=False),
     BranchMapper(name="hits",
                  key="hits",
                  extra={},
@@ -60,7 +63,8 @@ SUBBRANCH_MAPS = [
                      'hits.pure_a', 'hits.fUniqueID', 'hits.fBits'
                  ],
                  update={},
-                 attrparser=_nested_mapper),
+                 attrparser=_nested_mapper,
+                 flat=False),
     BranchMapper(name="mc_hits",
                  key="mc_hits",
                  extra={},
@@ -70,7 +74,8 @@ SUBBRANCH_MAPS = [
                      'mc_hits.fUniqueID', 'mc_hits.fBits'
                  ],
                  update={},
-                 attrparser=_nested_mapper),
+                 attrparser=_nested_mapper,
+                 flat=False),
 ]
 
 
@@ -123,37 +128,76 @@ class OfflineReader:
 
 class Usr:
     """Helper class to access AAObject `usr`` stuff"""
-    def __init__(self, name, tree, index=None):
-        # 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; the usr-format is simply a very bad design.
-        self._name = name
+
+    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 = {}
+
         try:
-            tree['usr']  # This will raise a KeyError in old aanet files
+            branch['usr']
+            # This will raise a KeyError in old aanet files
             # which has a different strucuter and key (usr_data)
-            # We do not support those...
-            self._usr_names = [
-                n.decode("utf-8") for n in tree['usr_names'].lazyarray(
-                    basketcache=BASKET_CACHE)[0]
-            ]
-        except (KeyError, IndexError):  # e.g. old aanet files
-            print("The `usr` fields could not be parsed for the '{}' branch."
-                  .format(name))
-            self._usr_names = []
+            # We do not support those
+        except (KeyError, IndexError):
+            print("The `usr` fields could not be parsed for the '{}' branch.".
+                  format(self._name))
+            return
+
+        if mapper.flat:
+            self._initialise_flat()
         else:
-            self._usr_idx_lookup = {
-                name: index
-                for index, name in enumerate(self._usr_names)
-            }
-            data = tree['usr'].lazyarray(basketcache=BASKET_CACHE)
-            if index is not None:
-                data = data[index]
-            self._usr_data = data
-            for name in self._usr_names:
-                setattr(self, name, self[name])
+            # self._initialise_nested()
+            # branch[self._mapper.key + '.usr']
+            pass
+
+    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['usr_names'].lazyarray(
+                basketcache=BASKET_CACHE)[0]
+        ]
+        self._usr_idx_lookup = {
+            name: index
+            for index, name in enumerate(self._usr_names)
+        }
+
+        data = self._branch['usr'].lazyarray(basketcache=BASKET_CACHE)
+
+        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]
+        ]
+        self.__getitem__ = self.__getitem_nested__
 
     def __getitem__(self, item):
-        return self._usr_data[:, self._usr_idx_lookup[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 keys(self):
         return self._usr_names
@@ -259,7 +303,7 @@ class Branch:
 
     @cached_property
     def usr(self):
-        return Usr(self._mapper.name, self._branch, index=self._index)
+        return Usr(self._mapper, self._branch, index=self._index)
 
     def __getattribute__(self, attr):
         if attr.startswith("_"):  # let all private and magic methods pass