From 483eab58e6c5b41fda28dca989b3e79b1cf680ff Mon Sep 17 00:00:00 2001
From: Tamas Gal <tgal@km3net.de>
Date: Thu, 10 Dec 2020 23:24:58 +0100
Subject: [PATCH] First iteration prototype
---
km3io/offline.py | 4 ++--
km3io/rootio.py | 22 +++++++++++++++++-----
km3io/tools.py | 7 ++++++-
3 files changed, 25 insertions(+), 8 deletions(-)
diff --git a/km3io/offline.py b/km3io/offline.py
index 47c1f82..c593024 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -19,8 +19,8 @@ class OfflineReader(EventReader):
item_name = "OfflineEvent"
skip_keys = ["t", "AAObject"]
aliases = {
- "t_sec": "t.fSec",
- "t_ns": "t.fNanoSec",
+ "t_sec": "t/t.fSec",
+ "t_ns": "t/t.fNanoSec",
"usr": "AAObject/usr",
"usr_names": "AAObject/usr_names",
}
diff --git a/km3io/rootio.py b/km3io/rootio.py
index 1c2d1cf..25054ae 100644
--- a/km3io/rootio.py
+++ b/km3io/rootio.py
@@ -95,13 +95,28 @@ class EventReader:
def _initialise_keys(self):
skip_keys = set(self.skip_keys)
- toplevel_keys = set(k.split("/")[0] for k in self._fobj[self.event_path].keys())
+ all_keys = set(self._fobj[self.event_path].keys())
+ toplevel_keys = set(k.split("/")[0] for k in all_keys)
+ valid_aliases = {}
+ for fromkey, tokey in self.aliases.items():
+ if tokey in all_keys:
+ valid_aliases[fromkey] = tokey
+ self.aliases = valid_aliases
keys = (toplevel_keys - skip_keys).union(
- list(self.aliases.keys()) + list(self.nested_aliases)
+ list(valid_aliases) + list(self.nested_aliases)
)
for key in list(self.nested_branches) + list(self.nested_aliases):
keys.add("n_" + key)
# self._grouped_branches = {k for k in toplevel_keys - skip_keys if isinstance(self._fobj[self.event_path][k].interpretation, uproot.AsGrouped)}
+ valid_nested_branches = {}
+ for nested_key, aliases in self.nested_branches.items():
+ if nested_key in toplevel_keys:
+ valid_nested_branches[nested_key] = {}
+ subbranch_keys = self._fobj[self.event_path][nested_key].keys()
+ for fromkey, tokey in aliases.items():
+ if tokey in subbranch_keys:
+ valid_nested_branches[nested_key][fromkey] = tokey
+ self.nested_branches = valid_nested_branches
self._keys = keys
def keys(self):
@@ -187,7 +202,6 @@ class EventReader:
) # all top-level keys for regular branches
log.debug("keys: %s", keys)
log.debug("aliases: %s", self.aliases)
- # check for valid keys, e.g. `usr` is not always there and raises recursion error!
events_it = events.iterate(
keys, aliases=self.aliases, step_size=self._step_size
)
@@ -197,8 +211,6 @@ class EventReader:
) # dict-key ordering is an implementation detail
log.debug("nested_keys: %s", nested_keys)
for key in nested_keys:
- # print(f"adding {key} with keys {self.nested_branches[key].keys()} and aliases {self.nested_branches[key]}")
-
nested.append(
events[key].iterate(
self.nested_branches[key].keys(),
diff --git a/km3io/tools.py b/km3io/tools.py
index 4fcdc38..1246cd1 100644
--- a/km3io/tools.py
+++ b/km3io/tools.py
@@ -264,7 +264,10 @@ def best_track(tracks, startend=None, minmax=None, stages=None):
if minmax is not None:
m1 = mask(tracks.rec_stages, minmax=minmax)
- original_ndim = tracks.ndim
+ try:
+ original_ndim = tracks.ndim
+ except AttributeError:
+ original_ndim = 1
axis = 1 if original_ndim == 2 else 0
tracks = tracks[m1]
@@ -278,6 +281,8 @@ def best_track(tracks, startend=None, minmax=None, stages=None):
out = tracks[m3]
if original_ndim == 1:
+ if isinstance(out, ak.Record):
+ return out[:, 0]
return out[0]
return out[:, 0]
--
GitLab