from __future__ import absolute_import, division, print_function
import pickle
import NuRadioReco.framework.station
import NuRadioReco.framework.radio_shower
import NuRadioReco.framework.emitter
import NuRadioReco.framework.sim_emitter
import NuRadioReco.framework.hybrid_information
import NuRadioReco.framework.particle
import NuRadioReco.framework.parameters as parameters
import NuRadioReco.utilities.version
from six import itervalues
import collections
import logging
logger = logging.getLogger('Event')
[docs]class Event:
def __init__(self, run_number, event_id):
self._parameters = {}
self.__run_number = run_number
self._id = event_id
self.__stations = collections.OrderedDict()
self.__radio_showers = collections.OrderedDict()
self.__sim_showers = collections.OrderedDict()
self.__sim_emitters = collections.OrderedDict()
self.__event_time = 0
self.__particles = collections.OrderedDict() # stores a dictionary of simulated MC particles in an event
self._generator_info = {} # copies over the relevant information on event generation from the input file attributes
self.__hybrid_information = NuRadioReco.framework.hybrid_information.HybridInformation()
self.__modules_event = [] # saves which modules were executed with what parameters on event level
self.__modules_station = {} # saves which modules were executed with what parameters on station level
[docs] def register_module_event(self, instance, name, kwargs):
"""
registers modules applied to this event
Parameters
----------
instance: module instance
the instance of the module that should be registered
name: module name
the name of the module
kwargs:
the key word arguments of the run method
"""
self.__modules_event.append([name, instance, kwargs])
[docs] def register_module_station(self, station_id, instance, name, kwargs):
"""
registers modules applied to this event
Parameters
----------
station_id: int
the station id
instance: module instance
the instance of the module that should be registered
name: module name
the name of the module
kwargs:
the key word arguments of the run method
"""
if station_id not in self.__modules_station:
self.__modules_station[station_id] = []
iE = len(self.__modules_event)
self.__modules_station[station_id].append([iE, name, instance, kwargs])
[docs] def iter_modules(self, station_id=None):
"""
returns an interator that loops over all modules. If a station id is provided it loops
over all modules that are applied on event or station level (on this particular station). If no
station_id is provided, the loop is only over the event modules.
The order follows the sequence these modules were applied
"""
iE = 0
iS = 0
while True:
if(station_id in self.__modules_station and (len(self.__modules_station[station_id]) > iS) and self.__modules_station[station_id][iS][0] == iE):
iS += 1
yield self.__modules_station[station_id][iS - 1][1:]
else:
if(len(self.__modules_event) == iE):
break
iE += 1
yield self.__modules_event[iE - 1]
[docs] def get_parameter(self, key):
if not isinstance(key, parameters.eventParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.eventParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.eventParameters")
return self._parameters[key]
[docs] def set_parameter(self, key, value):
if not isinstance(key, parameters.eventParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.eventParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.eventParameters")
self._parameters[key] = value
[docs] def has_parameter(self, key):
if not isinstance(key, parameters.eventParameters):
logger.error("parameter key needs to be of type NuRadioReco.framework.parameters.eventParameters")
raise ValueError("parameter key needs to be of type NuRadioReco.framework.parameters.eventParameters")
return key in self._parameters
[docs] def get_generator_info(self, key):
if not isinstance(key, parameters.generatorAttributes):
logger.error("generator information key needs to be of type NuRadioReco.framework.parameters.generatorAttributes")
raise ValueError("generator information key needs to be of type NuRadioReco.framework.parameters.generatorAttributes")
return self._generator_info[key]
[docs] def set_generator_info(self, key, value):
if not isinstance(key, parameters.generatorAttributes):
logger.error("generator information key needs to be of type NuRadioReco.framework.parameters.generatorAttributes")
raise ValueError("generator information key needs to be of type NuRadioReco.framework.parameters.generatorAttributes")
self._generator_info[key] = value
[docs] def has_generator_info(self, key):
if not isinstance(key, parameters.generatorAttributes):
logger.error("generator information key needs to be of type NuRadioReco.framework.parameters.generatorAttributes")
raise ValueError("generator information key needs to be of type NuRadioReco.framework.parameters.generatorAttributes")
return key in self._generator_info
[docs] def get_id(self):
return self._id
[docs] def set_id(self, evt_id):
self._id = evt_id
[docs] def get_run_number(self):
return self.__run_number
[docs] def get_station(self, station_id=None):
"""
Returns the station for a given station id.
Parameters
----------
station_id: int
Id of the station you want to get. If None and event has only one station
return it, otherwise raise error. (Default: None)
Returns
-------
station: NuRadioReco.framework.station
"""
if station_id is None:
if len(self.get_station_ids()) == 1:
return self.__stations[self.get_station_ids()[0]]
else:
err = "Event has more than one station, you have to specify \"station_id\""
logger.error(err)
raise ValueError(err)
return self.__stations[station_id]
[docs] def get_stations(self):
for station in itervalues(self.__stations):
yield station
[docs] def get_station_ids(self):
return list(self.__stations.keys())
[docs] def set_station(self, station):
self.__stations[station.get_id()] = station
[docs] def has_triggered(self, trigger_name=None):
"""
Returns true if any station has been triggered.
Parameters
----------
trigger_name: string or None (default None)
* if None: The function returns False if not trigger was set. If one or multiple triggers were set,
it returns True if any of those triggers triggered
* if trigger name is set: return if the trigger with name 'trigger_name' has a trigger
Returns
-------
has_triggered : bool
"""
for station in self.get_stations():
if station.has_triggered(trigger_name):
return True
# if it reaches this point, no station has a trigger
return False
[docs] def add_particle(self, particle):
"""
Adds a MC particle to the event
Parameters
----------
particle : NuRadioReco.framework.particle.Particle
The MC particle to be added to the event
"""
if not isinstance(particle, NuRadioReco.framework.particle.Particle):
logger.error("Requested to add non-Particle item to the list of particles. {particle} needs to be an instance of Particle.")
raise TypeError("Requested to add non-Particle item to the list of particles. {particle} needs to be an instance of Particle.")
if particle.get_id() in self.__particles:
logger.error("MC particle with id {particle.get_id()} already exists. Simulated particle id needs to be unique per event")
raise AttributeError("MC particle with id {particle.get_id()} already exists. Simulated particle id needs to be unique per event")
self.__particles[particle.get_id()] = particle
[docs] def get_particles(self):
"""
Returns an iterator over the MC particles stored in the event
"""
for particle in self.__particles.values():
yield particle
[docs] def get_particle(self, particle_id):
"""
returns a specific MC particle identified by its unique id
"""
if particle_id not in self.__particles:
raise AttributeError(f"MC particle with id {particle_id} not present")
return self.__particles[particle_id]
[docs] def get_primary(self):
"""
returns a first MC particle
"""
if len(self.__particles) == 0:
return None
return self.get_particle(0)
[docs] def get_parent(self, particle_or_shower):
"""
returns the parent of a particle or a shower
"""
if isinstance(particle_or_shower, NuRadioReco.framework.base_shower.BaseShower):
par_id = particle_or_shower[parameters.showerParameters.parent_id]
elif isinstance(particle_or_shower, NuRadioReco.framework.particle.Particle):
par_id = particle_or_shower[parameters.particleParameters.parent_id]
else:
raise ValueError("particle_or_shower needs to be an instance of NuRadioReco.framework.base_shower.BaseShower or NuRadioReco.framework.particle.Particle")
if par_id is None:
logger.info("did not find parent for {particle_or_shower}")
return None
return self.get_particle(par_id)
[docs] def has_particle(self, particle_id=None):
"""
Returns true if at least one MC particle is stored in the event
If particle_id is given, it checks if this particular MC particle exists
"""
if particle_id is None:
return len(self.__particles) > 0
return particle_id in self.__particles.keys()
[docs] def get_interaction_products(self, parent_particle, showers=True, particles=True):
"""
Return all the daughter particles and showers generated in the interaction of the <parent_particle>
Parameters
----------
showers: bool
Include simulated showers in the list
showers: bool
Include simulated particles in the list
"""
parent_id = parent_particle.get_id()
# iterate over sim_showers to look for parent id
if showers is True:
for shower in self.get_showers():
if shower[parameters.showerParameters.parent_id] == parent_id:
yield shower
# iterate over secondary particles to look for parent id
if particles is True:
for particle in self.get_particles():
if particle[parameters.particleParameters.parent_id] == parent_id:
yield particle
[docs] def add_shower(self, shower):
"""
Adds a radio shower to the event
Parameters
----------
shower: RadioShower object
The shower to be added to the event
"""
if(shower.get_id() in self.__radio_showers):
logger.error("shower with id {shower.get_id()} already exists. Shower id needs to be unique per event")
raise AttributeError("shower with id {shower.get_id()} already exists. Shower id needs to be unique per event")
self.__radio_showers[shower.get_id()] = shower
[docs] def get_showers(self, ids=None):
"""
Returns an iterator over the showers stored in the event
Parameters
----------
ids: list of integers
A list of station IDs. Only showers that are associated with
all stations in the list are returned
"""
for shower in self.__radio_showers.values():
if ids is None:
yield shower
elif shower.has_station_ids(ids):
yield shower
[docs] def get_shower(self, shower_id):
"""
returns a specific shower identified by its unique id
"""
if(shower_id not in self.__radio_showers):
raise AttributeError(f"shower with id {shower_id} not present")
return self.__radio_showers[shower_id]
[docs] def has_shower(self, shower_id=None):
"""
Returns true if at least one shower is stored in the event
If shower_id is given, it checks if this particular shower exists
"""
if(shower_id is None):
return len(self.__radio_showers) > 0
else:
return shower_id in self.__radio_showers.keys()
[docs] def get_first_shower(self, ids=None):
"""
Returns only the first shower stored in the event. Useful in cases
when there is only one shower in the event.
Parameters
----------
ids: list of integers
A list of station IDs. The first shower that is associated with
all stations in the list is returned
"""
if len(self.__radio_showers) == 0:
return None
if ids is None:
shower_ids = list(self.__radio_showers.keys())
return self.__radio_showers[shower_ids[0]]
for shower in self.__radio_showers:
if shower.has_station_ids(ids):
return shower
return None
[docs] def add_sim_shower(self, sim_shower):
"""
Add a simulated shower to the event
Parameters
----------
sim_shower: RadioShower object
The shower to be added to the event
"""
if not isinstance(sim_shower, NuRadioReco.framework.radio_shower.RadioShower):
raise AttributeError("sim_shower needs to be of type NuRadioReco.framework.radio_shower.RadioShower")
if(sim_shower.get_id() in self.__sim_showers):
logger.error(f"sim shower with id {sim_shower.get_id()} already exists. Shower id needs to be unique per event")
raise AttributeError(f"sim shower with id {sim_shower.get_id()} already exists. Shower id needs to be unique per event")
self.__sim_showers[sim_shower.get_id()] = sim_shower
[docs] def get_sim_showers(self):
"""
Get an iterator over all simulated showers in the event
"""
for shower in self.__sim_showers.values():
yield shower
[docs] def get_sim_shower(self, shower_id):
"""
returns a specific shower identified by its unique id
"""
if(shower_id not in self.__sim_showers):
raise AttributeError(f"sim shower with id {shower_id} not present")
return self.__sim_showers[shower_id]
[docs] def get_first_sim_shower(self, ids=None):
"""
Returns only the first sim shower stored in the event. Useful in cases
when there is only one shower in the event.
Parameters
----------
ids: list of integers
A list of station IDs. The first shower that is associated with
all stations in the list is returned
"""
if len(self.__sim_showers) == 0:
return None
if ids is None:
shower_ids = list(self.__sim_showers.keys())
return self.__sim_showers[shower_ids[0]]
for shower in self.__sim_showers:
if shower.has_station_ids(ids):
return shower
return None
[docs] def has_sim_shower(self, shower_id=None):
"""
Returns true if at least one simulated shower is stored in the event
If shower_id is given, it checks if this particular shower exists
"""
if(shower_id is None):
return shower_id in self.__sim_showers.keys()
else:
return len(self.__sim_showers) > 0
[docs] def add_sim_emitter(self, sim_emitter):
"""
Add a simulated emitter to the event
Parameters
----------
sim_emitter: SimEmitter object
The emitter to be added to the event
"""
if not isinstance(sim_emitter, NuRadioReco.framework.sim_emitter.SimEmitter):
raise AttributeError(f"emitter needs to be of type NuRadioReco.framework.sim_emitter.SimEmitter but is of type {type(sim_emitter)}")
if(sim_emitter.get_id() in self.__sim_emitters):
logger.error(f"sim emitter with id {sim_emitter.get_id()} already exists. Emitter id needs to be unique per event")
raise AttributeError(f"sim emitter with id {sim_emitter.get_id()} already exists. Emitter id needs to be unique per event")
self.__sim_emitters[sim_emitter.get_id()] = sim_emitter
[docs] def get_sim_emitters(self):
"""
Get an iterator over all simulated emitters in the event
"""
for emitter in self.__sim_emitters.values():
yield emitter
[docs] def get_sim_emitter(self, emitter_id):
"""
returns a specific emitter identified by its unique id
"""
if(emitter_id not in self.__sim_emitters):
raise AttributeError(f"sim emitter with id {emitter_id} not present")
return self.__sim_emitters[emitter_id]
[docs] def get_first_sim_emitter(self, ids=None):
"""
Returns only the first sim emitter stored in the event. Useful in cases
when there is only one emitter in the event.
Parameters
----------
station_ids: list of integers
A list of station IDs. The first emitter that is associated with
all stations in the list is returned
"""
if len(self.__sim_emitters) == 0:
return None
if ids is None:
emitter_ids = list(self.__sim_emitters.keys())
return self.__sim_emitters[emitter_ids[0]]
for emitter in self.__sim_emitters:
if emitter.has_station_ids(ids):
return emitter
return None
[docs] def has_sim_emitter(self, emitter_id=None):
"""
Returns true if at least one simulated emitter is stored in the event
If emitter_id is given, it checks if this particular emitter exists
"""
if(emitter_id is None):
return emitter_id in self.__sim_emitters.keys()
else:
return len(self.__sim_emitters) > 0
[docs] def serialize(self, mode):
stations_pkl = []
try:
commit_hash = NuRadioReco.utilities.version.get_NuRadioMC_commit_hash()
self.set_parameter(parameters.eventParameters.hash_NuRadioMC, commit_hash)
except:
logger.warning("Event is serialized without commit hash!")
self.set_parameter(parameters.eventParameters.hash_NuRadioMC, None)
for station in self.get_stations():
stations_pkl.append(station.serialize(mode))
showers_pkl = [shower.serialize() for shower in self.get_showers()]
sim_showers_pkl = [shower.serialize() for shower in self.get_sim_showers()]
sim_emitters_pkl = [emitter.serialize() for emitter in self.get_sim_emitters()]
particles_pkl = [particle.serialize() for particle in self.get_particles()]
hybrid_info = self.__hybrid_information.serialize()
modules_out_event = []
for value in self.__modules_event: # remove module instances (this will just blow up the file size)
modules_out_event.append([value[0], None, value[2]])
invalid_keys = [key for key,val in value[2].items() if isinstance(val, BaseException)]
if len(invalid_keys):
logger.warning(f"The following arguments to module {value[0]} could not be serialized and will not be stored: {invalid_keys}")
modules_out_station = {}
for key in self.__modules_station: # remove module instances (this will just blow up the file size)
modules_out_station[key] = []
for value in self.__modules_station[key]:
modules_out_station[key].append([value[0], value[1], None, value[3]])
invalid_keys = [key for key,val in value[3].items() if isinstance(val, BaseException)]
if len(invalid_keys):
logger.warning(f"The following arguments to module {value[0]} could not be serialized and will not be stored: {invalid_keys}")
data = {'_parameters': self._parameters,
'__run_number': self.__run_number,
'_id': self._id,
'__event_time': self.__event_time,
'stations': stations_pkl,
'showers': showers_pkl,
'sim_showers': sim_showers_pkl,
'sim_emitters': sim_emitters_pkl,
'particles': particles_pkl,
'hybrid_info': hybrid_info,
'generator_info': self._generator_info,
'__modules_event': modules_out_event,
'__modules_station': modules_out_station
}
return pickle.dumps(data, protocol=4)
[docs] def deserialize(self, data_pkl):
data = pickle.loads(data_pkl)
for station_pkl in data['stations']:
station = NuRadioReco.framework.station.Station(0)
station.deserialize(station_pkl)
self.set_station(station)
if 'showers' in data.keys():
for shower_pkl in data['showers']:
shower = NuRadioReco.framework.radio_shower.RadioShower(None)
shower.deserialize(shower_pkl)
self.add_shower(shower)
if 'sim_showers' in data.keys():
for shower_pkl in data['sim_showers']:
shower = NuRadioReco.framework.radio_shower.RadioShower(None)
shower.deserialize(shower_pkl)
self.add_sim_shower(shower)
if 'sim_emitters' in data.keys():
for emmitter_pkl in data['sim_emitters']:
emitter = NuRadioReco.framework.sim_emitter.SimEmitter(None)
emitter.deserialize(emmitter_pkl)
self.add_sim_emitter(emitter)
if 'particles' in data.keys():
for particle_pkl in data['particles']:
particle = NuRadioReco.framework.particle.Particle(None)
particle.deserialize(particle_pkl)
self.add_particle(particle)
self.__hybrid_information = NuRadioReco.framework.hybrid_information.HybridInformation()
if 'hybrid_info' in data.keys():
self.__hybrid_information.deserialize(data['hybrid_info'])
self._parameters = data['_parameters']
self.__run_number = data['__run_number']
self._id = data['_id']
self.__event_time = data['__event_time']
if 'generator_info' in data.keys():
self._generator_info = data['generator_info']
if "__modules_event" in data:
self.__modules_event = data['__modules_event']
if "__modules_station" in data:
self.__modules_station = data['__modules_station']