Skip to content
Snippets Groups Projects

Api update

Merged Api update
api-update into master
Merged Zineb Aly requested to merge api-update into master
Compare
and
Show latest version
1 file
+ 250
− 62
Compare changes
  • Side-by-side
  • Inline
km3io/aanet.py
+ 250
− 62
import uproot
BASKET_CACHE_SIZE = 23 * 1024**2 # [byte] DON T FORGET TO INCLUDE THIS FOR LARGE DATA SET!!!
# 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2
class AanetKeys:
"wrapper for aanet keys"
"""wrapper for aanet keys"""
def __init__(self, file_path):
"""AanetKeys is a class that reads all the available keys in an aanet
file and adapts the keys format to Python format.
Parameters
----------
file_path : path-like object
Path to the aanet file of interest. It can be a str or any python
path-like object that points to the file of ineterst.
"""
self._file_path = file_path
self._events_keys = None
self._hits_keys = None
@@ -17,6 +26,7 @@ class AanetKeys:
self._fit_keys = None
self._cut_hits_keys = None
self._cut_tracks_keys = None
self._cut_events_keys = None
def __str__(self):
return '\n'.join([
@@ -32,6 +42,14 @@ class AanetKeys:
@property
def events_keys(self):
"""reads events keys from an aanet file.
Returns
-------
list of str
list of all events keys found in an aanet file,
except those found in fake branches.
"""
if self._events_keys is None:
fake_branches = ['Evt', 'AAObject', 'TObject', 't']
t_baskets = ['t.fSec', 't.fNanoSec']
@@ -44,17 +62,40 @@ class AanetKeys:
@property
def hits_keys(self):
"""reads hits keys from an aanet file.
Returns
-------
list of str
list of all hits keys found in an aanet file,
except those found in fake branches.
"""
if self._hits_keys is None:
fake_branches = [
'hits.usr'
] # to be treated like trks.usr and trks.usr_names
tree = uproot.open(self._file_path)['E']['hits']
self._hits_keys = [key.decode('utf8') for key in tree.keys()]
self._hits_keys = [
key.decode('utf8') for key in tree.keys()
if key.decode('utf8') not in fake_branches
]
return self._hits_keys
@property
def tracks_keys(self):
"""reads tracks keys from an aanet file.
Returns
-------
list of str
list of all tracks keys found in an aanet file,
except those found in fake branches.
"""
if self._tracks_keys is None:
# a solution can be tree['trks.usr_data'].array(
# uproot.asdtype(">i4"))
fake_branches = [
'trks.usr_data', # a solution can be tree['trks.usr_data'].array(uproot.asdtype(">i4"))
'trks.usr_names'
'trks.usr_data', 'trks.usr'
] # can be accessed using tree['trks.usr_names'].array()
tree = uproot.open(self._file_path)['E']['Evt']['trks']
self._tracks_keys = [
@@ -65,15 +106,35 @@ class AanetKeys:
@property
def mc_hits_keys(self):
"""reads mc hits keys from an aanet file.
Returns
-------
list of str
list of all mc hits keys found in an aanet file,
except those found in fake branches.
"""
if self._mc_hits_keys is None:
fake_branches = ['mc_hits.usr']
tree = uproot.open(self._file_path)['E']['Evt']['mc_hits']
self._mc_hits_keys = [key.decode('utf8') for key in tree.keys()]
self._mc_hits_keys = [
key.decode('utf8') for key in tree.keys()
if key.decode('utf8') not in fake_branches
]
return self._mc_hits_keys
@property
def mc_tracks_keys(self):
"""reads mc tracks keys from an aanet file.
Returns
-------
list of str
list of all mc tracks keys found in an aanet file,
except those found in fake branches.
"""
if self._mc_tracks_keys is None:
fake_branches = ['mc_trks.usr_data', 'mc_trks.usr_names'
fake_branches = ['mc_trks.usr_data', 'mc_trks.usr'
] # same solution as above can be used
tree = uproot.open(self._file_path)['E']['Evt']['mc_trks']
self._mc_tracks_keys = [
@@ -87,7 +148,7 @@ class AanetKeys:
"""constructs a list of all valid keys to be read from an Aanet event file.
Returns
-------
list
list of str
list of all valid keys.
"""
if self._valid_keys is None:
@@ -98,7 +159,15 @@ class AanetKeys:
@property
def fit_keys(self):
"""constructs a list of fit parameters, not yet outsourced in an aanet file.
Returns
-------
list of str
list of all "trks.fitinf" keys.
"""
if self._fit_keys is None:
# these are hardcoded because they are not outsourced in aanet file
self._fit_keys = [
'JGANDALF_BETA0_RAD', 'JGANDALF_BETA1_RAD', 'JGANDALF_CHI2',
'JGANDALF_NUMBER_OF_HITS', 'JENERGY_ENERGY', 'JENERGY_CHI2',
@@ -112,6 +181,13 @@ class AanetKeys:
@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.
"""
if self._cut_hits_keys is None:
self._cut_hits_keys = [
k.split('hits.')[1].replace('.', '_') for k in self.hits_keys
@@ -120,17 +196,40 @@ class AanetKeys:
@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.
"""
if self._cut_tracks_keys is None:
self._cut_tracks_keys = [
k.split('trks.')[1].replace('.', '_') for k in self.tracks_keys
]
return self._cut_tracks_keys
@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.
"""
if self._cut_events_keys is None:
self._cut_events_keys = [
k.replace('.', '_') for k in self.events_keys
]
return self._cut_events_keys
class Reader:
"""Reader for one Aanet ROOT file"""
def __init__(self, file_path):
""" AanetReader class is a Aanet ROOT file wrapper
""" AanetReader class is a Aanet ROOT file reader. This class is a
"very" low level I/O.
Parameters
----------
@@ -139,11 +238,12 @@ class Reader:
path-like object that points to the file of ineterst.
"""
self._file_path = file_path
self._data = uproot.open(self._file_path)['E'].lazyarrays()
self._data = uproot.open(self._file_path)['E'].lazyarrays(
basketcache=uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE))
self._keys = None
def __getitem__(self, key):
"""reads data stored in the branch of interest in an event tree.
"""reads data stored in the branch of interest in an Evt tree.
Parameters
----------
@@ -174,14 +274,22 @@ class Reader:
@property
def keys(self):
"""wrapper for all keys in an aanet file.
Returns
-------
Class
AanetKeys.
"""
if self._keys is None:
self._keys = AanetKeys(self._file_path)
return self._keys
class AanetReader:
"""reader for Aanet ROOT files"""
def __init__(self, file_path, data=None):
""" AanetReader class is a Aanet ROOT file wrapper
""" AanetReader class is an aanet ROOT file wrapper
Parameters
----------
@@ -193,7 +301,9 @@ class AanetReader:
if data is not None:
self._data = data
else:
self._data = uproot.open(self._file_path)['E'].lazyarrays()
self._data = uproot.open(self._file_path)['E'].lazyarrays(
basketcache=uproot.cache.ThreadSafeArrayCache(
BASKET_CACHE_SIZE))
self._events = None
self._hits = None
self._tracks = None
@@ -201,28 +311,46 @@ class AanetReader:
self._mc_tracks = None
self._keys = None
# def __getitem__(self, item):
# return AanetEvents(self._events_keys, [self._data[key] for key in self.events])
def __getitem__(self, item):
return AanetReader(file_path=self._file_path, data=self._data[item])
@property
def keys(self):
"""wrapper for all keys in an aanet file.
Returns
-------
Class
AanetKeys.
"""
if self._keys is None:
self._keys = AanetKeys(self._file_path)
return self._keys
@property
def events(self):
"""wrapper for aanet events.
Returns
-------
Class
AanetEvents.
"""
if self._events is None:
self._events = AanetEvents(
self.keys.events_keys,
self.keys.cut_events_keys,
[self._data[key] for key in self.keys.events_keys])
return self._events
@property
def hits(self):
"""wrapper for aanet hits.
Returns
-------
Class
AanetHits.
"""
if self._hits is None:
self._hits = AanetHits(
self.keys.cut_hits_keys,
@@ -231,6 +359,13 @@ class AanetReader:
@property
def tracks(self):
"""wrapper for aanet tracks.
Returns
-------
Class
AanetTracks.
"""
if self._tracks is None:
self._tracks = AanetTracks(
self.keys.cut_tracks_keys,
@@ -240,6 +375,13 @@ class AanetReader:
@property
def mc_hits(self):
"""wrapper for aanet mc hits.
Returns
-------
Class
AanetHits.
"""
if self._mc_hits is None:
self._mc_hits = AanetHits(
self.keys.cut_hits_keys,
@@ -248,6 +390,13 @@ class AanetReader:
@property
def mc_tracks(self):
"""wrapper for aanet mc tracks.
Returns
-------
Class
AanetTracks.
"""
if self._mc_tracks is None:
self._mc_tracks = AanetTracks(
self.keys.cut_tracks_keys,
@@ -256,10 +405,18 @@ class AanetReader:
class AanetEvents:
"wrapper for Aanet events"
"""wrapper for Aanet events"""
def __init__(self, keys, values):
"""wrapper for aanet events.
def __init__(self, keys, values): # values is a list of lists
self._keys = keys # list of keys
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)
@@ -268,25 +425,32 @@ class AanetEvents:
return AanetEvent(self._keys, [v[item] for v in self._values])
def __len__(self):
return len(
self._values[0]
) # I don't like this being explicit, what if values is empty ...
try:
return len(self._values[0])
except IndexError:
return 0
def __str__(self):
return "Number of events: {}".format(len(self))
# def __repr__(self):
# return f'{self.__class__.__name__}({self._keys}, {self._values})'
def __repr__(self):
return str(self)
return "<{}: {} parsed events>".format(self.__class__.__name__,
len(self))
class AanetEvent:
"wrapper for an Aanet event"
"""wrapper for an Aanet event"""
def __init__(self, keys, values):
"""wrapper for one aanet event.
def __init__(self, keys, values): # both inputs are lists
self._keys = keys # list of keys
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)
@@ -297,18 +461,23 @@ class AanetEvent:
for k, v in zip(self._keys, self._values)
])
# def __repr__(self):
# return f'{self.__class__.__name__}({self._keys}, {self._values})'
def __repr__(self):
return str(self)
class AanetHits:
"wrapper for Aanet hits, manages the display of all hits in one event"
"""wrapper for Aanet hits, manages the display of all hits in one event"""
def __init__(self, keys, values):
"""wrapper for aanet hits.
def __init__(self, keys, values): # values is a list of lists
self._keys = keys # list of keys
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)
@@ -321,26 +490,31 @@ class AanetHits:
try:
return len(
self._values[0]
) # this is missleading and it sometimes prints the number of events ...
) # this is missleading and it sometimes prints the # of events
except IndexError:
return 0
def __str__(self):
return "Number of hits: {}".format(len(self))
# def __repr__(self):
# return f'{self.__class__.__name__}({self._keys}, {self._values})'
def __repr__(self):
return "<{}: {} parsed elements>".format(self.__class__.__name__,
len(self))
class AanetHit:
"wrapper for an Aanet hit"
"""wrapper for an Aanet hit"""
def __init__(self, keys, values):
"""wrapper for one aanet hit.
def __init__(self, keys, values): # both inputs are lists
self._keys = keys # list of keys
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)
@@ -352,22 +526,33 @@ class AanetHit:
])
def __getitem__(self, item):
# return self._values[item]
return self._values[item]
# def __repr__(self):
# return f'{self.__class__.__name__}({self._keys}, {self._values})'
def __repr__(self):
return str(self)
# def _is_empty(array):
# if array.size:
# return False
# else:
# return True
class AanetTracks:
"wrapper for Aanet tracks, manages the display of all tracks in one event"
"""wrapper for Aanet tracks"""
def __init__(self, keys, values, fit_keys=None):
"""Summary
def __init__(self, keys, values,
fit_keys=None): # values is a list of lists
self._keys = keys # list of keys
Parameters
----------
keys : TYPE
Description
values : TYPE
Description
fit_keys : None, optional
list of tracks fit information (not yet outsourced in aanet files)
"""
self._keys = keys
self._values = values
if fit_keys is not None:
self._fit_keys = fit_keys
@@ -375,7 +560,6 @@ class AanetTracks:
setattr(self, k, v)
def __getitem__(self, item):
# return self._values[item]
return AanetTrack(self._keys, [v[item] for v in self._values],
fit_keys=self._fit_keys)
@@ -389,19 +573,26 @@ class AanetTracks:
return "Number of tracks: {}".format(
len(self)) # this is not correct when reader.tracks is called
# def __repr__(self):
# return f'{self.__class__.__name__}({self._keys}, {self._values})'
def __repr__(self):
return "<{}: {} parsed elements>".format(self.__class__.__name__,
len(self))
class AanetTrack:
"wrapper for an Aanet track"
"""wrapper for an Aanet track"""
def __init__(self, keys, values, fit_keys=None):
"""wrapper for one aanet track.
def __init__(self, keys, values, fit_keys=None): # both inputs are lists
self._keys = keys # list of keys
Parameters
----------
keys : list of str
list of cropped tracks keys.
values : list of arrays
list of arrays containting track data.
fit_keys : None, optional
list of tracks fit information (not yet outsourced in aanet files).
"""
self._keys = keys
self._values = values
if fit_keys is not None:
self._fit_keys = fit_keys
@@ -421,8 +612,5 @@ class AanetTrack:
def __getitem__(self, item):
return self._values[item]
# def __repr__(self):
# return f'{self.__class__.__name__}({self._keys}, {self._values})'
def __repr__(self):
return str(self)