from __future__ import absolute_import, division, print_function
import NuRadioReco.framework.base_station
import NuRadioReco.framework.sim_station
import NuRadioReco.framework.channel
from six import iteritems
import pickle
import logging
import collections
logger = logging.getLogger('Station')
[docs]class Station(NuRadioReco.framework.base_station.BaseStation):
def __init__(self, station_id):
NuRadioReco.framework.base_station.BaseStation.__init__(self, station_id)
self.__channels = collections.OrderedDict()
self.__reference_reconstruction = 'RD'
self.__sim_station = None
[docs] def set_sim_station(self, sim_station):
self.__sim_station = sim_station
[docs] def get_sim_station(self):
return self.__sim_station
[docs] def has_sim_station(self):
return self.__sim_station is not None
[docs] def iter_channels(self, use_channels=None):
for channel_id, channel in iteritems(self.__channels):
if(use_channels is None):
yield channel
else:
if channel_id in use_channels:
yield channel
[docs] def get_channel(self, channel_id):
return self.__channels[channel_id]
[docs] def iter_channel_group(self, channel_group_id):
found_channel_group = False
for channel_id, channel in iteritems(self.__channels):
if(channel.get_group_id() == channel_group_id):
found_channel_group = True
yield channel
if found_channel_group == False:
msg = f"channel group id {channel_group_id} is not present"
raise ValueError(msg)
[docs] def get_number_of_channels(self):
return len(self.__channels)
[docs] def get_channel_ids(self):
return list(self.__channels.keys())
[docs] def add_channel(self, channel):
self.__channels[channel.get_id()] = channel
[docs] def remove_channel(self, channel_id):
"""
Removes a channel from the station by deleting is from the channels dictionary. The `channel_id`
should be the id of the Channel, but supplying the Channel object itself is also supported.
Parameters
----------
channel_id : int or NuRadioReco.framework.channel.Channel
The Channel (id) to remove from the Station.
"""
if isinstance(channel_id, NuRadioReco.framework.channel.Channel):
del self.__channels[channel_id.get_id()]
else:
del self.__channels[channel_id]
[docs] def set_reference_reconstruction(self, reference):
if reference not in ['RD', 'MC']:
import sys
logger.error("reference reconstructions other than RD and MC are not supported")
sys.exit(-1)
self.__reference_reconstruction = reference
[docs] def get_reference_reconstruction(self):
return self.__reference_reconstruction
[docs] def get_reference_direction(self):
if(self.__reference_reconstruction == 'RD'):
return self.get_parameter('zenith'), self.get_parameter('azimuth')
if(self.__reference_reconstruction == 'MC'):
return self.get_sim_station().get_parameter('zenith'), self.get_sim_station().get_parameter('azimuth')
[docs] def get_magnetic_field_vector(self, time=None):
if(self.__reference_reconstruction == 'MC'):
return self.get_sim_station().get_magnetic_field_vector()
if(self.__reference_reconstruction == 'RD'):
if time is not None:
logger.warning("time dependent magnetic field model not yet implemented, returning static magnetic field for the ARIANNA site")
from radiotools import helper
return helper.get_magnetic_field_vector('arianna')
[docs] def serialize(self, mode):
save_efield_traces = 'ElectricFields' in mode and mode['ElectricFields'] is True
base_station_pkl = NuRadioReco.framework.base_station.BaseStation.serialize(self, save_efield_traces=save_efield_traces)
channels_pkl = []
save_channel_trace = 'Channels' in mode and mode['Channels'] is True
for channel in self.iter_channels():
channels_pkl.append(channel.serialize(save_channel_trace))
save_sim_channel_trace = 'SimChannels' in mode and mode['SimChannels'] is True
save_sim_efield_trace = 'SimElectricFields' in mode and mode['SimElectricFields'] is True
sim_station_pkl = None
if(self.has_sim_station()):
sim_station_pkl = self.get_sim_station().serialize(save_channel_traces=save_sim_channel_trace,
save_efield_traces=save_sim_efield_trace)
data = {'__reference_reconstruction': self.__reference_reconstruction,
'channels': channels_pkl,
'base_station': base_station_pkl,
'sim_station': sim_station_pkl}
return pickle.dumps(data, protocol=4)
[docs] def deserialize(self, data_pkl):
data = pickle.loads(data_pkl)
NuRadioReco.framework.base_station.BaseStation.deserialize(self, data['base_station'])
if(data['sim_station'] is None):
self.__sim_station = None
else:
self.__sim_station = NuRadioReco.framework.sim_station.SimStation(None)
self.__sim_station.deserialize(data['sim_station'])
for channel_pkl in data['channels']:
channel = NuRadioReco.framework.channel.Channel(0)
channel.deserialize(channel_pkl)
self.add_channel(channel)
self.__reference_reconstruction = data['__reference_reconstruction']