from __future__ import absolute_import, division, print_function
import NuRadioReco.framework.base_trace
import NuRadioReco.framework.parameters as parameters
import NuRadioReco.framework.parameter_serialization
try:
import cPickle as pickle
except ImportError:
import pickle
import logging
logger = logging.getLogger('NuRadioReco.ElectricField')
[docs]class ElectricField(NuRadioReco.framework.base_trace.BaseTrace):
def __init__(self, channel_ids, position=None,
shower_id=None, ray_tracing_id=None):
"""
Initialize a new electric field object
This object stores a 3 dimensional trace plus additional meta parameters
Parameters
----------
channel_ids: array of ints
the channels ids this electric field is valid for.
(For cosmic rays one electric field is typically valid
for several channels. For neutrino simulations, we typically
simulate the electric field for each
channel separately)
position: 3-dim array/list of floats
the position of the electric field
shower_id: int or None
the id of the corresponding shower object
ray_tracing_id: int or None
the id of the corresponding ray tracing solution
"""
NuRadioReco.framework.base_trace.BaseTrace.__init__(self)
self._channel_ids = channel_ids
self._parameters = {}
self._parameter_covariances = {}
self._position = position
if position is None:
self._position = [0, 0, 0]
self._shower_id = shower_id
self._ray_tracing_id = ray_tracing_id
[docs] def get_unique_identifier(self):
"""
returns a unique identifier consisting of the tuple channel_ids, shower_id and ray_tracing_id
"""
return (self._channel_ids, self._shower_id, self._ray_tracing_id)
[docs] def get_parameter(self, key):
if not isinstance(key, parameters.electricFieldParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
return self._parameters[key]
[docs] def get_parameters(self):
return self._parameters
[docs] def set_parameter(self, key, value):
if not isinstance(key, parameters.electricFieldParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
self._parameters[key] = value
[docs] def has_parameter(self, key):
if not isinstance(key, parameters.electricFieldParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
return key in self._parameters
[docs] def has_parameter_error(self, key):
if not isinstance(key, parameters.electricFieldParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
return (key, key) in self._parameter_covariances
[docs] def set_parameter_error(self, key, value):
if not isinstance(key, parameters.electricFieldParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
self._parameter_covariances[(key, key)] = value ** 2
[docs] def get_parameter_error(self, key):
if not isinstance(key, parameters.electricFieldParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.electricFieldParameters")
return self._parameter_covariances[(key, key)] ** 0.5
def __setitem__(self, key, value):
self.set_parameter(key, value)
def __getitem__(self, key):
return self.get_parameter(key)
[docs] def set_channel_ids(self, channel_ids):
self._channel_ids = channel_ids
[docs] def get_channel_ids(self):
return self._channel_ids
[docs] def has_channel_ids(self, channel_ids):
for channel_id in channel_ids:
if channel_id not in self._channel_ids:
return False
return True
[docs] def get_shower_id(self):
return self._shower_id
[docs] def get_ray_tracing_solution_id(self):
return self._ray_tracing_id
[docs] def get_position(self):
"""
get position of the electric field relative to station position
"""
return self._position
[docs] def set_position(self, position):
"""
set position of the electric field relative to station position
"""
self._position = position
[docs] def serialize(self, save_trace):
if save_trace:
base_trace_pkl = NuRadioReco.framework.base_trace.BaseTrace.serialize(self)
else:
base_trace_pkl = None
data = {'parameters': NuRadioReco.framework.parameter_serialization.serialize(self._parameters),
'parameter_covariances': NuRadioReco.framework.parameter_serialization.serialize_covariances(
self._parameter_covariances),
'channel_ids': self._channel_ids,
'_shower_id': self._shower_id,
'_ray_tracing_id': self._ray_tracing_id,
'position': self._position,
'base_trace': base_trace_pkl}
return pickle.dumps(data, protocol=4)
[docs] def deserialize(self, data_pkl):
data = pickle.loads(data_pkl)
if data['base_trace'] is not None:
NuRadioReco.framework.base_trace.BaseTrace.deserialize(self, data['base_trace'])
if 'position' in data: # for backward compatibility
self._position = data['position']
self._parameters = NuRadioReco.framework.parameter_serialization.deserialize(
data['parameters'], parameters.electricFieldParameters)
if 'parameter_covariances' in data:
self._parameter_covariances = NuRadioReco.framework.parameter_serialization.deserialize_covariances(
data['parameter_covariances'], parameters.electricFieldParameters)
self._channel_ids = data['channel_ids']
self._shower_id = data.get('_shower_id', None)
self._ray_tracing_id = data.get('_ray_tracing_id', None)