From af2d04e424214d005ccdd2ca35ea98689bb72222 Mon Sep 17 00:00:00 2001
From: Tamas Gal <tgal@km3net.de>
Date: Wed, 15 Apr 2020 16:59:22 +0200
Subject: [PATCH] Add option to auto-convert keys to awkward
---
km3io/offline.py | 3 ++-
km3io/tools.py | 10 +++++++++-
2 files changed, 11 insertions(+), 2 deletions(-)
diff --git a/km3io/offline.py b/km3io/offline.py
index 73cb13c..fb2c05a 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 e7bd86d..478faad 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):
--
GitLab