diff --git a/km3io/offline.py b/km3io/offline.py index 73cb13c3bf9484180d03ecd27b4e42d60d143337..fb2c05a29613278833a8f163ede2b668d0f6f58e 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -38,7 +38,8 @@ SUBBRANCH_MAPS = [ exclude=EXCLUDE_KEYS + ['trks.usr_data', 'trks.usr', 'trks.fUniqueID', 'trks.fBits'], attrparser=_nested_mapper, - flat=False), + flat=False, + toawkward=['fitinf', 'rec_stages']), BranchMapper(name="mc_tracks", key="mc_trks", exclude=EXCLUDE_KEYS + [ diff --git a/km3io/tools.py b/km3io/tools.py index e7bd86dc295cf16fbae60185f1dbb4b74b5a002b..478faad65f078f5b2c77eceb902381f2d7b82240 100644 --- a/km3io/tools.py +++ b/km3io/tools.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 from collections import namedtuple import numpy as np +import awkward1 as ak import uproot # 110 MB based on the size of the largest basket found so far in km3net BASKET_CACHE_SIZE = 110 * 1024**2 @@ -44,7 +45,8 @@ class BranchMapper: update=None, attrparser=None, flat=True, - interpretations=None): + interpretations=None, + toawkward=None): """ Mapper helper for keys in a ROOT branch. @@ -68,6 +70,9 @@ class BranchMapper: The function to be used to create attribute names. This is only needed if unsupported characters are present, like ``.``, which would prevent setting valid Python attribute names. + toawkward: ``None``, ``list(str)`` + List of keys to convert to awkward arrays (recommended for + doubly ragged arrays) """ self.name = name self.key = key @@ -78,6 +83,7 @@ class BranchMapper: self.attrparser = (lambda x: x) if attrparser is None else attrparser self.flat = flat self.interpretations = {} if interpretations is None else interpretations + self.toawkward = [] if toawkward is None else toawkward class Branch: @@ -145,6 +151,8 @@ class Branch: out = self._branch[self._keymap[key]].lazyarray( interpretation=self._mapper.interpretations.get(key), basketcache=BASKET_CACHE) + if self._index_chain is not None and key in self._mapper.toawkward: + out = ak.from_iter(out) return _unfold_indices(out, self._index_chain) def __getitem__(self, item):