Skip to content
Snippets Groups Projects
Commit 1ed7365e authored by Tamas Gal's avatar Tamas Gal :speech_balloon:
Browse files

Merge branch 'api-update' into 'master' 

Api update

See merge request !1
parents 336faa30 62edbd50
No related branches found
No related tags found
1 merge request!1Api update
Pipeline #7612 passed with warnings
Stage: test
Stage: coverage
Stage: doc
Stage: release
Stage: deploy
Hide whitespace changes
Inline Side-by-side
km3io/aanet.py
+ 562
− 57
View file @ 1ed7365e
import uproot
# 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2
class AanetReader:
"""Reader for one Aanet ROOT file"""
class AanetKeys:
"""wrapper for aanet keys"""
def __init__(self, file_path):
""" AanetReader class is a Aanet ROOT file wrapper
"""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 file of interest. It can be a str or any python
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.data = uproot.open(self.file_path)['E']
self.lazy_data = self.data.lazyarrays()
self._file_path = file_path
self._events_keys = None
self._hits_keys = None
self._tracks_keys = None
self._mc_hits_keys = None
self._mc_tracks_keys = None
self._valid_keys = None
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([
"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 str(self)
# return f'{self.__class__.__name__}("{self._file_path}")'
@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']
tree = uproot.open(self._file_path)['E']['Evt']
self._events_keys = [
key.decode('utf-8') for key in tree.keys()
if key.decode('utf-8') not in fake_branches
] + t_baskets
return self._events_keys
@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', 'hits.usr_names'
] # 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()
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', 'trks.usr', 'trks.usr_names'
] # can be accessed using tree['trks.usr_names'].array()
tree = uproot.open(self._file_path)['E']['Evt']['trks']
self._tracks_keys = [
key.decode('utf8') for key in tree.keys()
if key.decode('utf8') not in fake_branches
]
return self._tracks_keys
@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', 'mc_hits.usr_names']
tree = uproot.open(self._file_path)['E']['Evt']['mc_hits']
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', 'mc_trks.usr_names'
] # same solution as above can be used
tree = uproot.open(self._file_path)['E']['Evt']['mc_trks']
self._mc_tracks_keys = [
key.decode('utf8') for key in tree.keys()
if key.decode('utf8') not in fake_branches
]
return self._mc_tracks_keys
@property
def valid_keys(self):
"""constructs a list of all valid keys to be read from an Aanet event file.
Returns
-------
list of str
list of all valid keys.
"""
if self._valid_keys is None:
self._valid_keys = (self.events_keys + self.hits_keys +
self.tracks_keys + self.mc_tracks_keys +
self.mc_hits_keys)
return self._valid_keys
@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',
'JGANDALF_LAMBDA', 'JGANDALF_NUMBER_OF_ITERATIONS',
'JSTART_NPE_MIP', 'JSTART_NPE_MIP_TOTAL',
'JSTART_LENGTH_METRES', 'JVETO_NPE', 'JVETO_NUMBER_OF_HITS',
'JENERGY_MUON_RANGE_METRES', 'JENERGY_NOISE_LIKELIHOOD',
'JENERGY_NDF', 'JENERGY_NUMBER_OF_HITS', 'JCOPY_Z_M'
]
return self._fit_keys
@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
]
return self._cut_hits_keys
@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 reader. This class is a
"very" low level I/O.
Parameters
----------
file_path : path-like object
Path to the 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._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
----------
......@@ -41,78 +264,360 @@ class AanetReader:
Raises
------
KeyEroor
KeyError
Some branches in an Aanet file structure are "fake branches" and do
not contain data. Therefore, the keys corresponding to these fake
branches are not read.
"""
if key not in self.keys() and not isinstance(key, int):
if key not in self.keys.valid_keys and not isinstance(key, int):
raise KeyError(
"'{}' is not a valid key or is a fake branch.".format(key))
return self.lazy_data[key]
return self._data[key]
def __len__(self):
return len(self.lazy_data)
return len(self._data)
def __repr__(self):
return '\n'.join([
"Number of events: {}".format(self.__len__()),
"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)
])
return "<{}: {} entries>".format(self.__class__.__name__, len(self))
@property
def keys(self):
"""constructs a list of all valid keys to be read from an Aanet event file.
"""wrapper for all keys in an aanet file.
Returns
-------
list
list of all valid keys.
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 an aanet ROOT file wrapper
Parameters
----------
file_path : path-like object
Path to the file of interest. It can be a str or any python
path-like object that points to the file of ineterst.
"""
return self.events_keys + self.hits_keys + self.tracks_keys + self.mc_tracks_keys + self.mc_hits_keys
self._file_path = file_path
if data is not None:
self._data = data
else:
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
self._mc_hits = None
self._mc_tracks = None
self._keys = None
def __getitem__(self, item):
return AanetReader(file_path=self._file_path, data=self._data[item])
def __len__(self):
return len(self._data)
@property
def events_keys(self):
if self._events_keys is None:
fake_branches = ['Evt', 'AAObject', 'TObject', 't']
t_baskets = ['t.fSec', 't.fNanoSec']
self._events_keys = [
key.decode('utf-8') for key in self.data['Evt'].keys()
if key.decode('utf-8') not in fake_branches
] + t_baskets
return self._events_keys
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 hits_keys(self):
if self._hits_keys is None:
hits_tree = self.data['Evt']['hits']
self._hits_keys = [key.decode('utf8') for key in hits_tree.keys()]
return self._hits_keys
def events(self):
"""wrapper for aanet events.
Returns
-------
Class
AanetEvents.
"""
if self._events is None:
self._events = AanetEvents(
self.keys.cut_events_keys,
[self._data[key] for key in self.keys.events_keys])
return self._events
@property
def tracks_keys(self):
if self._tracks_keys is None:
tracks_tree = self.data['Evt']['trks']
self._tracks_keys = [
key.decode('utf8') for key in tracks_tree.keys()
]
return self._tracks_keys
def hits(self):
"""wrapper for aanet hits.
Returns
-------
Class
AanetHits.
"""
if self._hits is None:
self._hits = AanetHits(
self.keys.cut_hits_keys,
[self._data[key] for key in self.keys.hits_keys])
return self._hits
@property
def mc_hits_keys(self):
if self._mc_hits_keys is None:
mc_hits_tree = self.data['Evt']['mc_hits']
self._mc_hits_keys = [key.decode('utf8') for key in mc_hits_tree.keys()]
return self._mc_hits_keys
def tracks(self):
"""wrapper for aanet tracks.
Returns
-------
Class
AanetTracks.
"""
if self._tracks is None:
self._tracks = AanetTracks(
self.keys.cut_tracks_keys,
[self._data[key] for key in self.keys.tracks_keys],
fit_keys=self.keys.fit_keys)
return self._tracks
@property
def mc_tracks_keys(self):
if self._mc_tracks_keys is None:
mc_tracks_tree = self.data['Evt']['mc_trks']
self._mc_tracks_keys = [
key.decode('utf8') for key in mc_tracks_tree.keys()
]
return self._mc_tracks_keys
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,
[self._data[key] for key in self.keys.mc_hits_keys])
return self._mc_hits
@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,
[self._data[key] for key in self.keys.mc_tracks_keys])
return self._mc_tracks
class AanetEvents:
"""wrapper for Aanet events"""
def __init__(self, keys, values):
"""wrapper for aanet 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 AanetEvent(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 AanetEvent:
"""wrapper for an Aanet event"""
def __init__(self, keys, values):
"""wrapper for one aanet 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 "Aanet event:\n\t" + "\n\t".join([
"{:15} {:^10} {:>10}".format(k, ':', str(v))
for k, v in zip(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"""
def __init__(self, keys, values):
"""wrapper for aanet 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 AanetHit(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 AanetHit:
"""wrapper for an Aanet hit"""
def __init__(self, keys, values):
"""wrapper for one aanet 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 "Aanet 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]
def __repr__(self):
return str(self)
# def _is_empty(array):
# if array.size:
# return False
# else:
# return True
class AanetTracks:
"""wrapper for Aanet tracks"""
def __init__(self, keys, values, fit_keys=None):
"""Summary
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
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __getitem__(self, item):
return AanetTrack(self._keys, [v[item] for v in self._values],
fit_keys=self._fit_keys)
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))
class AanetTrack:
"""wrapper for an Aanet track"""
def __init__(self, keys, values, fit_keys=None):
"""wrapper for one aanet track.
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
for k, v in zip(self._keys, self._values):
setattr(self, k, v)
def __str__(self):
return "Aanet 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(v))
for k, v in zip(self._fit_keys, self._values[18]
) # I don't like 18 being explicit here
])
def __getitem__(self, item):
return self._values[item]
def __repr__(self):
return str(self)
notebooks/AanetReader_tutorial.ipynb
+ 136
− 48
View file @ 1ed7365e
%% Cell type:code id: tags:
``` python
# Add file to current python path
from pathlib import Path
import sys
sys.path.append(str(Path.cwd().parent))
Path.cwd()
print(sys.path)
```
%% Output
['/home/zineb/km3net/km3net/km3io/notebooks', '/home/zineb/miniconda3/envs/km3pipe/lib/python37.zip', '/home/zineb/miniconda3/envs/km3pipe/lib/python3.7', '/home/zineb/miniconda3/envs/km3pipe/lib/python3.7/lib-dynload', '', '/home/zineb/miniconda3/envs/km3pipe/lib/python3.7/site-packages', '/home/zineb/km3net/km3net/km3io', '/home/zineb/miniconda3/envs/km3pipe/lib/python3.7/site-packages/IPython/extensions', '/home/zineb/.ipython', '/home/zineb/km3net/km3net/km3io']
%% Cell type:code id: tags:
``` python
from km3io import AanetReader
```
%% Cell type:code id: tags:
``` python
# test samples directory - aanet test file
files_path = Path.cwd().parent / 'tests/samples'
aanet_file = files_path / 'aanet_v2.0.0.root'
```
%% Cell type:code id: tags:
``` python
reader = AanetReader(aanet_file)
reader
```
%% Output
Number of events: 10
Events keys are:
id
det_id
mc_id
run_id
mc_run_id
frame_index
trigger_mask
trigger_counter
overlays
hits
trks
w
w2list
w3list
mc_t
mc_hits
mc_trks
comment
index
flags
t.fSec
t.fNanoSec
Hits keys are:
hits.id
hits.dom_id
hits.channel_id
hits.tdc
hits.tot
hits.trig
hits.pmt_id
hits.t
hits.a
hits.pos.x
hits.pos.y
hits.pos.z
hits.dir.x
hits.dir.y
hits.dir.z
hits.pure_t
hits.pure_a
hits.type
hits.origin
hits.pattern_flags
Tracks keys are:
trks.fUniqueID
trks.fBits
trks.usr_data
trks.usr_names
trks.id
trks.pos.x
trks.pos.y
trks.pos.z
trks.dir.x
trks.dir.y
trks.dir.z
trks.t
trks.E
trks.len
trks.lik
trks.type
trks.rec_type
trks.rec_stages
trks.status
trks.mother_id
trks.fitinf
trks.hit_ids
trks.error_matrix
trks.comment
Mc hits keys are:
mc_hits.id
mc_hits.dom_id
mc_hits.channel_id
mc_hits.tdc
mc_hits.tot
mc_hits.trig
mc_hits.pmt_id
mc_hits.t
mc_hits.a
mc_hits.pos.x
mc_hits.pos.y
mc_hits.pos.z
mc_hits.dir.x
mc_hits.dir.y
mc_hits.dir.z
mc_hits.pure_t
mc_hits.pure_a
mc_hits.type
mc_hits.origin
mc_hits.pattern_flags
Mc tracks keys are:
mc_trks.fUniqueID
mc_trks.fBits
mc_trks.usr_data
mc_trks.usr_names
mc_trks.id
mc_trks.pos.x
mc_trks.pos.y
mc_trks.pos.z
mc_trks.dir.x
mc_trks.dir.y
mc_trks.dir.z
mc_trks.t
mc_trks.E
mc_trks.len
mc_trks.lik
mc_trks.type
mc_trks.rec_type
mc_trks.rec_stages
mc_trks.status
mc_trks.mother_id
mc_trks.fitinf
mc_trks.hit_ids
mc_trks.error_matrix
mc_trks.comment
%% Cell type:code id: tags:
``` python
reader.tracks_keys
```
%% Output
['trks.fUniqueID',
'trks.fBits',
'trks.usr_data',
'trks.usr_names',
'trks.id',
'trks.pos.x',
'trks.pos.y',
'trks.pos.z',
'trks.dir.x',
'trks.dir.y',
'trks.dir.z',
'trks.t',
'trks.E',
'trks.len',
'trks.lik',
'trks.type',
'trks.rec_type',
'trks.rec_stages',
'trks.status',
'trks.mother_id',
'trks.fitinf',
'trks.hit_ids',
'trks.error_matrix',
'trks.comment']
%% Cell type:code id: tags:
``` python
# big lazyarray with ALL file data!
lazy_data = reader.lazy_data
lazy_data = reader._lazy_data
lazy_data
```
%% Output
<Table [<Row 0> <Row 1> <Row 2> ... <Row 7> <Row 8> <Row 9>] at 0x7fb2342f1f10>
<Table [<Row 0> <Row 1> <Row 2> ... <Row 7> <Row 8> <Row 9>] at 0x7f83f9d62c10>
%% Cell type:code id: tags:
``` python
# getting the run_id for a specific event (event 5 for example)
reader[5]['run_id']
reader['run_id']
```
%% Output
5971
<ChunkedArray [5971 5971 5971 ... 5971 5971 5971] at 0x7f83f9d2cdd0>
%% Cell type:code id: tags:
``` python
# one can check how many hits are in event 5
reader[5]['hits']
```
%% Output
60
%% Cell type:code id: tags:
``` python
# one can also check how many tracks are in event 5
reader[5]['trks']
```
%% Output
56
%% Cell type:code id: tags:
``` python
# the user is reminded to always specify the "correct" event/hits/tracks
# key in the Aanet event file
try:
reader['whatever']
except KeyError as e:
print(e)
# # the user is reminded to always specify the "correct" event/hits/tracks
# # key in the Aanet event file
# try:
# reader['whatever']
# except KeyError as e:
# print(e)
```
%% Output
"'whatever' is not a valid key or is a fake branch."
%% Cell type:markdown id: tags:
# Now let's explore in more details the hits:
%% Cell type:code id: tags:
``` python
# reading all data from a specific branch in hits data: for example
# 'hits.dom_id'
reader['hits.dom_id']
```
%% Output
<ChunkedArray [[806451572 806451572 806451572 ... 809544061 809544061 809544061] [806451572 806451572 806451572 ... 809524432 809526097 809544061] [806451572 806451572 806451572 ... 809544061 809544061 809544061] ... [806451572 806455814 806465101 ... 809526097 809544058 809544061] [806455814 806455814 806455814 ... 809544061 809544061 809544061] [806455814 806455814 806455814 ... 809544058 809544058 809544061]] at 0x7fb23426e190>
<ChunkedArray [[806451572 806451572 806451572 ... 809544061 809544061 809544061] [806451572 806451572 806451572 ... 809524432 809526097 809544061] [806451572 806451572 806451572 ... 809544061 809544061 809544061] ... [806451572 806455814 806465101 ... 809526097 809544058 809544061] [806455814 806455814 806455814 ... 809544061 809544061 809544061] [806455814 806455814 806455814 ... 809544058 809544058 809544061]] at 0x7f83f9a13890>
%% Cell type:code id: tags:
``` python
# the user can access hits.dom_id data for a specific event (for example
# event 5)
reader['hits.dom_id'][5]
```
%% Output
array([806455814, 806487219, 806487219, 806487219, 806487226, 808432835,
808432835, 808432835, 808432835, 808432835, 808432835, 808432835,
808451904, 808451904, 808451907, 808451907, 808469129, 808469129,
808469129, 808493910, 808949744, 808949744, 808951460, 808951460,
808956908, 808961655, 808964908, 808969848, 808969857, 808972593,
808972593, 808972598, 808972598, 808972698, 808972698, 808974758,
808974811, 808976377, 808981510, 808981523, 808981812, 808982005,
808982005, 808982018, 808982077, 808982077, 808982547, 809007627,
809521500, 809521500, 809521500, 809524432, 809526097, 809526097,
809526097, 809526097, 809526097, 809526097, 809526097, 809544058],
dtype=int32)
%% Cell type:code id: tags:
``` python
# We previsouly checked (using reader[5]['hits']) that event
# 5 has 60 hits, now we can see that reader['hits.dom_id'][5]
# has exaclty 60 dom ids as well!
len(reader['hits.dom_id'][5])
```
%% Output
60
%% Cell type:code id: tags:
``` python
# one can access a dom id of the first hit
reader['hits.dom_id'][5][0]
```
%% Output
806455814
%% Cell type:markdown id: tags:
# Now let's explore in more details the tracks:
%% Cell type:code id: tags:
``` python
# reading all data from a specific branch in tracks data:
reader['trks.dir.z']
```
%% Output
<ChunkedArray [[-0.872885221293917 -0.872885221293917 -0.872885221293917 ... -0.6631226836266504 -0.5680647731737454 -0.5680647731737454] [-0.8351996698137462 -0.8351996698137462 -0.8351996698137462 ... -0.7485107718446855 -0.8229838871876581 -0.239315690284641] [-0.989148723802379 -0.989148723802379 -0.989148723802379 ... -0.9350162572437829 -0.88545604390297 -0.88545604390297] ... [-0.5704611045902105 -0.5704611045902105 -0.5704611045902105 ... -0.9350162572437829 -0.4647231989130516 -0.4647231989130516] [-0.9779941383490359 -0.9779941383490359 -0.9779941383490359 ... -0.88545604390297 -0.88545604390297 -0.8229838871876581] [-0.7396916780974963 -0.7396916780974963 -0.7396916780974963 ... -0.6631226836266504 -0.7485107718446855 -0.7485107718446855]] at 0x7fb234277590>
<ChunkedArray [[-0.872885221293917 -0.872885221293917 -0.872885221293917 ... -0.6631226836266504 -0.5680647731737454 -0.5680647731737454] [-0.8351996698137462 -0.8351996698137462 -0.8351996698137462 ... -0.7485107718446855 -0.8229838871876581 -0.239315690284641] [-0.989148723802379 -0.989148723802379 -0.989148723802379 ... -0.9350162572437829 -0.88545604390297 -0.88545604390297] ... [-0.5704611045902105 -0.5704611045902105 -0.5704611045902105 ... -0.9350162572437829 -0.4647231989130516 -0.4647231989130516] [-0.9779941383490359 -0.9779941383490359 -0.9779941383490359 ... -0.88545604390297 -0.88545604390297 -0.8229838871876581] [-0.7396916780974963 -0.7396916780974963 -0.7396916780974963 ... -0.6631226836266504 -0.7485107718446855 -0.7485107718446855]] at 0x7f83f9b65f90>
%% Cell type:code id: tags:
``` python
# the user can access trks.dir.z data for a specific event (for example
# event 5)
reader['trks.dir.z'][5]
```
%% Output
array([-0.60246049, -0.60246049, -0.60246049, -0.51420541, -0.5475772 ,
-0.5772408 , -0.56068238, -0.64907684, -0.67781799, -0.66565114,
-0.63014839, -0.64566464, -0.62691012, -0.58465493, -0.59287533,
-0.63655091, -0.63771247, -0.73446841, -0.7456636 , -0.70941246,
-0.66312268, -0.66312268, -0.56806477, -0.56806477, -0.66312268,
-0.66312268, -0.74851077, -0.74851077, -0.66312268, -0.74851077,
-0.56806477, -0.74851077, -0.66312268, -0.74851077, -0.56806477,
-0.66312268, -0.56806477, -0.66312268, -0.56806477, -0.56806477,
-0.66312268, -0.74851077, -0.66312268, -0.93501626, -0.56806477,
-0.74851077, -0.66312268, -0.56806477, -0.82298389, -0.74851077,
-0.66312268, -0.56806477, -0.82298389, -0.56806477, -0.66312268,
-0.97094183])
%% Cell type:code id: tags:
``` python
# We previsouly checked (using reader[5]['trks']) that event
# 5 has 56 tracks, now we can see that reader['trks.dir.z'][5]
# has exaclty 56 values of trks.dir.z as well!
len(reader['trks.dir.z'][5])
```
%% Output
56
%% Cell type:code id: tags:
``` python
# one can access the first trks.dir.z from event 5 using
reader['trks.dir.z'][5][0]
```
%% Output
-0.6024604933159441
......
tests/samples/numucc.root 0 → 100644
+ 0
− 0
View file @ 1ed7365e
File added
tests/test_aanet.py
+ 259
− 23
View file @ 1ed7365e
import unittest
from pathlib import Path
from km3io.aanet import Reader, AanetEvents, AanetHits, AanetTracks
from km3io import AanetReader
SAMPLES_DIR = Path(__file__).parent / 'samples'
AANET_FILE = SAMPLES_DIR / 'aanet_v2.0.0.root'
AANET_NUMUCC = SAMPLES_DIR / "numucc.root" # with mc data
class TestAanetReader(unittest.TestCase):
class TestAanetKeys(unittest.TestCase):
def setUp(self):
self.r = AanetReader(AANET_FILE)
self.keys = AanetReader(AANET_FILE).keys
def test_repr(self):
reader_repr = repr(self.keys)
# check that there are 106 keys + 5 extra str
self.assertEqual(len(reader_repr.split('\n')), 111)
def test_events_keys(self):
# there are 22 "valid" events keys
self.assertEqual(len(self.keys.events_keys), 22)
self.assertEqual(len(self.keys.cut_events_keys), 22)
def test_hits_keys(self):
# there are 20 "valid" hits keys
self.assertEqual(len(self.keys.hits_keys), 20)
self.assertEqual(len(self.keys.mc_hits_keys), 20)
self.assertEqual(len(self.keys.cut_hits_keys), 20)
def test_tracks_keys(self):
# there are 22 "valid" tracks keys
self.assertEqual(len(self.keys.tracks_keys), 22)
self.assertEqual(len(self.keys.mc_tracks_keys), 22)
self.assertEqual(len(self.keys.cut_tracks_keys), 22)
def test_valid_keys(self):
# there are 106 valid keys: 22*2 + 22 + 20*2
# (fit keys are excluded)
self.assertEqual(len(self.keys.valid_keys), 106)
def test_fit_keys(self):
# there are 18 fit keys
self.assertEqual(len(self.keys.fit_keys), 18)
class TestReader(unittest.TestCase):
def setUp(self):
self.r = Reader(AANET_FILE)
self.lengths = {0: 176, 1: 125, -1: 105}
self.total_item_count = 1434
......@@ -48,27 +87,18 @@ class TestAanetReader(unittest.TestCase):
self.assertListEqual([70104010.0, 70104016.0, 70104192.0],
list(ts[0][:3]))
def test_reading_hits_keys(self):
keys = self.r.hits_keys
mc_keys = self.r.mc_hits_keys
# there are 20 hits keys
self.assertEqual(len(keys), 20)
self.assertEqual(len(mc_keys), 20)
def test_reading_tracks_keys(self):
keys = self.r.tracks_keys
mc_keys = self.r.mc_tracks_keys
# there are 24 tracks keys
self.assertEqual(len(keys), 24)
self.assertEqual(len(mc_keys), 24)
def test_reading_keys(self):
all_keys = self.r.keys()
# there are 106 "valid" keys in Aanet file
self.assertEqual(len(self.r.keys.valid_keys), 106)
# there are 20 hits keys
self.assertEqual(len(self.r.keys.hits_keys), 20)
self.assertEqual(len(self.r.keys.mc_hits_keys), 20)
# there are 110 valid keys in Aanet file
self.assertEqual(len(all_keys), 110)
# there are 22 tracks keys
self.assertEqual(len(self.r.keys.tracks_keys), 22)
self.assertEqual(len(self.r.keys.mc_tracks_keys), 22)
def test_raising_KeyError(self):
# non valid keys must raise a KeyError
......@@ -77,9 +107,215 @@ class TestAanetReader(unittest.TestCase):
def test_number_events(self):
Nevents = len(self.r)
reader_repr = repr(self.r)
# check that there are 10 events
self.assertEqual(Nevents, 10)
# check that there are 110 keys + 6 extra str
self.assertEqual(len(reader_repr.split('\n')), 116)
class TestAanetReader(unittest.TestCase):
def setUp(self):
self.r = AanetReader(AANET_FILE)
self.Nevents = 10
self.selected_data = AanetReader(AANET_FILE,
data=self.r._data[0])._data
def test_item_selection(self):
# test class instance with data=None option
self.assertEqual(len(self.selected_data), len(self.r._data[0]))
# test item selection (here we test with hits=176)
self.assertEqual(self.r[0].events.hits, self.selected_data['hits'])
def test_number_events(self):
Nevents = len(self.r)
# check that there are 10 events
self.assertEqual(Nevents, self.Nevents)
class TestAanetEvents(unittest.TestCase):
def setUp(self):
self.events = AanetReader(AANET_FILE).events
self.hits = {0: 176, 1: 125, -1: 105}
self.Nevents = 10
def test_reading_hits(self):
# test item selection
for event_id, hit in self.hits.items():
self.assertEqual(hit, self.events.hits[event_id])
def reading_tracks(self):
self.assertListEqual(list(self.events.trks[:3]), [56, 55, 56])
def test_item_selection(self):
for event_id, hit in self.hits.items():
self.assertEqual(hit, self.events[event_id].hits)
def test_len(self):
self.assertEqual(len(self.events), self.Nevents)
def test_IndexError(self):
# test handling IndexError with empty lists/arrays
self.assertEqual(len(AanetEvents(['whatever'], [])), 0)
def test_str(self):
self.assertEqual(str(self.events), 'Number of events: 10')
def test_repr(self):
self.assertEqual(repr(self.events), '<AanetEvents: 10 parsed events>')
class TestAanetEvent(unittest.TestCase):
def setUp(self):
self.event = AanetReader(AANET_FILE).events[0]
def test_str(self):
self.assertEqual(repr(self.event).split('\n\t')[0], 'Aanet event:')
self.assertEqual(
repr(self.event).split('\n\t')[2],
'det_id : 44')
class TestAanetHits(unittest.TestCase):
def setUp(self):
self.hits = AanetReader(AANET_FILE).hits
self.lengths = {0: 176, 1: 125, -1: 105}
self.total_item_count = 1434
self.r_mc = AanetReader(AANET_NUMUCC)
self.Nevents = 10
def test_item_selection(self):
self.assertListEqual(list(self.hits[0].dom_id[:3]),
[806451572, 806451572, 806451572])
def test_IndexError(self):
# test handling IndexError with empty lists/arrays
self.assertEqual(len(AanetHits(['whatever'], [])), 0)
def test_repr(self):
self.assertEqual(repr(self.hits), '<AanetHits: 10 parsed elements>')
def test_str(self):
self.assertEqual(str(self.hits), 'Number of hits: 10')
def test_reading_dom_id(self):
dom_ids = self.hits.dom_id
for event_id, length in self.lengths.items():
self.assertEqual(length, len(dom_ids[event_id]))
self.assertEqual(self.total_item_count, sum(dom_ids.count()))
self.assertListEqual([806451572, 806451572, 806451572],
list(dom_ids[0][:3]))
def test_reading_channel_id(self):
channel_ids = self.hits.channel_id
for event_id, length in self.lengths.items():
self.assertEqual(length, len(channel_ids[event_id]))
self.assertEqual(self.total_item_count, sum(channel_ids.count()))
self.assertListEqual([8, 9, 14], list(channel_ids[0][:3]))
# channel IDs are always between [0, 30]
self.assertTrue(all(c >= 0 for c in channel_ids.min()))
self.assertTrue(all(c < 31 for c in channel_ids.max()))
def test_reading_times(self):
ts = self.hits.t
for event_id, length in self.lengths.items():
self.assertEqual(length, len(ts[event_id]))
self.assertEqual(self.total_item_count, sum(ts.count()))
self.assertListEqual([70104010.0, 70104016.0, 70104192.0],
list(ts[0][:3]))
def test_reading_mc_pmt_id(self):
pmt_ids = self.r_mc.mc_hits.pmt_id
lengths = {0: 58, 2: 28, -1: 48}
for hit_id, length in lengths.items():
self.assertEqual(length, len(pmt_ids[hit_id]))
self.assertEqual(self.Nevents, len(pmt_ids))
self.assertListEqual([677, 687, 689], list(pmt_ids[0][:3]))
class TestAanetHit(unittest.TestCase):
def setUp(self):
self.hit = AanetReader(AANET_FILE)[0].hits[0]
def test_item_selection(self):
self.assertEqual(self.hit[0], self.hit.id)
self.assertEqual(self.hit[1], self.hit.dom_id)
def test_str(self):
self.assertEqual(repr(self.hit).split('\n\t')[0], 'Aanet hit:')
self.assertEqual(
repr(self.hit).split('\n\t')[2],
'dom_id : 806451572')
class TestAanetTracks(unittest.TestCase):
def setUp(self):
self.tracks = AanetReader(AANET_FILE).tracks
self.r_mc = AanetReader(AANET_NUMUCC)
self.Nevents = 10
def test_item_selection(self):
self.assertListEqual(list(self.tracks[0].dir_z[:2]),
[-0.872885221293917, -0.872885221293917])
def test_IndexError(self):
# test handling IndexError with empty lists/arrays
self.assertEqual(len(AanetTracks(['whatever'], [])), 0)
def test_repr(self):
self.assertEqual(repr(self.tracks),
'<AanetTracks: 10 parsed elements>')
def test_str(self):
self.assertEqual(str(self.tracks), 'Number of tracks: 10')
def test_reading_tracks_dir_z(self):
dir_z = self.tracks.dir_z
tracks_dir_z = {0: 56, 1: 55, 8: 54}
for track_id, n_dir in tracks_dir_z.items():
self.assertEqual(n_dir, len(dir_z[track_id]))
# check that there are 10 arrays of tracks.dir_z info
self.assertEqual(len(dir_z), self.Nevents)
def test_reading_mc_tracks_dir_z(self):
dir_z = self.r_mc.mc_tracks.dir_z
tracks_dir_z = {0: 11, 1: 25, 8: 13}
for track_id, n_dir in tracks_dir_z.items():
self.assertEqual(n_dir, len(dir_z[track_id]))
# check that there are 10 arrays of tracks.dir_z info
self.assertEqual(len(dir_z), self.Nevents)
self.assertListEqual([0.230189, 0.230189, 0.218663],
list(dir_z[0][:3]))
class TestAanetTrack(unittest.TestCase):
def setUp(self):
self.track = AanetReader(AANET_FILE)[0].tracks[0]
def test_item_selection(self):
self.assertEqual(self.track[0], self.track.fUniqueID)
self.assertEqual(self.track[10], self.track.E)
def test_str(self):
self.assertEqual(repr(self.track).split('\n\t')[0], 'Aanet track:')
self.assertEqual(
repr(self.track).split('\n\t')[28],
'JGANDALF_LAMBDA : 4.2409761837248484e-12')
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment