from __future__ import absolute_import, division, print_function
import NuRadioReco.framework.parameters as parameters
import NuRadioReco.framework.parameter_serialization
import pickle
import collections
import math
import logging
logger = logging.getLogger('Particle')
[docs]class Particle:
def __init__(self, particle_index):
# "_id" is not the PDG code but a hierarchical index
# (PDG code is stored in _parameters["flavor"])
self._id = particle_index
self._parameters = {}
def __setitem__(self, key, value):
self.set_parameter(key, value)
def __getitem__(self, key):
return self.get_parameter(key)
[docs] def get_id(self):
""" Returns hierarchical index """
return self._id
def __str__(self):
msg = (
"Particle ({}): "
"Flavor: {: 3}, lgE = {:.1f}, cos(theta) = {:.2f}".format(
hex(id(self)),
self.get_parameter(parameters.particleParameters.flavor),
math.log10(self.get_parameter(parameters.particleParameters.energy)),
math.cos(self.get_parameter(parameters.particleParameters.zenith)))
)
return msg
[docs] def get_parameter(self, key):
if not isinstance(key, parameters.particleParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.particleParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.particleParameters")
return self._parameters[key]
[docs] def set_parameter(self, key, value):
if not isinstance(key, parameters.particleParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.particleParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.particleParameters")
self._parameters[key] = value
[docs] def has_parameter(self, key):
if not isinstance(key, parameters.particleParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.particleParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.particleParameters")
return key in self._parameters
[docs] def as_hdf5_dict(self):
hdf5_dict = collections.OrderedDict()
hdf5_dict['azimuths'] = self.get_parameter(parameters.particleParameters.azimuth)
hdf5_dict['energies'] = self.get_parameter(parameters.particleParameters.energy)
hdf5_dict['event_group_ids'] = self.get_id()
hdf5_dict['flavors'] = self.get_parameter(parameters.particleParameters.flavor)
hdf5_dict['inelasticity'] = self.get_parameter(parameters.particleParameters.inelasticity)
hdf5_dict['interaction_type'] = self.get_parameter(parameters.particleParameters.interaction_type)
hdf5_dict['n_interaction'] = self.get_parameter(parameters.particleParameters.n_interaction)
hdf5_dict['vertex_times'] = self.get_parameter(parameters.particleParameters.vertex_time)
hdf5_dict['weights'] = self.get_parameter(parameters.particleParameters.weight)
hdf5_dict['xx'] = self.get_parameter(parameters.particleParameters.vertex[0])
hdf5_dict['yy'] = self.get_parameter(parameters.particleParameters.vertex[1])
hdf5_dict['zeniths'] = self.get_parameter(parameters.particleParameters.zenith)
hdf5_dict['zz'] = self.get_parameter(parameters.particleParameters.vertex[2])
return hdf5_dict
[docs] def serialize(self):
data = {'_parameters': NuRadioReco.framework.parameter_serialization.serialize(self._parameters),
'_id': self._id}
return pickle.dumps(data, protocol=4)
[docs] def deserialize(self, data_pkl):
data = pickle.loads(data_pkl)
if '_id' in data.keys():
self._id = data['_id']
else:
self._id = None
self._parameters = NuRadioReco.framework.parameter_serialization.deserialize(
data['_parameters'],
parameters.particleParameters
)