Skip to content
Snippets Groups Projects

Refactor offline

Closed Refactor offline
refactor-offline into master
Closed Tamas Gal requested to merge refactor-offline into master
Viewing commit 71c1fb39
Prev Next
Show latest version
2 files
+ 41
− 27
Compare changes
  • Side-by-side
  • Inline
Files
2
km3io/offline.py
+ 199
− 834
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
from .tools import Branch, BranchMapper, cached_property, _to_num, _unfold_indices
MAIN_TREE_NAME = "E"
EXCLUDE_KEYS = ["AAObject", "t", "fBits", "fUniqueID"]
# 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2
class cached_property:
"""A simple cache decorator for properties."""
def __init__(self, function):
self.function = function
def __get__(self, obj, cls):
if obj is None:
return self
prop = obj.__dict__[self.function.__name__] = self.function(obj)
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
BASKET_CACHE = uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)
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=EXCLUDE_KEYS,
update={
'n_hits': 'hits',
'n_mc_hits': 'mc_hits',
'n_tracks': 'trks',
'n_mc_tracks': 'mc_trks'
},
attrparser=lambda a: a,
flat=True)
SUBBRANCH_MAPS = [
BranchMapper(name="tracks",
key="trks",
extra={},
exclude=EXCLUDE_KEYS +
['trks.usr_data', 'trks.usr', 'trks.fUniqueID', 'trks.fBits'],
update={},
attrparser=_nested_mapper,
flat=False),
BranchMapper(name="mc_tracks",
key="mc_trks",
extra={},
exclude=EXCLUDE_KEYS + [
'mc_trks.usr_data', 'mc_trks.usr', 'mc_trks.rec_stages',
'mc_trks.fitinf', 'mc_trks.fUniqueID', 'mc_trks.fBits'
],
update={},
attrparser=_nested_mapper,
flat=False),
BranchMapper(name="hits",
key="hits",
extra={},
exclude=EXCLUDE_KEYS + [
'hits.usr', 'hits.pmt_id', 'hits.origin', 'hits.a',
'hits.pure_a', 'hits.fUniqueID', 'hits.fBits'
],
update={},
attrparser=_nested_mapper,
flat=False),
BranchMapper(name="mc_hits",
key="mc_hits",
extra={},
exclude=EXCLUDE_KEYS + [
'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,
flat=False),
]
class OfflineBranch(Branch):
@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
def usr(self):
return Usr(self._mapper, self._branch, index_chain=self._index_chain)
@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)
class Usr:
"""Helper class to access AAObject `usr` stuff"""
def __init__(self, mapper, branch, index_chain=None):
self._mapper = mapper
self._name = mapper.name
self._index_chain = [] if index_chain is None else index_chain
self._branch = branch
self._usr_names = []
self._usr_idx_lookup = {}
@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)
self._usr_key = 'usr' if mapper.flat else mapper.key + '.usr'
@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'
self._initialise()
def _initialise(self):
try:
self._branch[self._usr_key]
# This will raise a KeyError in old aanet files
# which has a different strucuter and key (usr_data)
# We do not support those (yet)
except (KeyError, IndexError):
print("The `usr` fields could not be parsed for the '{}' branch.".
format(self._name))
return
if self._mapper.flat:
self._initialise_flat()
def _initialise_flat(self):
# 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.
self._usr_names = [
n.decode("utf-8")
for n in self._branch[self._usr_key + '_names'].lazyarray()[0]
]
return self._get_keys(self._tree['Evt']['mc_trks'], fake_branches)
self._usr_idx_lookup = {
name: index
for index, name in enumerate(self._usr_names)
}
@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)
data = self._branch[self._usr_key].lazyarray()
@cached_property
def fit_keys(self):
"""constructs a list of fit parameters, not yet outsourced in an offline file.
if self._index_chain:
data = _unfold_indices(data, self._index_chain)
Returns
-------
list of str
list of all "trks.fitinf" keys.
"""
return sorted(km3io.definitions.fitparameters.data,
key=km3io.definitions.fitparameters.data.get,
reverse=False)
self._usr_data = data
@cached_property
def cut_hits_keys(self):
"""adapts hits keys for instance variables format in a Python class.
for name in self._usr_names:
setattr(self, name, self[name])
Returns
-------
list of str
list of adapted hits keys.
"""
return [k.split('hits.')[1].replace('.', '_') for k in self.hits_keys]
# def _initialise_nested(self):
# self._usr_names = [
# n.decode("utf-8") for n in self.branch['usr_names'].lazyarray(
# # TODO this will be fixed soon in uproot,
# # see https://github.com/scikit-hep/uproot/issues/465
# uproot.asgenobj(
# uproot.SimpleArray(uproot.STLVector(uproot.STLString())),
# self.branch['usr_names']._context, 6),
# basketcache=BASKET_CACHE)[0]
# ]
@cached_property
def cut_tracks_keys(self):
"""adapts tracks keys for instance variables format in a Python class.
def __getitem__(self, item):
if self._mapper.flat:
return self.__getitem_flat__(item)
return self.__getitem_nested__(item)
def __getitem_flat__(self, item):
if self._index_chain:
return _unfold_indices(
self._usr_data, self._index_chain)[:,
self._usr_idx_lookup[item]]
else:
return self._usr_data[:, self._usr_idx_lookup[item]]
def __getitem_nested__(self, item):
data = self._branch[self._usr_key + '_names'].lazyarray(
# TODO this will be fixed soon in uproot,
# see https://github.com/scikit-hep/uproot/issues/465
uproot.asgenobj(
uproot.SimpleArray(uproot.STLVector(uproot.STLString())),
self._branch[self._usr_key + '_names']._context, 6),
basketcache=BASKET_CACHE)
return _unfold_indices(data, self._index_chain)
Returns
-------
list of str
list of adapted tracks keys.
"""
return [
k.split('trks.')[1].replace('.', '_') for k in self.tracks_keys
]
def keys(self):
return self._usr_names
@cached_property
def cut_events_keys(self):
"""adapts events keys for instance variables format in a Python class.
def __str__(self):
entries = []
for name in self.keys():
entries.append("{}: {}".format(name, self[name]))
return '\n'.join(entries)
Returns
-------
list of str
list of adapted events keys.
"""
return [k.replace('.', '_') for k in self.events_keys]
def __repr__(self):
return "<{}[{}]>".format(self.__class__.__name__, self._name)
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,681 +192,67 @@ 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
self._fobj = uproot.open(file_path)
self._tree = self._fobj[MAIN_TREE_NAME]
@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)
def __len__(self):
return len(self._data)
@cached_property
def events(self):
"""The `E` branch, containing all offline events."""
return OfflineBranch(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]).
class Header:
"""The header"""
def __init__(self, header):
self._data = {}
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 = {}
for attribute, fields in header.items():
values = fields.split()
fields = mc_header.get(attribute, [])
if mc is False:
rec_stages = np.array(
[match for match in self._find_rec_stages(stages, mc=False)])
hits_data = self.hits
n_values = len(values)
n_fields = len(fields)
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
if n_values == 1 and n_fields == 0:
self._data[attribute] = _to_num(values[0])
continue
yield trk_index, stages_index
n_max = max(n_values, n_fields)
values += [None] * (n_max - n_values)
fields += ["field_{}".format(i) for i in range(n_fields, n_max)]
def _find_empty(self, array):
"""finds empty lists/arrays in an awkward array
Constructor = namedtuple(attribute, fields)
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
if not values:
continue
yield i, j
class Usr:
"""Helper class to access AAObject usr stuff"""
def __init__(self, tree):
# 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.
try:
self._usr_names = [
n.decode("utf-8") for n in tree['Evt']['usr_names'].array()[0]
]
except (KeyError, IndexError): # e.g. old aanet files
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))
for name in self._usr_names:
setattr(self, name, self[name])
def __getitem__(self, item):
return self._usr_data[:, self._usr_idx_lookup[item]]
def keys(self):
return self._usr_names
def __str__(self):
entries = []
for name in self.keys():
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)
])
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):
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))
class OfflineHit:
"""wrapper for an offline hit"""
def __init__(self, keys, values):
"""wrapper for one offline hit.
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 __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)
])
def __getitem__(self, item):
return self._values[item]
class OfflineTracks:
"""wrapper for offline tracks"""
def __init__(self, keys, values):
"""wrapper for offline tracks
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 __getitem__(self, item):
return OfflineTrack(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 tracks: {}".format(len(self))
def __repr__(self):
return "<{}: {} parsed elements>".format(self.__class__.__name__,
len(self))
self._data[attribute] = Constructor(
**{f: _to_num(v)
for (f, v) in zip(fields, values)})
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)
for attribute, value in self._data.items():
setattr(self, attribute, value)
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]
lines = ["MC Header:"]
keys = set(mc_header.keys())
for key, value in self._data.items():
if key in keys:
lines.append(" {}".format(value))
else:
lines.append(" {}: {}".format(key, value))
return "\n".join(lines)