Skip to content
Snippets Groups Projects
Commit 54d7afb7 authored by Zineb Aly's avatar Zineb Aly Committed by Tamas Gal
Browse files

Add km3net dataformat

parent 35c2fe3f
No related branches found
No related tags found
No related merge requests found
import uproot import uproot
import numpy as np import numpy as np
import warnings import warnings
import km3io.definitions.trigger
import km3io.definitions.fitparameters
import km3io.definitions.reconstruction
# 110 MB based on the size of the largest basket found so far in km3net # 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2 BASKET_CACHE_SIZE = 110 * 1024**2
...@@ -29,6 +33,9 @@ class OfflineKeys: ...@@ -29,6 +33,9 @@ class OfflineKeys:
self._cut_hits_keys = None self._cut_hits_keys = None
self._cut_tracks_keys = None self._cut_tracks_keys = None
self._cut_events_keys = None self._cut_events_keys = None
self._trigger = None
self._fitparameters = None
self._reconstruction = None
def __str__(self): def __str__(self):
return '\n'.join([ return '\n'.join([
...@@ -171,16 +178,8 @@ class OfflineKeys: ...@@ -171,16 +178,8 @@ class OfflineKeys:
list of all "trks.fitinf" keys. list of all "trks.fitinf" keys.
""" """
if self._fit_keys is None: if self._fit_keys is None:
# these are hardcoded because they are not outsourced in offline self._fit_keys = sorted(self.fitparameters, key=self.fitparameters.get, reverse=False)
# files # self._fit_keys = [*fit.keys()]
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 return self._fit_keys
@property @property
...@@ -228,6 +227,48 @@ class OfflineKeys: ...@@ -228,6 +227,48 @@ class OfflineKeys:
] ]
return self._cut_events_keys return self._cut_events_keys
@property
def trigger(self):
"""trigger parameters and their index from km3net-Dataformat.
Returns
-------
dict
dictionary of trigger parameters and their index in an Offline
file.
"""
if self._trigger is None:
self._trigger = km3io.definitions.trigger.data
return self._trigger
@property
def reconstruction(self):
"""reconstruction parameters and their index from km3net-Dataformat.
Returns
-------
dict
dictionary of reconstruction parameters and their index in an
Offline file.
"""
if self._reconstruction is None:
self._reconstruction = km3io.definitions.reconstruction.data
return self._reconstruction
@property
def fitparameters(self):
"""fit parameters parameters and their index from km3net-Dataformat.
Returns
-------
dict
dictionary of fit parameters and their index in an Offline
file.
"""
if self._fitparameters is None:
self._fitparameters = km3io.definitions.fitparameters.data
return self._fitparameters
class Reader: class Reader:
"""Reader for one offline ROOT file""" """Reader for one offline ROOT file"""
...@@ -399,7 +440,7 @@ class OfflineReader: ...@@ -399,7 +440,7 @@ class OfflineReader:
self._tracks = OfflineTracks( self._tracks = OfflineTracks(
self.keys.cut_tracks_keys, self.keys.cut_tracks_keys,
[self._data[key] for key in self.keys.tracks_keys], [self._data[key] for key in self.keys.tracks_keys],
fit_keys=self.keys.fit_keys) fitparameters=self.keys.fitparameters)
return self._tracks return self._tracks
@property @property
...@@ -429,7 +470,8 @@ class OfflineReader: ...@@ -429,7 +470,8 @@ class OfflineReader:
if self._mc_tracks is None: if self._mc_tracks is None:
self._mc_tracks = OfflineTracks( self._mc_tracks = OfflineTracks(
self.keys.cut_tracks_keys, self.keys.cut_tracks_keys,
[self._data[key] for key in self.keys.mc_tracks_keys]) [self._data[key] for key in self.keys.mc_tracks_keys],
fitparameters=self.keys.fitparameters)
return self._mc_tracks return self._mc_tracks
@property @property
...@@ -717,7 +759,7 @@ class OfflineHit: ...@@ -717,7 +759,7 @@ class OfflineHit:
class OfflineTracks: class OfflineTracks:
"""wrapper for offline tracks""" """wrapper for offline tracks"""
def __init__(self, keys, values, fit_keys=None): def __init__(self, keys, values, fitparameters=None):
"""wrapper for offline tracks """wrapper for offline tracks
Parameters Parameters
...@@ -726,26 +768,26 @@ class OfflineTracks: ...@@ -726,26 +768,26 @@ class OfflineTracks:
list of cropped tracks keys. list of cropped tracks keys.
values : list of arrays values : list of arrays
list of arrays containting tracks data. list of arrays containting tracks data.
fit_keys : None, optional fitparameters : None, optional
list of tracks fit information (not yet outsourced in offline dictionary of tracks fit information (not yet outsourced in offline
files). files).
""" """
self._keys = keys self._keys = keys
self._values = values self._values = values
if fit_keys is not None: if fitparameters is not None:
self._fit_keys = fit_keys self._fitparameters = fitparameters
for k, v in zip(self._keys, self._values): for k, v in zip(self._keys, self._values):
setattr(self, k, v) setattr(self, k, v)
def __getitem__(self, item): def __getitem__(self, item):
if isinstance(item, int): if isinstance(item, int):
return OfflineTrack(self._keys, [v[item] for v in self._values], return OfflineTrack(self._keys, [v[item] for v in self._values],
fit_keys=self._fit_keys) fitparameters=self._fitparameters)
else: else:
return OfflineTracks( return OfflineTracks(
self._keys, self._keys,
[v[item] for v in self._values], [v[item] for v in self._values],
fit_keys=self._fit_keys fitparameters=self._fitparameters
) )
def __len__(self): def __len__(self):
...@@ -764,7 +806,7 @@ class OfflineTracks: ...@@ -764,7 +806,7 @@ class OfflineTracks:
class OfflineTrack: class OfflineTrack:
"""wrapper for an offline track""" """wrapper for an offline track"""
def __init__(self, keys, values, fit_keys=None): def __init__(self, keys, values, fitparameters=None):
"""wrapper for one offline track. """wrapper for one offline track.
Parameters Parameters
...@@ -773,14 +815,14 @@ class OfflineTrack: ...@@ -773,14 +815,14 @@ class OfflineTrack:
list of cropped tracks keys. list of cropped tracks keys.
values : list of arrays values : list of arrays
list of arrays containting track data. list of arrays containting track data.
fit_keys : None, optional fitparameters : None, optional
list of tracks fit information (not yet outsourced in offline dictionary of tracks fit information (not yet outsourced in offline
files). files).
""" """
self._keys = keys self._keys = keys
self._values = values self._values = values
if fit_keys is not None: if fitparameters is not None:
self._fit_keys = fit_keys self._fitparameters = fitparameters
for k, v in zip(self._keys, self._values): for k, v in zip(self._keys, self._values):
setattr(self, k, v) setattr(self, k, v)
...@@ -789,10 +831,9 @@ class OfflineTrack: ...@@ -789,10 +831,9 @@ class OfflineTrack:
"{:30} {:^2} {:>26}".format(k, ':', str(v)) "{:30} {:^2} {:>26}".format(k, ':', str(v))
for k, v in zip(self._keys, self._values) if k not in ['fitinf'] for k, v in zip(self._keys, self._values) if k not in ['fitinf']
]) + "\n\t" + "\n\t".join([ ]) + "\n\t" + "\n\t".join([
"{:30} {:^2} {:>26}".format(k, ':', str(v)) "{:30} {:^2} {:>26}".format(k, ':', str(self._values[18][v]))
for k, v in zip(self._fit_keys, self._values[18] for k, v in self._fitparameters.items() if len(self._values[18])>v
) # I don't like 18 being explicit here ]) # I don't like 18 being explicit here
])
def __getitem__(self, item): def __getitem__(self, item):
return self._values[item] return self._values[item]
......
...@@ -46,6 +46,36 @@ class TestOfflineKeys(unittest.TestCase): ...@@ -46,6 +46,36 @@ class TestOfflineKeys(unittest.TestCase):
# there are 18 fit keys # there are 18 fit keys
self.assertEqual(len(self.keys.fit_keys), 18) self.assertEqual(len(self.keys.fit_keys), 18)
def test_trigger(self):
# there are 4 trigger keys in v1.1.2 of km3net-Dataformat
trigger = self.keys.trigger
keys = ['JTRIGGER3DSHOWER', 'JTRIGGERMXSHOWER',
'JTRIGGER3DMUON', 'JTRIGGERNB']
values = [1, 2, 4, 5]
for k, v in zip(keys, values):
self.assertEqual(v, trigger[k])
def test_reconstruction(self):
# there are 34 parameters in v1.1.2 of km3net-Dataformat
reco = self.keys.reconstruction
keys = ['JPP_RECONSTRUCTION_TYPE', 'JMUONFIT', 'JMUONBEGIN',
'JMUONPREFIT', 'JMUONSIMPLEX', 'JMUONGANDALF',
'JMUONENERGY', 'JMUONSTART']
values = [4000, 0, 0, 1, 2, 3, 4, 5]
self.assertEqual(34, len([*reco.keys()]))
for k, v in zip(keys, values):
self.assertEqual(v, reco[k])
def test_fitparameters(self):
# there are 18 parameters in v1.1.2 of km3net-Dataformat
fit = self.keys.fitparameters
values = [i for i in range(18)]
self.assertEqual(18, len([*fit.keys()]))
for k, v in fit.items():
self.assertEqual(values[v], fit[k])
class TestReader(unittest.TestCase): class TestReader(unittest.TestCase):
def setUp(self): def setUp(self):
...@@ -399,6 +429,4 @@ class TestOfflineTrack(unittest.TestCase): ...@@ -399,6 +429,4 @@ class TestOfflineTrack(unittest.TestCase):
def test_str(self): def test_str(self):
self.assertEqual(repr(self.track).split('\n\t')[0], 'offline track:') self.assertEqual(repr(self.track).split('\n\t')[0], 'offline track:')
self.assertEqual( self.assertTrue("JGANDALF_LAMBDA" in repr(self.track))
repr(self.track).split('\n\t')[28],
'JGANDALF_LAMBDA : 4.2409761837248484e-12')
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment