Skip to content
Snippets Groups Projects

WIP: Slicing and refactoring offline

Closed WIP: Slicing and refactoring offline
37-user-parameters-seem-to-be-transposed into master
Closed Tamas Gal requested to merge 37-user-parameters-seem-to-be-transposed into master
Viewing commit f882f362
Prev Next
Show latest version
1 file
+ 1
− 1
Compare changes
  • Side-by-side
  • Inline
km3io/offline.py
+ 225
− 801
from collections import namedtuple
import uproot
import numpy as np
import warnings
import km3io.definitions.trigger
import km3io.definitions.fitparameters
import km3io.definitions.reconstruction
from .definitions import mc_header
MAIN_TREE_NAME = "E"
# 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2
BASKET_CACHE = uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)
EXCLUDE_KEYS = set(["AAObject", "t", "fBits", "fUniqueID"])
BranchMapper = namedtuple(
"BranchMapper",
['name', 'key', 'extra', 'exclude', 'update', 'attrparser'])
def _nested_mapper(key):
"""Maps a key in the ROOT file to another key (e.g. trks.pos.x -> pos_x)"""
return '_'.join(key.split('.')[1:])
EVENTS_MAP = BranchMapper(name="events",
key="Evt",
extra={
't_sec': 't.fSec',
't_ns': 't.fNanoSec'
},
exclude=[],
update={
'n_hits': 'hits',
'n_mc_hits': 'mc_hits',
'n_tracks': 'trks',
'n_mc_tracks': 'mc_trks'
},
attrparser=lambda a: a)
SUBBRANCH_MAPS = [
BranchMapper(
name="tracks",
key="trks",
extra={},
exclude=['trks.usr_data', 'trks.usr', 'trks.fUniqueID', 'trks.fBits'],
update={},
attrparser=_nested_mapper),
BranchMapper(name="mc_tracks",
key="mc_trks",
extra={},
exclude=[
'mc_trks.usr_data', 'mc_trks.usr', 'mc_trks.rec_stages',
'mc_trks.fitinf', 'mc_trks.fUniqueID', 'mc_trks.fBits'
],
update={},
attrparser=_nested_mapper),
BranchMapper(name="hits",
key="hits",
extra={},
exclude=[
'hits.usr', 'hits.pmt_id', 'hits.origin', 'hits.a',
'hits.pure_a', 'hits.fUniqueID', 'hits.fBits'
],
update={},
attrparser=_nested_mapper),
BranchMapper(name="mc_hits",
key="mc_hits",
extra={},
exclude=[
'mc_hits.usr', 'mc_hits.dom_id', 'mc_hits.channel_id',
'mc_hits.tdc', 'mc_hits.tot', 'mc_hits.trig',
'mc_hits.fUniqueID', 'mc_hits.fBits'
],
update={},
attrparser=_nested_mapper),
]
class cached_property:
@@ -22,188 +86,9 @@ class cached_property:
return prop
class OfflineKeys:
"""wrapper for offline keys"""
def __init__(self, tree):
"""OfflineKeys is a class that reads all the available keys in an offline
file and adapts the keys format to Python format.
Parameters
----------
tree : uproot.TTree
The main ROOT tree.
"""
self._tree = tree
def __str__(self):
return '\n'.join([
"Events keys are:\n\t" + "\n\t".join(self.events_keys),
"Hits keys are:\n\t" + '\n\t'.join(self.hits_keys),
"Tracks keys are:\n\t" + '\n\t'.join(self.tracks_keys),
"Mc hits keys are:\n\t" + '\n\t'.join(self.mc_hits_keys),
"Mc tracks keys are:\n\t" + '\n\t'.join(self.mc_tracks_keys)
])
def __repr__(self):
return "<{}>".format(self.__class__.__name__)
def _get_keys(self, tree, fake_branches=None):
"""Get tree keys except those in fake_branches
Parameters
----------
tree : uproot.Tree
The tree to look for keys
fake_branches : list of str or None
The fake branches to ignore
Returns
-------
list of str
The keys of the tree.
"""
keys = []
for key in tree.keys():
key = key.decode('utf-8')
if fake_branches is not None and key in fake_branches:
continue
keys.append(key)
return keys
@cached_property
def events_keys(self):
"""reads events keys from an offline file.
Returns
-------
list of str
list of all events keys found in an offline file,
except those found in fake branches.
"""
fake_branches = ['Evt', 'AAObject', 'TObject', 't']
t_baskets = ['t.fSec', 't.fNanoSec']
tree = self._tree['Evt']
return self._get_keys(self._tree['Evt'], fake_branches) + t_baskets
@cached_property
def hits_keys(self):
"""reads hits keys from an offline file.
Returns
-------
list of str
list of all hits keys found in an offline file,
except those found in fake branches.
"""
fake_branches = ['hits.usr', 'hits.usr_names']
return self._get_keys(self._tree['hits'], fake_branches)
@cached_property
def tracks_keys(self):
"""reads tracks keys from an offline file.
Returns
-------
list of str
list of all tracks keys found in an offline file,
except those found in fake branches.
"""
# a solution can be tree['trks.usr_data'].array(
# uproot.asdtype(">i4"))
fake_branches = ['trks.usr_data', 'trks.usr', 'trks.usr_names']
return self._get_keys(self._tree['Evt']['trks'], fake_branches)
@cached_property
def mc_hits_keys(self):
"""reads mc hits keys from an offline file.
Returns
-------
list of str
list of all mc hits keys found in an offline file,
except those found in fake branches.
"""
fake_branches = ['mc_hits.usr', 'mc_hits.usr_names']
return self._get_keys(self._tree['Evt']['mc_hits'], fake_branches)
@cached_property
def mc_tracks_keys(self):
"""reads mc tracks keys from an offline file.
Returns
-------
list of str
list of all mc tracks keys found in an offline file,
except those found in fake branches.
"""
fake_branches = [
'mc_trks.usr_data', 'mc_trks.usr', 'mc_trks.usr_names'
]
return self._get_keys(self._tree['Evt']['mc_trks'], fake_branches)
@cached_property
def valid_keys(self):
"""constructs a list of all valid keys to be read from an offline event file.
Returns
-------
list of str
list of all valid keys.
"""
return (self.events_keys + self.hits_keys + self.tracks_keys +
self.mc_tracks_keys + self.mc_hits_keys)
@cached_property
def fit_keys(self):
"""constructs a list of fit parameters, not yet outsourced in an offline file.
Returns
-------
list of str
list of all "trks.fitinf" keys.
"""
return sorted(km3io.definitions.fitparameters.data,
key=km3io.definitions.fitparameters.data.get,
reverse=False)
@cached_property
def cut_hits_keys(self):
"""adapts hits keys for instance variables format in a Python class.
Returns
-------
list of str
list of adapted hits keys.
"""
return [k.split('hits.')[1].replace('.', '_') for k in self.hits_keys]
@cached_property
def cut_tracks_keys(self):
"""adapts tracks keys for instance variables format in a Python class.
Returns
-------
list of str
list of adapted tracks keys.
"""
return [
k.split('trks.')[1].replace('.', '_') for k in self.tracks_keys
]
@cached_property
def cut_events_keys(self):
"""adapts events keys for instance variables format in a Python class.
Returns
-------
list of str
list of adapted events keys.
"""
return [k.replace('.', '_') for k in self.events_keys]
class OfflineReader:
"""reader for offline ROOT files"""
def __init__(self, file_path=None, fobj=None, data=None):
def __init__(self, file_path=None):
""" OfflineReader class is an offline ROOT file wrapper
Parameters
@@ -213,485 +98,57 @@ class OfflineReader:
path-like object that points to the file.
"""
if file_path is not None:
self._fobj = uproot.open(file_path)
self._tree = self._fobj[MAIN_TREE_NAME]
self._data = self._tree.lazyarrays(
basketcache=uproot.cache.ThreadSafeArrayCache(
BASKET_CACHE_SIZE))
else:
self._fobj = fobj
self._tree = self._fobj[MAIN_TREE_NAME]
self._data = data
@classmethod
def from_index(cls, source, index):
"""Create an instance with a subtree of a given index
Parameters
----------
source: ROOTDirectory
The source file.
index: index or slice
The index or slice to create the subtree.
"""
instance = cls(fobj=source._fobj, data=source._data[index])
return instance
def __getitem__(self, index):
return OfflineReader.from_index(source=self, index=index)
self._fobj = uproot.open(file_path)
self._tree = self._fobj[MAIN_TREE_NAME]
def __len__(self):
return len(self._data)
@cached_property
def events(self):
"""The `E` branch, containing all offline events."""
return Branch(self._tree,
mapper=EVENTS_MAP,
subbranchmaps=SUBBRANCH_MAPS)
@cached_property
def header(self):
"""The file header"""
if 'Head' in self._fobj:
header = {}
for n, x in self._fobj['Head']._map_3c_string_2c_string_3e_.items(
):
header[n.decode("utf-8")] = x.decode("utf-8").strip()
return header
return Header(header)
else:
warnings.warn("Your file header has an unsupported format")
@cached_property
def keys(self):
"""wrapper for all keys in an offline file.
Returns
-------
Class
OfflineKeys.
"""
return OfflineKeys(self._tree)
@cached_property
def events(self):
"""wrapper for offline events.
Returns
-------
Class
OfflineEvents.
"""
return OfflineEvents(
self.keys.cut_events_keys,
[self._data[key] for key in self.keys.events_keys])
@cached_property
def hits(self):
"""wrapper for offline hits.
Returns
-------
Class
OfflineHits.
"""
return OfflineHits(self.keys.cut_hits_keys,
[self._data[key] for key in self.keys.hits_keys])
@cached_property
def tracks(self):
"""wrapper for offline tracks.
Returns
-------
Class
OfflineTracks.
"""
return OfflineTracks(
self.keys.cut_tracks_keys,
[self._data[key] for key in self.keys.tracks_keys])
@cached_property
def mc_hits(self):
"""wrapper for offline mc hits.
Returns
-------
Class
OfflineHits.
"""
return OfflineHits(self.keys.cut_hits_keys,
[self._data[key] for key in self.keys.mc_hits_keys])
@cached_property
def mc_tracks(self):
"""wrapper for offline mc tracks.
Returns
-------
Class
OfflineTracks.
"""
return OfflineTracks(
self.keys.cut_tracks_keys,
[self._data[key] for key in self.keys.mc_tracks_keys])
@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
returned, it means that the reconstruction parameter of interest is not
found. for example, in the case of muon simulations: if [1, 2] are the
reconstruction stages, then only the fit parameters corresponding to the
stages [1, 2] are found in the Offline files, the remaining fit parameters
corresponding to the stages 3, 4, 5 are all filled with nan.
Returns
-------
numpy recarray
a recarray of the best track fit data (reconstruction data).
"""
keys = ", ".join(self.keys.fit_keys[:-1])
empty_fit_info = np.array(
[match for match in self._find_empty(self.tracks.fitinf)])
fit_info = [
i for i, j in zip(self.tracks.fitinf, empty_fit_info[:, 1])
if j is not None
]
stages = self._get_max_reco_stages(self.tracks.rec_stages)
fit_data = np.array([i[j] for i, j in zip(fit_info, stages[:, 2])])
rows_size = len(max(fit_data, key=len))
equal_size_data = np.vstack([
np.hstack([i, np.zeros(rows_size - len(i)) + np.nan])
for i in fit_data
])
return np.core.records.fromarrays(equal_size_data.transpose(),
names=keys)
def _get_max_reco_stages(self, reco_stages):
"""find the longest reconstructed track based on the maximum size of
reconstructed stages.
Parameters
----------
reco_stages : chunked array
chunked array of all the reconstruction stages of all tracks.
In km3io, it is accessed with
km3io.OfflineReader(my_file).tracks.rec_stages .
Returns
-------
numpy array
array with 3 columns: *list of the maximum reco_stages
*lentgh of the maximum reco_stages
*position of the maximum reco_stages
"""
empty_reco_stages = np.array(
[match for match in self._find_empty(reco_stages)])
max_reco_stages = np.array(
[[max(i, key=len),
len(max(i, key=len)),
i.index(max(i, key=len))]
for i, j in zip(reco_stages, empty_reco_stages[:, 1])
if j is not None])
return max_reco_stages
def get_reco_fit(self, stages, mc=False):
"""construct a numpy recarray of the fit information (reconstruction
data) of the tracks reconstructed following the reconstruction stages
of interest.
Parameters
----------
stages : list
list of reconstruction stages of interest. for example
[1, 2, 3, 4, 5].
mc : bool, optional
default is False to look for fit data in the tracks tree in offline files
(not the mc tracks tree). mc=True to look for fit data from the mc tracks
tree in offline files.
Returns
-------
numpy recarray
a recarray of the fit information (reconstruction data) of
the tracks of interest.
Raises
------
ValueError
ValueError raised when the reconstruction stages of interest
are not found in the file.
"""
keys = ", ".join(self.keys.fit_keys[:-1])
if mc is False:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=False)])
fitinf = self.tracks.fitinf
if mc is True:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=True)])
fitinf = self.mc_tracks.fitinf
mask = rec_stages[:, 1] != None
if np.all(rec_stages[:, 1] == None):
raise ValueError(
"The stages {} are not found in your file.".format(
str(stages)))
else:
fit_data = np.array(
[i[k] for i, k in zip(fitinf[mask], rec_stages[:, 1][mask])])
rec_array = np.core.records.fromarrays(fit_data.transpose(),
names=keys)
return rec_array
def get_reco_hits(self, stages, keys, mc=False):
"""construct a dictionary of hits class data based on the reconstruction
stages of interest. For example, if the reconstruction stages of interest
are [1, 2, 3, 4, 5], then get_reco_hits method will select the hits data
from the events that were reconstructed following these stages (i.e
[1, 2, 3, 4, 5]).
Parameters
----------
stages : list
list of reconstruction stages of interest. for example
[1, 2, 3, 4, 5].
keys : list of str
list of the hits class attributes.
mc : bool, optional
default is False to look for hits data in the hits tree in offline files
(not the mc_hits tree). mc=True to look for mc hits data in the mc hits
tree in offline files.
Returns
-------
dict
dictionary of lazyarrays containing data for each hits attribute requested.
Raises
------
ValueError
ValueError raised when the reconstruction stages of interest
are not found in the file.
"""
lazy_d = {}
if mc is False:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=False)])
hits_data = self.hits
if mc is True:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=True)])
hits_data = self.mc_hits
mask = rec_stages[:, 1] != None
if np.all(rec_stages[:, 1] == None):
raise ValueError(
"The stages {} are not found in your file.".format(
str(stages)))
else:
for key in keys:
lazy_d[key] = getattr(hits_data, key)[mask]
return lazy_d
def get_reco_events(self, stages, keys, mc=False):
"""construct a dictionary of events class data based on the reconstruction
stages of interest. For example, if the reconstruction stages of interest
are [1, 2, 3, 4, 5], then get_reco_events method will select the events data
that were reconstructed following these stages (i.e [1, 2, 3, 4, 5]).
Parameters
----------
stages : list
list of reconstruction stages of interest. for example
[1, 2, 3, 4, 5].
keys : list of str
list of the events class attributes.
mc : bool, optional
default is False to look for the reconstruction stages in the tracks tree
in offline files (not the mc tracks tree). mc=True to look for the reconstruction
data in the mc tracks tree in offline files.
Returns
-------
dict
dictionary of lazyarrays containing data for each events attribute requested.
Raises
------
ValueError
ValueError raised when the reconstruction stages of interest
are not found in the file.
"""
lazy_d = {}
if mc is False:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=False)])
if mc is True:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=True)])
mask = rec_stages[:, 1] != None
if np.all(rec_stages[:, 1] == None):
raise ValueError(
"The stages {} are not found in your file.".format(
str(stages)))
else:
for key in keys:
lazy_d[key] = getattr(self.events, key)[mask]
return lazy_d
def get_reco_tracks(self, stages, keys, mc=False):
"""construct a dictionary of tracks class data based on the reconstruction
stages of interest. For example, if the reconstruction stages of interest
are [1, 2, 3, 4, 5], then get_reco_tracks method will select tracks data
from the events that were reconstructed following these stages (i.e
[1, 2, 3, 4, 5]).
Parameters
----------
stages : list
list of reconstruction stages of interest. for example
[1, 2, 3, 4, 5].
keys : list of str
list of the tracks class attributes.
mc : bool, optional
default is False to look for tracks data in the tracks tree in offline files
(not the mc tracks tree). mc=True to look for tracks data in the mc tracks
tree in offline files.
Returns
-------
dict
dictionary of lazyarrays containing data for each tracks attribute requested.
Raises
------
ValueError
ValueError raised when the reconstruction stages of interest
are not found in the file.
"""
lazy_d = {}
if mc is False:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=False)])
tracks_data = self.tracks
if mc is True:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=True)])
tracks_data = self.mc_tracks
mask = rec_stages[:, 1] != None
if np.all(rec_stages[:, 1] == None):
raise ValueError(
"The stages {} are not found in your file.".format(
str(stages)))
else:
for key in keys:
lazy_d[key] = np.array([
i[k] for i, k in zip(
getattr(tracks_data, key)[mask], rec_stages[:,
1][mask])
])
return lazy_d
def _find_rec_stages(self, stages, mc=False):
"""find the index of reconstruction stages of interest in a
list of multiple reconstruction stages.
Parameters
----------
stages : list
list of reconstruction stages of interest. for example
[1, 2, 3, 4, 5].
mc : bool, optional
default is False to look for reconstruction stages in the tracks tree in
offline files (not the mc tracks tree). mc=True to look for reconstruction
stages in the mc tracks tree in offline files.
Yields
------
generator
the track id and the index of the reconstruction stages of
interest if found. If the reconstruction stages of interest
are not found, None is returned as the stages index.
"""
if mc is False:
stages_data = self.tracks.rec_stages
if mc is True:
stages_data = self.mc_tracks.rec_stages
for trk_index, rec_stages in enumerate(stages_data):
try:
stages_index = rec_stages.index(stages)
except ValueError:
stages_index = None
yield trk_index, stages_index
continue
yield trk_index, stages_index
def _find_empty(self, array):
"""finds empty lists/arrays in an awkward array
Parameters
----------
array : awkward array
Awkward array of data of interest. For example:
km3io.OfflineReader(my_file).tracks.fitinf .
Yields
------
generator
the empty list id and the index of the empty list. When
data structure (list) is simply empty, None is written in the
corresponding index. However, when data structure (list) is not
empty and does not contain an empty list, then False is written in the
corresponding index.
"""
for i, rs in enumerate(array):
try:
if len(rs) == 0:
j = None
if len(rs) != 0:
j = rs.index([])
except ValueError:
j = False # rs not empty but [] not found
yield i, j
continue
yield i, j
class Usr:
"""Helper class to access AAObject usr stuff"""
def __init__(self, tree):
"""Helper class to access AAObject `usr`` stuff"""
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.
# to massively increase the performance. This needs triple check if
# it's always the case; the usr-format is simply a very bad design.
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(
basketcache=BASKET_CACHE)[0]
]
except (KeyError, IndexError): # e.g. old aanet files
print("The `usr` fields could not be parsed for the '{}' branch."
.format(name))
self._usr_names = []
else:
self._usr_idx_lookup = {
name: index
for index, name in enumerate(self._usr_names)
}
self._usr_data = tree['Evt']['usr'].lazyarray(
basketcache=uproot.cache.ThreadSafeArrayCache(
BASKET_CACHE_SIZE))
data = tree['usr'].lazyarray(basketcache=BASKET_CACHE)
if index is not None:
data = data[index]
self._usr_data = data
for name in self._usr_names:
setattr(self, name, self[name])
@@ -707,187 +164,154 @@ class Usr:
entries.append("{}: {}".format(name, self[name]))
return '\n'.join(entries)
class OfflineEvents:
"""wrapper for offline events"""
def __init__(self, keys, values):
"""wrapper for offline events.
Parameters
----------
keys : list of str
list of valid events keys.
values : list of arrays
list of arrays containting events data.
"""
self._keys = keys
self._values = values
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __getitem__(self, item):
return OfflineEvent(self._keys, [v[item] for v in self._values])
def __len__(self):
try:
return len(self._values[0])
except IndexError:
return 0
def __str__(self):
return "Number of events: {}".format(len(self))
def __repr__(self):
return "<{}: {} parsed events>".format(self.__class__.__name__,
len(self))
class OfflineEvent:
"""wrapper for an offline event"""
def __init__(self, keys, values):
"""wrapper for one offline event.
Parameters
----------
keys : list of str
list of valid events keys.
values : list of arrays
list of arrays containting event data.
"""
self._keys = keys
self._values = values
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __str__(self):
return "offline event:\n\t" + "\n\t".join([
"{:15} {:^10} {:>10}".format(k, ':', str(v))
for k, v in zip(self._keys, self._values)
])
return "<{}[{}]>".format(self.__class__.__name__, self._name)
class OfflineHits:
"""wrapper for offline hits"""
def __init__(self, keys, values):
"""wrapper for offline hits.
Parameters
----------
keys : list of str
list of cropped hits keys.
values : list of arrays
list of arrays containting hits data.
"""
self._keys = keys
self._values = values
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __getitem__(self, item):
return OfflineHit(self._keys, [v[item] for v in self._values])
def __len__(self):
def _to_num(value):
"""Convert a value to a numerical one if possible"""
if value is None:
return
try:
return int(value)
except ValueError:
try:
return len(self._values[0])
except IndexError:
return 0
def __str__(self):
return "Number of hits: {}".format(len(self))
def __repr__(self):
return "<{}: {} parsed elements>".format(self.__class__.__name__,
len(self))
return float(value)
except ValueError:
pass
else:
return value
class Header:
"""The header"""
def __init__(self, header):
self._data = {}
for attribute, fields in mc_header.items():
values = header.get(attribute, '').split()
if not values:
continue
Constructor = namedtuple(attribute, fields)
if len(values) < len(fields):
values += [None] * (len(fields) - len(values))
self._data[attribute] = Constructor(
**{f: _to_num(v)
for (f, v) in zip(fields, values)})
for attribute, value in self._data.items():
setattr(self, attribute, value)
class OfflineHit:
"""wrapper for an offline hit"""
def __init__(self, keys, values):
"""wrapper for one offline hit.
def __str__(self):
lines = ["MC Header:"]
for value in self._data.values():
lines.append(" {}".format(value))
return "\n".join(lines)
class Branch:
"""Branch accessor class"""
def __init__(self,
tree,
mapper,
index=None,
subbranchmaps=None,
keymap=None):
self._tree = tree
self._mapper = mapper
self._index = index
self._keymap = None
self._branch = tree[mapper.key]
self._subbranches = []
if keymap is None:
self._initialise_keys() #
else:
self._keymap = keymap
if subbranchmaps is not None:
for mapper in subbranchmaps:
subbranch = Branch(self._tree,
mapper=mapper,
index=self._index)
self._subbranches.append(subbranch)
for subbranch in self._subbranches:
setattr(self, subbranch._mapper.name, subbranch)
def _initialise_keys(self):
"""Create the keymap and instance attributes for branch keys"""
# TODO: this could be a cached property
keys = set(k.decode('utf-8') for k in self._branch.keys()) - set(
self._mapper.exclude) - EXCLUDE_KEYS
self._keymap = {
**{self._mapper.attrparser(k): k
for k in keys},
**self._mapper.extra
}
self._keymap.update(self._mapper.update)
for k in self._mapper.update.values():
del self._keymap[k]
for key in self._keymap.keys():
setattr(self, key, None)
Parameters
----------
keys : list of str
list of cropped hits keys.
values : list of arrays
list of arrays containting hit data.
"""
self._keys = keys
self._values = values
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def keys(self):
return self._keymap.keys()
def __str__(self):
return "offline hit:\n\t" + "\n\t".join([
"{:15} {:^10} {:>10}".format(k, ':', str(v))
for k, v in zip(self._keys, self._values)
])
@cached_property
def usr(self):
return Usr(self._mapper.name, self._branch, index=self._index)
def __getitem__(self, item):
return self._values[item]
def __getattribute__(self, attr):
if attr.startswith("_"): # let all private and magic methods pass
return object.__getattribute__(self, attr)
if attr in self._keymap.keys(): # intercept branch key lookups
return self.__getkey__(attr)
class OfflineTracks:
"""wrapper for offline tracks"""
def __init__(self, keys, values):
"""wrapper for offline tracks
return object.__getattribute__(self, attr)
Parameters
----------
keys : list of str
list of cropped tracks keys.
values : list of arrays
list of arrays containting tracks data.
"""
self._keys = keys
self._values = values
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __getkey__(self, key):
out = self._branch[self._keymap[key]].lazyarray(
basketcache=BASKET_CACHE)
if self._index is not None:
out = out[self._index]
return out
def __getitem__(self, item):
return OfflineTrack(self._keys, [v[item] for v in self._values])
"""Slicing magic"""
if isinstance(item, (int, slice)):
return self.__class__(self._tree,
self._mapper,
index=item,
keymap=self._keymap,
subbranchmaps=SUBBRANCH_MAPS)
if isinstance(item, tuple):
return self[item[0]][item[1]]
if isinstance(item, str):
return self.__getkey__(item)
return self.__class__(self._tree,
self._mapper,
index=np.array(item),
keymap=self._keymap,
subbranchmaps=SUBBRANCH_MAPS)
def __len__(self):
try:
return len(self._values[0])
except IndexError:
return 0
if self._index is None:
return len(self._branch)
elif isinstance(self._index, int):
return 1
else:
return len(self._branch[self._keymap['id']].lazyarray(
basketcache=BASKET_CACHE)[self._index])
def __str__(self):
return "Number of tracks: {}".format(len(self))
return "Number of elements: {}".format(len(self._branch))
def __repr__(self):
return "<{}: {} parsed elements>".format(self.__class__.__name__,
len(self))
class OfflineTrack:
"""wrapper for an offline track"""
def __init__(self, keys, values):
"""wrapper for one offline track.
Parameters
----------
keys : list of str
list of cropped tracks keys.
values : list of arrays
list of arrays containting track data.
"""
self._keys = keys
self._values = values
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __str__(self):
return "offline track:\n\t" + "\n\t".join([
"{:30} {:^2} {:>26}".format(k, ':', str(v))
for k, v in zip(self._keys, self._values) if k not in ['fitinf']
]) + "\n\t" + "\n\t".join([
"{:30} {:^2} {:>26}".format(k, ':', str(
getattr(self, 'fitinf')[v]))
for k, v in km3io.definitions.fitparameters.data.items()
if len(getattr(self, 'fitinf')) > v
]) # I don't like 18 being explicit here
def __getitem__(self, item):
return self._values[item]
length = len(self)
return "<{}[{}]: {} element{}>".format(self.__class__.__name__,
self._mapper.name, length,
's' if length > 1 else '')