Skip to content
Snippets Groups Projects
Commit e7eb425b authored by Zineb Aly's avatar Zineb Aly
Browse files

adapted AanetReader format to pep8

parent 1756e937
No related branches found
No related tags found
No related merge requests found
Pipeline #6939 failed
import uproot import uproot
class AanetReader: class AanetReader:
"""Reader for one Aanet ROOT file""" """Reader for one Aanet ROOT file"""
def __init__(self, file_path): def __init__(self, file_path):
""" AanetReader class is a Aanet ROOT file wrapper """ AanetReader class is a Aanet ROOT file wrapper
Parameters Parameters
---------- ----------
file_path : path-like object file_path : path-like object
Path to the file of interest. It can be a str or any python path-like object that points Path to the file of interest. It can be a str or any python
to the file of ineterst. path-like object that points to the file of ineterst.
""" """
self.file_path = file_path self.file_path = file_path
self.data = uproot.open(self.file_path)['E'] self.data = uproot.open(self.file_path)['E']
...@@ -18,52 +19,65 @@ class AanetReader: ...@@ -18,52 +19,65 @@ class AanetReader:
self._hits_keys = None self._hits_keys = None
self._tracks_keys = None self._tracks_keys = None
def __getitem__(self, key): 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 event tree.
Parameters Parameters
---------- ----------
key : str key : str
name of the branch of interest in event data. name of the branch of interest in event data.
Returns Returns
------- -------
lazyarray lazyarray
Lazyarray of all data stored in the branch of interest. A lazyarray is an array-like Lazyarray of all data stored in the branch of interest. A lazyarray
object that reads data on demand. Here, only the first and last chunks of data are is an array-like object that reads data on demand. Here, only the
read in memory, and not all data in the array. The output can be used first and last chunks of data are read in memory, and not all data
with all `Numpy's universal functions <https://docs.scipy.org/doc/numpy/reference/ufuncs.html>`. in the array. The output can be used with all `Numpy's universal
functions <https://docs.scipy.org/doc/numpy/reference/ufuncs.html>`
.
Raises Raises
------ ------
KeyEroor KeyEroor
Some branches in an Aanet file structure are "fake branches" and do not contain data. Therefore, Some branches in an Aanet file structure are "fake branches" and do
the keys corresponding to these fake branches are not read. 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() and not isinstance(key, int):
raise KeyError(f"'{key}' is not a valid key or is a fake branch.") raise KeyError(f"'{key}' is not a valid key or is a fake branch.")
return self.lazy_data[key] return self.lazy_data[key]
def __len__(self): def __len__(self):
return len(self.lazy_data) return len(self.lazy_data)
def __repr__(self): def __repr__(self):
return '\n'.join([f"Number of events: {self.__len__()}", return '\n'.join([
"Events keys are:\n\t" + '\n\t'.join(self.events_keys), f"Number of events: {self.__len__()}",
"Hits keys are:\n\t" + '\n\t'.join(self.hits_keys), "Events keys are:\n\t" + '\n\t'.join(self.events_keys),
"Tracks keys are:\n\t" + '\n\t'.join(self.tracks_keys)]) "Hits keys are:\n\t" + '\n\t'.join(self.hits_keys),
"Tracks keys are:\n\t" + '\n\t'.join(self.tracks_keys)
])
def keys(self): def keys(self):
return self.events_keys + self.hits_keys + self.tracks_keys """constructs a list of all valid keys to be read from an Aanet event file.
Returns
-------
list
list of all valid keys.
"""
return self.events_keys + self.hits_keys + self.tracks_keys
@property @property
def events_keys(self): def events_keys(self):
if self._events_keys is None: if self._events_keys is None:
fake_branches = ['Evt', 'AAObject', 'TObject','t'] fake_branches = ['Evt', 'AAObject', 'TObject', 't']
t_baskets = ['t.fSec', 't.fNanoSec'] 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 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 return self._events_keys
@property @property
...@@ -77,5 +91,7 @@ class AanetReader: ...@@ -77,5 +91,7 @@ class AanetReader:
def tracks_keys(self): def tracks_keys(self):
if self._tracks_keys is None: if self._tracks_keys is None:
tracks_tree = self.data['Evt']['trks'] tracks_tree = self.data['Evt']['trks']
self._tracks_keys = [key.decode('utf8') for key in tracks_tree.keys()] self._tracks_keys = [
key.decode('utf8') for key in tracks_tree.keys()
]
return self._tracks_keys return self._tracks_keys
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# Add file to current python path # Add file to current python path
from pathlib import Path from pathlib import Path
import sys import sys
sys.path.append(str(Path.cwd().parent)) sys.path.append(str(Path.cwd().parent))
Path.cwd() Path.cwd()
print(sys.path) print(sys.path)
``` ```
%% Output %% 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/miniconda3/envs/km3pipe/lib/python3.7/site-packages/IPython/extensions', '/home/zineb/.ipython', '/home/zineb/km3net/km3net/km3io'] ['/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/miniconda3/envs/km3pipe/lib/python3.7/site-packages/IPython/extensions', '/home/zineb/.ipython', '/home/zineb/km3net/km3net/km3io']
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
from km3io import AanetReader from km3io import AanetReader
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# test samples directory - aanet test file # test samples directory - aanet test file
files_path = Path.cwd().parent / 'tests/samples' files_path = Path.cwd().parent / 'tests/samples'
aanet_file = files_path / 'aanet_v2.0.0.root' aanet_file = files_path / 'aanet_v2.0.0.root'
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
reader = AanetReader(aanet_file) reader = AanetReader(aanet_file)
reader reader
``` ```
%% Output %% Output
Number of events: 10 Number of events: 10
Events keys are: Events keys are:
id id
det_id det_id
mc_id mc_id
run_id run_id
mc_run_id mc_run_id
frame_index frame_index
trigger_mask trigger_mask
trigger_counter trigger_counter
overlays overlays
hits hits
trks trks
w w
w2list w2list
w3list w3list
mc_t mc_t
mc_hits mc_hits
mc_trks mc_trks
comment comment
index index
flags flags
t.fSec t.fSec
t.fNanoSec t.fNanoSec
Hits keys are: Hits keys are:
hits.id hits.id
hits.dom_id hits.dom_id
hits.channel_id hits.channel_id
hits.tdc hits.tdc
hits.tot hits.tot
hits.trig hits.trig
hits.pmt_id hits.pmt_id
hits.t hits.t
hits.a hits.a
hits.pos.x hits.pos.x
hits.pos.y hits.pos.y
hits.pos.z hits.pos.z
hits.dir.x hits.dir.x
hits.dir.y hits.dir.y
hits.dir.z hits.dir.z
hits.pure_t hits.pure_t
hits.pure_a hits.pure_a
hits.type hits.type
hits.origin hits.origin
hits.pattern_flags hits.pattern_flags
Tracks keys are: Tracks keys are:
trks.fUniqueID trks.fUniqueID
trks.fBits trks.fBits
trks.usr_data trks.usr_data
trks.usr_names trks.usr_names
trks.id trks.id
trks.pos.x trks.pos.x
trks.pos.y trks.pos.y
trks.pos.z trks.pos.z
trks.dir.x trks.dir.x
trks.dir.y trks.dir.y
trks.dir.z trks.dir.z
trks.t trks.t
trks.E trks.E
trks.len trks.len
trks.lik trks.lik
trks.type trks.type
trks.rec_type trks.rec_type
trks.rec_stages trks.rec_stages
trks.status trks.status
trks.mother_id trks.mother_id
trks.fitinf trks.fitinf
trks.hit_ids trks.hit_ids
trks.error_matrix trks.error_matrix
trks.comment trks.comment
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# big lazyarray with ALL file data! # big lazyarray with ALL file data!
lazy_data = reader.lazy_data lazy_data = reader.lazy_data
lazy_data lazy_data
``` ```
%% Output %% Output
<Table [<Row 0> <Row 1> <Row 2> ... <Row 7> <Row 8> <Row 9>] at 0x7f57ac134810> <Table [<Row 0> <Row 1> <Row 2> ... <Row 7> <Row 8> <Row 9>] at 0x7f9011f02b50>
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# getting the run_id for a specific event (event 5 for example) # getting the run_id for a specific event (event 5 for example)
reader[5]['run_id'] reader[5]['run_id']
``` ```
%% Output %% Output
5971 5971
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# one can check how many hits are in event 5 # one can check how many hits are in event 5
reader[5]['hits'] reader[5]['hits']
``` ```
%% Output %% Output
60 60
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# one can also check how many tracks are in event 5 # one can also check how many tracks are in event 5
reader[5]['trks'] reader[5]['trks']
``` ```
%% Output %% Output
56 56
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# the user is reminded to always specify the "correct" event/hits/tracks # the user is reminded to always specify the "correct" event/hits/tracks
# key in the Aanet event file # key in the Aanet event file
try: try:
reader['whatever'] reader['whatever']
except KeyError as e: except KeyError as e:
print(e) print(e)
``` ```
%% Output %% Output
"'whatever' is not a valid key or is a fake branch." "'whatever' is not a valid key or is a fake branch."
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
# Now let's explore in more details the hits: # Now let's explore in more details the hits:
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# reading all data from a specific branch in hits data: for example # reading all data from a specific branch in hits data: for example
# 'hits.dom_id' # 'hits.dom_id'
reader['hits.dom_id'] reader['hits.dom_id']
``` ```
%% Output %% 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 0x7f57997ee990> <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 0x7f9011eb3b10>
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# the user can access hits.dom_id data for a specific event (for example # the user can access hits.dom_id data for a specific event (for example
# event 5) # event 5)
reader['hits.dom_id'][5] reader['hits.dom_id'][5]
``` ```
%% Output %% Output
array([806455814, 806487219, 806487219, 806487219, 806487226, 808432835, array([806455814, 806487219, 806487219, 806487219, 806487226, 808432835,
808432835, 808432835, 808432835, 808432835, 808432835, 808432835, 808432835, 808432835, 808432835, 808432835, 808432835, 808432835,
808451904, 808451904, 808451907, 808451907, 808469129, 808469129, 808451904, 808451904, 808451907, 808451907, 808469129, 808469129,
808469129, 808493910, 808949744, 808949744, 808951460, 808951460, 808469129, 808493910, 808949744, 808949744, 808951460, 808951460,
808956908, 808961655, 808964908, 808969848, 808969857, 808972593, 808956908, 808961655, 808964908, 808969848, 808969857, 808972593,
808972593, 808972598, 808972598, 808972698, 808972698, 808974758, 808972593, 808972598, 808972598, 808972698, 808972698, 808974758,
808974811, 808976377, 808981510, 808981523, 808981812, 808982005, 808974811, 808976377, 808981510, 808981523, 808981812, 808982005,
808982005, 808982018, 808982077, 808982077, 808982547, 809007627, 808982005, 808982018, 808982077, 808982077, 808982547, 809007627,
809521500, 809521500, 809521500, 809524432, 809526097, 809526097, 809521500, 809521500, 809521500, 809524432, 809526097, 809526097,
809526097, 809526097, 809526097, 809526097, 809526097, 809544058], 809526097, 809526097, 809526097, 809526097, 809526097, 809544058],
dtype=int32) dtype=int32)
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# We previsouly checked (using reader[5]['hits']) that event # We previsouly checked (using reader[5]['hits']) that event
# 5 has 60 hits, now we can see that reader['hits.dom_id'][5] # 5 has 60 hits, now we can see that reader['hits.dom_id'][5]
# has exaclty 60 dom ids as well! # has exaclty 60 dom ids as well!
len(reader['hits.dom_id'][5]) len(reader['hits.dom_id'][5])
``` ```
%% Output %% Output
60 60
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# one can access a dom id of the first hit # one can access a dom id of the first hit
reader['hits.dom_id'][5][0] reader['hits.dom_id'][5][0]
``` ```
%% Output %% Output
806455814 806455814
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
# Now let's explore in more details the tracks: # Now let's explore in more details the tracks:
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# reading all data from a specific branch in tracks data: # reading all data from a specific branch in tracks data:
reader['trks.dir.z'] reader['trks.dir.z']
``` ```
%% Output %% 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 0x7f57997f8d90> <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 0x7f9011ebdf10>
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# the user can access trks.dir.z data for a specific event (for example # the user can access trks.dir.z data for a specific event (for example
# event 5) # event 5)
reader['trks.dir.z'][5] reader['trks.dir.z'][5]
``` ```
%% Output %% Output
array([-0.60246049, -0.60246049, -0.60246049, -0.51420541, -0.5475772 , array([-0.60246049, -0.60246049, -0.60246049, -0.51420541, -0.5475772 ,
-0.5772408 , -0.56068238, -0.64907684, -0.67781799, -0.66565114, -0.5772408 , -0.56068238, -0.64907684, -0.67781799, -0.66565114,
-0.63014839, -0.64566464, -0.62691012, -0.58465493, -0.59287533, -0.63014839, -0.64566464, -0.62691012, -0.58465493, -0.59287533,
-0.63655091, -0.63771247, -0.73446841, -0.7456636 , -0.70941246, -0.63655091, -0.63771247, -0.73446841, -0.7456636 , -0.70941246,
-0.66312268, -0.66312268, -0.56806477, -0.56806477, -0.66312268, -0.66312268, -0.66312268, -0.56806477, -0.56806477, -0.66312268,
-0.66312268, -0.74851077, -0.74851077, -0.66312268, -0.74851077, -0.66312268, -0.74851077, -0.74851077, -0.66312268, -0.74851077,
-0.56806477, -0.74851077, -0.66312268, -0.74851077, -0.56806477, -0.56806477, -0.74851077, -0.66312268, -0.74851077, -0.56806477,
-0.66312268, -0.56806477, -0.66312268, -0.56806477, -0.56806477, -0.66312268, -0.56806477, -0.66312268, -0.56806477, -0.56806477,
-0.66312268, -0.74851077, -0.66312268, -0.93501626, -0.56806477, -0.66312268, -0.74851077, -0.66312268, -0.93501626, -0.56806477,
-0.74851077, -0.66312268, -0.56806477, -0.82298389, -0.74851077, -0.74851077, -0.66312268, -0.56806477, -0.82298389, -0.74851077,
-0.66312268, -0.56806477, -0.82298389, -0.56806477, -0.66312268, -0.66312268, -0.56806477, -0.82298389, -0.56806477, -0.66312268,
-0.97094183]) -0.97094183])
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# We previsouly checked (using reader[5]['trks']) that event # We previsouly checked (using reader[5]['trks']) that event
# 5 has 56 tracks, now we can see that reader.hits_reader('hits.dom_id')[5] # 5 has 56 tracks, now we can see that reader['trks.dir.z'][5]
# has exaclty 56 values of trks.dir.z as well! # has exaclty 56 values of trks.dir.z as well!
len(reader.read_tracks('trks.dir.z')[5]) len(reader['trks.dir.z'][5])
``` ```
%% Cell type:code id: tags: %% Output
``` python
# one can access the first trks.dir.z from event 5 using
reader.read_tracks('trks.dir.z')[5][0]
```
%% Cell type:code id: tags:
``` python
# once again, the user is reminded to always use the tracks keys stored in Aanet
# event file
try:
reader.read_tracks('whatever')
except KeyError as e:
print(e)
```
%% Cell type:code id: tags:
``` python
reader['trks.dir.z']
```
%% Cell type:code id: tags:
``` python
import numpy as np
np.array(reader[0])
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:code id: tags:
``` python 56
```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# one can access the first trks.dir.z from event 5 using
reader['trks.dir.z'][5][0]
``` ```
%% Cell type:code id: tags: %% Output
``` python -0.6024604933159441
```
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment