Skip to content
Snippets Groups Projects

Resolve "uproot4 integration"

Merged Tamas Gal requested to merge 58-uproot4-integration-2 into master
Compare and Show latest version
5 files
+ 136
68
Compare changes
  • Side-by-side
  • Inline
Files
5
+ 17
9
@@ -79,7 +79,15 @@ class OfflineReader:
"mc_tracks": "mc_trks",
}
def __init__(self, f, index_chain=None, step_size=2000, keys=None, aliases=None, event_ctor=None):
def __init__(
self,
f,
index_chain=None,
step_size=2000,
keys=None,
aliases=None,
event_ctor=None,
):
"""OfflineReader class is an offline ROOT file wrapper
Parameters
@@ -187,10 +195,12 @@ class OfflineReader:
step_size=self._step_size,
aliases=self.aliases,
keys=self.keys(),
event_ctor=self._event_ctor
event_ctor=self._event_ctor,
)
if isinstance(key, str) and key.startswith("n_"): # group counts, for e.g. n_events, n_hits etc.
if isinstance(key, str) and key.startswith(
"n_"
): # group counts, for e.g. n_events, n_hits etc.
key = self._keyfor(key.split("n_")[1])
arr = self._fobj[self.event_path][key].array(uproot.AsDtype(">i4"))
return unfold_indices(arr, self._index_chain)
@@ -207,22 +217,21 @@ class OfflineReader:
if from_field in branch[key].keys():
fields.append(to_field)
log.debug(fields)
out = branch[key].arrays(
fields, aliases=self.special_branches[key]
)
out = branch[key].arrays(fields, aliases=self.special_branches[key])
else:
out = branch[self.aliases.get(key, key)].array()
return unfold_indices(out, self._index_chain)
def __iter__(self):
self._iterator_index = 0
self._events = self._event_generator()
return self
def _event_generator(self):
events = self._fobj[self.event_path]
group_count_keys = set(k for k in self.keys() if k.startswith("n_")) # special keys to make it easy to count subbranch lengths
group_count_keys = set(
k for k in self.keys() if k.startswith("n_")
) # special keys to make it easy to count subbranch lengths
log.debug("group_count_keys: %s", group_count_keys)
keys = set(
list(
@@ -293,7 +302,6 @@ class OfflineReader:
"""The raw number of events without any indexing/slicing magic"""
return len(self._fobj[self.event_path]["id"].array())
def __repr__(self):
length = len(self)
actual_length = self.__actual_len__()
Loading