WIP: Slicing and refactoring offline

Closes #37 (closed)

km3io/offline.py

@@ -11,7 +11,7 @@ BASKET_CACHE_SIZE = 110 * 1024**2
 
 BranchMapper = namedtuple(
     "BranchMapper",
-    ['name', 'key', 'extra', 'exclude', 'update', 'attrparser'])
+    ['name', 'key', 'extra', 'exclude', 'update', 'attrparser', 'flat'])
 
 
 def _nested_mapper(key):
@@ -20,14 +20,17 @@ def _nested_mapper(key):
 
 
 EXCLUDE_KEYS = set(["AAObject", "t", "fBits", "fUniqueID"])
-BRANCH_MAPS = [
+
+SUBBRANCH_MAPS = [
     BranchMapper("tracks", "trks", {}, ['trks.usr_data', 'trks.usr'], {},
-                 _nested_mapper),
+                 _nested_mapper, False),
     BranchMapper("mc_tracks", "mc_trks", {},
-                 ['mc_trks.usr_data', 'mc_trks.usr'], {}, _nested_mapper),
-    BranchMapper("hits", "hits", {}, ['hits.usr'], {}, _nested_mapper),
-    BranchMapper("mc_hits", "mc_hits", {}, ['mc_hits.usr'], {},
-                 _nested_mapper),
+                 ['mc_trks.usr_data', 'mc_trks.usr'], {}, _nested_mapper,
+                 False),
+    BranchMapper("hits", "hits", {}, ['hits.usr'], {}, _nested_mapper, False),
+    BranchMapper("mc_hits", "mc_hits", {},
+                 ['mc_hits.usr', 'mc_hits.dom_id', 'mc_hits.channel_id'], {},
+                 _nested_mapper, False),
     BranchMapper("events", "Evt", {
         't_sec': 't.fSec',
         't_ns': 't.fNanoSec'
@@ -36,9 +39,18 @@ BRANCH_MAPS = [
         'n_mc_hits': 'mc_hits',
         'n_tracks': 'trks',
         'n_mc_tracks': 'mc_trks'
-    }, lambda a: a),
+    }, lambda a: a, True),
 ]
 
+EVENTS_MAP = BranchMapper("events", "Evt", {
+        't_sec': 't.fSec',
+        't_ns': 't.fNanoSec'
+    }, [], {
+        'n_hits': 'hits',
+        'n_mc_hits': 'mc_hits',
+        'n_tracks': 'trks',
+        'n_mc_tracks': 'mc_trks'
+    }, lambda a: a, True)
 
 class cached_property:
     """A simple cache decorator for properties."""
@@ -54,11 +66,7 @@ class cached_property:
 
 class OfflineReader:
     """reader for offline ROOT files"""
-    def __init__(self,
-                 file_path=None,
-                 fobj=None,
-                 data=None,
-                 index=slice(None)):
+    def __init__(self, file_path=None, fobj=None, data=None, index=None):
         """ OfflineReader class is an offline ROOT file wrapper
 
         Parameters
@@ -80,10 +88,9 @@ class OfflineReader:
             self._tree = self._fobj[MAIN_TREE_NAME]
             self._data = data
 
-        for mapper in BRANCH_MAPS:
-            setattr(
-                self, mapper.name,
-                BranchElement(self._tree, mapper=mapper, index=self._index))
+    @cached_property
+    def events(self):
+        return Branch(self._tree, mapper=EVENTS_MAP, index=self._index, subbranches=SUBBRANCH_MAPS)
 
     @classmethod
     def from_index(cls, source, index):
@@ -106,7 +113,7 @@ class OfflineReader:
 
     def __len__(self):
         tree = self._fobj[MAIN_TREE_NAME]
-        if self._index == slice(None):
+        if self._index is None:
             return len(tree)
         else:
             return len(
@@ -124,10 +131,6 @@ class OfflineReader:
         else:
             warnings.warn("Your file header has an unsupported format")
 
-    @cached_property
-    def usr(self):
-        return Usr(self._tree)
-
     def get_best_reco(self):
         """returns the best reconstructed track fit data. The best fit is defined
         as the track fit with the maximum reconstruction stages. When "nan" is
@@ -421,10 +424,10 @@ class OfflineReader:
             are not found, None is returned as the stages index.
         """
         if mc is False:
-            stages_data = self.tracks.rec_stages
+            stages_data = self.events.tracks.rec_stages
 
         if mc is True:
-            stages_data = self.mc_tracks.rec_stages
+            stages_data = self.events.mc_tracks.rec_stages
 
         for trk_index, rec_stages in enumerate(stages_data):
             try:
@@ -469,25 +472,37 @@ class OfflineReader:
 
 class Usr:
     """Helper class to access AAObject usr stuff"""
-    def __init__(self, tree):
+    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.
+        # print("initialising usr for {}".format(name))
+        # print("Setting up usr")
+        self._name = name
         try:
+            tree['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['Evt']['usr_names'].array()[0]
+                n.decode("utf-8") for n in tree['usr_names'].lazyarray()[0]
             ]
         except (KeyError, IndexError):  # e.g. old aanet files
             self._usr_names = []
         else:
+            # print(" checking usr data")
             self._usr_idx_lookup = {
                 name: index
                 for index, name in enumerate(self._usr_names)
             }
-            self._usr_data = tree['Evt']['usr'].lazyarray(
+            data = tree['usr'].lazyarray(
                 basketcache=uproot.cache.ThreadSafeArrayCache(
                     BASKET_CACHE_SIZE))
+            if index is not None:
+                data = data[index]
+            self._usr_data = data
+            # print("    adding attributes")
             for name in self._usr_names:
+                # print("  setting {}".format(name))
                 setattr(self, name, self[name])
 
     def __getitem__(self, item):
@@ -502,6 +517,9 @@ class Usr:
             entries.append("{}: {}".format(name, self[name]))
         return '\n'.join(entries)
 
+    def __repr__(self):
+        return "<{}[{}]>".format(self.__class__.__name__, self._name)
+
 
 def _to_num(value):
     """Convert a value to a numerical one if possible"""
@@ -543,17 +561,24 @@ class Header:
         return "\n".join(lines)
 
 
-class BranchElement:
-    """wrapper for offline tracks"""
-    def __init__(self, tree, mapper, index=slice(None)):
+class Branch:
+    """Branch accessor class"""
+    def __init__(self, tree, mapper, index=None, subbranches=[]):
         self._tree = tree
         self._mapper = mapper
         self._index = index
         self._keymap = None
         self._branch = tree[mapper.key]
+        self._subbranches = subbranches
+        self._subbranch_keys = []
 
         self._initialise_keys()
 
+        for mapper in subbranches:
+            setattr(self, mapper.name,
+                    Branch(self._tree, mapper=mapper, index=self._index))
+            self._subbranch_keys.append(mapper.name)
+
     def _initialise_keys(self):
         """Create the keymap and instance attributes"""
         keys = set(k.decode('utf-8') for k in self._branch.keys()) - set(
@@ -569,28 +594,68 @@ class BranchElement:
 
         # self._EntryType = namedtuple(mapper.name[:-1], self.keys())
 
-        for key in self.keys():
+        for key in self._keymap.keys():
             # print("setting", self._mapper.name, key)
             setattr(self, key, self[key])
 
     def keys(self):
         return self._keymap.keys()
 
+    @cached_property
+    def usr(self):
+        return Usr(self._mapper.name, self._branch, index=self._index)
+
     def __getitem__(self, item):
         """Slicing magic a la numpy"""
         if isinstance(item, slice):
-            return self.__class__(self._tree, self._mapper, index=item)
+            return self.__class__(self._tree, self._mapper, index=item, subbranches=self._subbranches)
         if isinstance(item, int):
-            return {
-                key: self._branch[self._keymap[key]].array()[self._index, item]
-                for key in self.keys()
-            }
-        return self._branch[self._keymap[item]].lazyarray(
-            basketcache=uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE))[
-                self._index]
+            # A bit ugly, but whatever works
+            if self._mapper.flat:
+                if self._index is None:
+                    dct = {
+                        key: self._branch[self._keymap[key]].array()
+                        for key in self.keys()
+                    }
+                else:
+                    dct = {
+                        key:
+                        self._branch[self._keymap[key]].array()[self._index]
+                        for key in self.keys()
+                    }
+                return BranchElement(self._mapper.name, dct)[item]
+            else:
+                if self._index is None:
+                    dct = {
+                        key: self._branch[self._keymap[key]].array()[item]
+                        for key in self.keys()
+                    }
+                else:
+                    dct = {
+                        key:
+                        self._branch[self._keymap[key]].array()[self._index,
+                                                                item]
+                        for key in self.keys()
+                    }
+                return BranchElement(self._mapper.name, dct)
+
+        if isinstance(item, tuple):
+            return self[item[0]][item[1]]
+
+        if isinstance(item, str):
+            item = self._keymap[item]
+
+            out = self._branch[item].lazyarray(
+                basketcache=uproot.cache.ThreadSafeArrayCache(
+                    BASKET_CACHE_SIZE))
+            if self._index is not None:
+                out = out[self._index]
+            return out
+
+        return self.__class__(self._tree, self._mapper, index=np.array(item), subbranches=self._subbranches)
 
     def __len__(self):
-        if self._index == slice(None):
+        if self._index is None:
             return len(self._branch)
         else:
             return len(
@@ -603,3 +668,44 @@ class BranchElement:
         return "<{}[{}]: {} parsed elements>".format(self.__class__.__name__,
                                                      self._mapper.name,
                                                      len(self))
+
+
+class BranchElement:
+    """Represents a single branch element
+
+    Parameters
+    ----------
+    name: str
+        The name of the branch
+    dct: dict (keys=attributes, values=arrays of values)
+        The data
+    index: slice
+        The slice mask to be applied to the sub-arrays
+    """
+    def __init__(self, name, dct, index=None, subbranches=[]):
+        self._dct = dct
+        self._name = name
+        self._index = index
+        self.ItemConstructor = namedtuple(self._name[:-1], dct.keys())
+        if index is None:
+            for key, values in dct.items():
+                setattr(self, key, values)
+        else:
+            for key, values in dct.items():
+                setattr(self, key, values[index])
+
+    def __getitem__(self, item):
+        if isinstance(item, slice):
+            return self.__class__(self._name, self._dct, index=item)
+        if isinstance(item, int):
+            if self._index is None:
+                return self.ItemConstructor(
+                    **{k: v[item]
+                    for k, v in self._dct.items()})
+            else:
+                return self.ItemConstructor(
+                    **{k: v[self._index][item]
+                    for k, v in self._dct.items()})
+
+    def __repr__(self):
+        return "<{}[{}]>".format(self.__class__.__name__, self._name)