import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt
from NuRadioReco.utilities import units
import NuRadioReco.utilities.io_utilities
import NuRadioReco.framework.electric_field
import NuRadioReco.detector.antennapattern
from NuRadioReco.framework.parameters import stationParameters as stnp
from NuRadioReco.framework.parameters import showerParameters as shp
from NuRadioReco.utilities import trace_utilities, fft, bandpass_filter
import radiotools.helper as hp
[docs]class neutrino3DVertexReconstructor:
def __init__(self, lookup_table_location):
"""
Constructor for the vertex reconstructor
Parameters
----------
lookup_table_location: string
path to the folder in which the lookup tables for the signal travel
times are stored
"""
self.__lookup_table_location = lookup_table_location
self.__detector = None
self.__lookup_table = {}
self.__header = {}
self.__channel_ids = None
self.__station_id = None
self.__channel_pairs = None
self.__rec_x = None
self.__rec_z = None
self.__sampling_rate = None
self.__passband = None
self.__channel_pair = None
self.__channel_positions = None
self.__correlation = None
self.__max_corr_index = None
self.__current_ray_types = None
self.__electric_field_template = None
self.__azimuths_2d = None
self.__distances_2d = None
self.__z_coordinates_2d = None
self.__distance_step_3d = None
self.__z_step_3d = None
self.__widths_3d = None
self.__heights_3d = None
self.__debug_folder = None
self.__current_distance = None
self.__pair_correlations = None
self.__self_correlations = None
self.__antenna_pattern_provider = NuRadioReco.detector.antennapattern.AntennaPatternProvider()
self.__ray_types = [
['direct', 'direct'],
['reflected', 'reflected'],
['refracted', 'refracted'],
['direct', 'reflected'],
['reflected', 'direct'],
['direct', 'refracted'],
['refracted', 'direct'],
['reflected', 'refracted'],
['refracted', 'reflected']
]
[docs] def begin(
self,
station_id,
channel_ids,
detector,
template,
distances_2d=None,
azimuths_2d=None,
z_coordinates_2d=None,
distance_step_3d=2 * units.m,
widths_3d=None,
z_step_3d=2 * units.m,
passband=None,
min_antenna_distance=5. * units.m,
debug_folder='.'
):
"""
General settings for vertex reconstruction
Parameters
----------
station_id: integer
ID of the station to be used for the reconstruction
channel_ids: array of integers
IDs of the channels to be used for the reconstruction
detector: Detector or GenericDetector
Detector description for the detector used in the reconstruction
template: BaseTrace object or object of child class
An electric field to be used to calculate the template which are correlated
with the channel traces in order to determine the timing difference between
channels
distances_2d: array of float
A list of horizontal distances from the center of the station at which the first
rough scan to determine the search volume is done. The minimum and maximum of this
list is later also used as the minimum and maximum distance for the finer search
azimuths_2d: array of float
Array of azimuths to be used in the first scan to determine the search volume
z_coordinates_2d: array of float
Array of the z coordinates relative to the surface to be used in the first scan
to determine the search volume. The maximum depth is also used as the maximum
depth for the finer search
distance_step_3d: float
Step size for the horizontal distances used in the finer scan
widths_3d: array of float
List of distances to the left and right of the line determined in the first rough
scan on which the finer scan should be performed
z_step_3d: float
Step size for the depts used in the finer scan
passband: array of float
Lower and upper bounds off the bandpass filter that is applied to the channel
waveform and the template before the correlations are determined. This filter
does not affect the voltages stored in the channels.
min_antenna_distance: float
Minimum distance two antennas need to have to be used as a pair in the reconstruction
debug_folder: string
Path to the folder in which debug plots should be saved if the debug=True option
is picked in the run() method.
"""
self.__detector = detector
self.__channel_ids = channel_ids
self.__station_id = station_id
self.__debug_folder = debug_folder
self.__channel_pairs = []
for i in range(len(channel_ids) - 1):
for j in range(i + 1, len(channel_ids)):
relative_positions = detector.get_relative_position(station_id, channel_ids[i]) - detector.get_relative_position(station_id, channel_ids[j])
if detector.get_antenna_type(station_id, channel_ids[i]) == detector.get_antenna_type(station_id, channel_ids[j]) \
and np.sqrt(np.sum(relative_positions**2)) > min_antenna_distance:
self.__channel_pairs.append([channel_ids[i], channel_ids[j]])
self.__lookup_table = {}
self.__header = {}
self.__electric_field_template = template
self.__sampling_rate = template.get_sampling_rate()
self.__passband = passband
if distances_2d is None:
self.__distances_2d = np.arange(100, 3000, 200)
else:
self.__distances_2d = distances_2d
if azimuths_2d is None:
self.__azimuths_2d = np.arange(0, 360.1, 2.5) * units.deg
else:
self.__azimuths_2d = azimuths_2d
if z_coordinates_2d is None:
self.__z_coordinates_2d = np.arange(-2700, -100, 25)
else:
self.__z_coordinates_2d = z_coordinates_2d
self.__distance_step_3d = distance_step_3d
if widths_3d is None:
self.__widths_3d = np.arange(-50, 50, 2.)
else:
self.__widths_3d = widths_3d
self.__z_step_3d = z_step_3d
for channel_id in channel_ids:
channel_z = abs(detector.get_relative_position(station_id, channel_id)[2])
if channel_z not in self.__lookup_table.keys():
f = NuRadioReco.utilities.io_utilities.read_pickle('{}/lookup_table_{}.p'.format(self.__lookup_table_location, int(abs(channel_z))))
self.__header[int(channel_z)] = f['header']
self.__lookup_table[int(abs(channel_z))] = f['antenna_{}'.format(channel_z)]
[docs] def run(
self,
event,
station,
det,
debug=False
):
azimuth_grid_2d, z_grid_2d = np.meshgrid(self.__azimuths_2d, self.__z_coordinates_2d)
distance_correlations = np.zeros(self.__distances_2d.shape)
full_correlations = np.zeros((len(self.__distances_2d), len(self.__z_coordinates_2d), len(self.__azimuths_2d)))
if debug:
plt.close('all')
fig1 = plt.figure(figsize=(12, (len(self.__channel_pairs) // 2 + len(self.__channel_pairs) % 2)))
self.__pair_correlations = np.zeros((len(self.__channel_pairs), station.get_channel(self.__channel_ids[0]).get_number_of_samples() + self.__electric_field_template.get_number_of_samples() - 1))
for i_pair, channel_pair in enumerate(self.__channel_pairs):
channel_1 = station.get_channel(channel_pair[0])
channel_2 = station.get_channel(channel_pair[1])
antenna_response = trace_utilities.get_efield_antenna_factor(
station=station,
frequencies=self.__electric_field_template.get_frequencies(),
channels=[channel_pair[0]],
detector=det,
zenith=90. * units.deg,
azimuth=0,
antenna_pattern_provider=self.__antenna_pattern_provider
)[0]
voltage_spec = (
antenna_response[0] * self.__electric_field_template.get_frequency_spectrum() + antenna_response[1] * self.__electric_field_template.get_frequency_spectrum()
) * det.get_amplifier_response(station.get_id(), channel_pair[0], self.__electric_field_template.get_frequencies())
if self.__passband is not None:
voltage_spec *= bandpass_filter.get_filter_response(self.__electric_field_template.get_frequencies(), self.__passband, 'butterabs', 10)
voltage_template = fft.freq2time(voltage_spec, self.__sampling_rate)
voltage_template /= np.max(np.abs(voltage_template))
if self.__passband is None:
corr_1 = np.abs(hp.get_normalized_xcorr(channel_1.get_trace(), voltage_template))
corr_2 = np.abs(hp.get_normalized_xcorr(channel_2.get_trace(), voltage_template))
else:
corr_1 = np.abs(hp.get_normalized_xcorr(channel_1.get_filtered_trace(self.__passband, 'butterabs', 10), voltage_template))
corr_2 = np.abs(hp.get_normalized_xcorr(channel_2.get_filtered_trace(self.__passband, 'butterabs', 10), voltage_template))
correlation_product = np.zeros_like(corr_1)
sample_shifts = np.arange(-len(corr_1) // 2, len(corr_1) // 2, dtype=int)
toffset = sample_shifts / channel_1.get_sampling_rate()
for i_shift, shift_sample in enumerate(sample_shifts):
correlation_product[i_shift] = np.max((corr_1 * np.roll(corr_2, shift_sample)))
self.__pair_correlations[i_pair] = correlation_product
if debug:
ax1_1 = fig1.add_subplot(len(self.__channel_pairs) // 2 + len(self.__channel_pairs) % 2, 2,
i_pair + 1)
ax1_1.grid()
ax1_1.plot(toffset, correlation_product)
if debug:
fig1.tight_layout()
fig1.savefig('{}/{}_{}_correlation.png'.format(self.__debug_folder, event.get_run_number(), event.get_id()))
for i_dist, distance in enumerate(self.__distances_2d):
self.__current_distance = distance
correlation_sum = np.zeros_like(azimuth_grid_2d)
for i_pair, channel_pair in enumerate(self.__channel_pairs):
self.__correlation = self.__pair_correlations[i_pair]
self.__channel_pair = channel_pair
self.__channel_positions = [self.__detector.get_relative_position(self.__station_id, channel_pair[0]),
self.__detector.get_relative_position(self.__station_id, channel_pair[1])]
correlation_map = np.zeros_like(correlation_sum)
for i_ray in range(len(self.__ray_types)):
self.__current_ray_types = self.__ray_types[i_ray]
correlation_map = np.maximum(self.get_correlation_array_2d(azimuth_grid_2d, z_grid_2d), correlation_map)
correlation_sum += correlation_map
distance_correlations[i_dist] = np.max(correlation_sum)
full_correlations[i_dist] = correlation_sum
corr_fit_threshold = .7 * np.max(full_correlations)
flattened_corr = np.max(full_correlations, axis=2).T
i_max_d = np.argmax(flattened_corr, axis=0)
corr_mask_d = np.max(flattened_corr, axis=0) > corr_fit_threshold
line_fit_d = np.polyfit(
self.__distances_2d[corr_mask_d],
self.__z_coordinates_2d[i_max_d][corr_mask_d],
1
)
residuals_d = np.sum((self.__z_coordinates_2d[i_max_d][corr_mask_d] - self.__distances_2d[corr_mask_d] * line_fit_d[0] - line_fit_d[1])**2) / np.sum(corr_mask_d.astype(int))
i_max_z = np.argmax(flattened_corr, axis=1)
corr_mask_z = np.max(flattened_corr, axis=1) > corr_fit_threshold
line_fit_z = np.polyfit(
self.__distances_2d[i_max_z][corr_mask_z],
self.__z_coordinates_2d[corr_mask_z],
1
)
residuals_z = np.sum((self.__z_coordinates_2d[corr_mask_z] - self.__distances_2d[i_max_z][corr_mask_z] * line_fit_z[0] - line_fit_z[1])**2) / np.sum(corr_mask_z.astype(int))
if residuals_d <= residuals_z:
slope = line_fit_d[0]
offset = line_fit_d[1]
max_z_offset = np.max([50, np.min([200, np.max(self.__z_coordinates_2d[i_max_d][corr_mask_d] - self.__distances_2d[corr_mask_d] * slope - offset)])])
min_z_offset = np.max([50, np.min([200, np.max(-self.__z_coordinates_2d[i_max_d][corr_mask_d] + self.__distances_2d[corr_mask_d] * slope + offset)])])
flattened_corr_theta = np.max(full_correlations, axis=1)
theta_corr_mask = np.max(flattened_corr_theta, axis=1) >= corr_fit_threshold
i_max_theta = np.argmax(flattened_corr_theta, axis=1)
median_theta = np.median(self.__azimuths_2d[i_max_theta][theta_corr_mask])
z_fit = False
else:
slope = line_fit_z[0]
offset = line_fit_z[1]
max_z_offset = np.max([50, np.min([200, np.max(self.__z_coordinates_2d[corr_mask_z] - self.__distances_2d[i_max_z][corr_mask_z] * slope - offset)])])
min_z_offset = np.max([50, np.min([200, np.max(-self.__z_coordinates_2d[corr_mask_z] + self.__distances_2d[i_max_z][corr_mask_z] * slope + offset)])])
flattened_corr_theta = np.max(full_correlations, axis=0)
theta_corr_mask = np.max(flattened_corr_theta, axis=1) >= corr_fit_threshold
i_max_theta = np.argmax(flattened_corr_theta, axis=1)
median_theta = np.median(self.__azimuths_2d[i_max_theta][theta_corr_mask])
z_fit = True
if debug:
self.__draw_2d_correlation_map(event, full_correlations)
self.__draw_search_zones(
event,
slope,
offset,
line_fit_d,
line_fit_z,
min_z_offset,
max_z_offset,
i_max_d,
i_max_z,
corr_mask_d,
corr_mask_z,
z_fit,
i_max_theta,
theta_corr_mask,
median_theta
)
# <--- 3D Fit ---> #
distances_3d = np.arange(self.__distances_2d[0], self.__distances_2d[-1], self.__distance_step_3d)
z_coords = slope * distances_3d + offset
distances_3d = distances_3d[(z_coords < 0) & (z_coords > -2700)]
search_heights = np.arange(-1.1 * min_z_offset, 1.1 * max_z_offset, self.__z_step_3d)
x_0, y_0, z_0 = np.meshgrid(distances_3d, self.__widths_3d, search_heights)
z_coords = z_0 + slope * x_0 + offset
x_coords = np.cos(median_theta) * x_0 - y_0 * np.sin(median_theta)
y_coords = np.sin(median_theta) * x_0 + y_0 * np.cos(median_theta)
correlation_sum = np.zeros_like(z_coords)
for i_pair, channel_pair in enumerate(self.__channel_pairs):
self.__correlation = self.__pair_correlations[i_pair]
self.__channel_pair = channel_pair
self.__channel_positions = [self.__detector.get_relative_position(self.__station_id, channel_pair[0]),
self.__detector.get_relative_position(self.__station_id, channel_pair[1])]
correlation_map = np.zeros_like(correlation_sum)
for i_ray in range(len(self.__ray_types)):
self.__current_ray_types = self.__ray_types[i_ray]
correlation_map = np.maximum(self.get_correlation_array_3d(x_coords, y_coords, z_coords), correlation_map)
correlation_sum += correlation_map
i_max = np.unravel_index(np.argmax(correlation_sum), correlation_sum.shape)
if debug:
self.__draw_vertex_reco(
event,
correlation_sum / np.max(correlation_sum),
x_0,
y_0,
z_0,
x_coords,
y_coords,
z_coords,
slope,
offset,
median_theta,
i_max
)
# <<--- DnR Reco --->> #
self.__self_correlations = np.zeros((len(self.__channel_ids), station.get_channel(self.__channel_ids[0]).get_number_of_samples() + self.__electric_field_template.get_number_of_samples() - 1))
self_correlation_sum = np.zeros_like(z_coords)
for i_channel, channel_id in enumerate(self.__channel_ids):
channel = station.get_channel(channel_id)
antenna_response = trace_utilities.get_efield_antenna_factor(
station=station,
frequencies=self.__electric_field_template.get_frequencies(),
channels=[channel_id],
detector=det,
zenith=90. * units.deg,
azimuth=0,
antenna_pattern_provider=self.__antenna_pattern_provider
)[0]
voltage_spec = (
antenna_response[0] * self.__electric_field_template.get_frequency_spectrum() + antenna_response[1] * self.__electric_field_template.get_frequency_spectrum()
) * det.get_amplifier_response(station.get_id(), channel_id, self.__electric_field_template.get_frequencies())
if self.__passband is not None:
voltage_spec *= bandpass_filter.get_filter_response(self.__electric_field_template.get_frequencies(), self.__passband, 'butter', 10)
voltage_template = fft.freq2time(voltage_spec, self.__sampling_rate)
voltage_template /= np.max(np.abs(voltage_template))
if self.__passband is None:
corr_1 = hp.get_normalized_xcorr(channel.get_trace(), voltage_template)
corr_2 = hp.get_normalized_xcorr(channel.get_trace(), voltage_template)
else:
corr_1 = np.abs(hp.get_normalized_xcorr(channel.get_filtered_trace(self.__passband, 'butter', 10), voltage_template))
corr_2 = np.abs(hp.get_normalized_xcorr(channel.get_filtered_trace(self.__passband, 'butter', 10), voltage_template))
correlation_product = np.zeros_like(corr_1)
sample_shifts = np.arange(-len(corr_1) // 2, len(corr_1) // 2, dtype=int)
toffset = sample_shifts / channel.get_sampling_rate()
for i_shift, shift_sample in enumerate(sample_shifts):
correlation_product[i_shift] = np.max((corr_1 * np.roll(corr_2, shift_sample)))
correlation_product = np.abs(correlation_product)
correlation_product[np.abs(toffset) < 20] = 0
self.__correlation = correlation_product
self.__channel_pair = [channel_id, channel_id]
self.__channel_positions = [self.__detector.get_relative_position(self.__station_id, channel_id),
self.__detector.get_relative_position(self.__station_id, channel_id)]
correlation_map = np.zeros_like(correlation_sum)
for i_ray in range(len(self.__ray_types)):
if self.__ray_types[i_ray][0] != self.__ray_types[i_ray][1]:
self.__current_ray_types = self.__ray_types[i_ray]
correlation_map = np.maximum(self.get_correlation_array_3d(x_coords, y_coords, z_coords), correlation_map)
self_correlation_sum += correlation_map
combined_correlations = correlation_sum / len(self.__channel_pairs) + self_correlation_sum / len(self.__channel_ids)
i_max_dnr = np.unravel_index(np.argmax(combined_correlations), combined_correlations.shape)
vertex_x = x_coords[i_max_dnr]
vertex_y = y_coords[i_max_dnr]
vertex_z = z_coords[i_max_dnr]
station.set_parameter(stnp.nu_vertex, np.array([vertex_x, vertex_y, vertex_z]))
if debug:
self.__draw_dnr_reco(
event,
correlation_sum / np.max(correlation_sum),
self_correlation_sum / np.max(self_correlation_sum),
combined_correlations / np.max(combined_correlations),
x_0,
y_0,
z_0,
slope,
offset,
median_theta,
i_max,
i_max_dnr
)
[docs] def get_correlation_array_2d(self, phi, z):
"""
Returns the correlations corresponding to the different
signal travel times between channels for the given positions.
This is done by correcting for the distance of the channels
from the station center and then calling
self.get_correlation_for_pos, which does the actual work.
Parameters
----------
x, z: array
Coordinates of the points for which calculations are
to be calculated. Correspond to the (r, z) pair
of cylindrical coordinates.
"""
channel_pos1 = self.__channel_positions[0]
channel_pos2 = self.__channel_positions[1]
x = self.__current_distance * np.cos(phi)
y = self.__current_distance * np.sin(phi)
d_hor1 = np.sqrt((x - channel_pos1[0])**2 + (y - channel_pos1[1])**2)
d_hor2 = np.sqrt((x - channel_pos2[0])**2 + (y - channel_pos2[1])**2)
res = self.get_correlation_for_pos(np.array([d_hor1, d_hor2]), z)
return res
[docs] def get_correlation_array_3d(self, x, y, z):
channel_pos1 = self.__channel_positions[0]
channel_pos2 = self.__channel_positions[1]
d_hor1 = np.sqrt((x - channel_pos1[0])**2 + (y - channel_pos1[1])**2)
d_hor2 = np.sqrt((x - channel_pos2[0])**2 + (y - channel_pos2[1])**2)
res = self.get_correlation_for_pos(np.array([d_hor1, d_hor2]), z)
res[np.abs(res) < .8 * np.max(np.abs(res))] = 0
return res
[docs] def get_correlation_for_pos(self, d_hor, z):
"""
Returns the correlations corresponding to the different
signal travel times between channels for the given positions.
Parameters
----------
d_hor, z: array
Coordinates of the points for which calculations are
to be calculated. Correspond to the (r, z) pair
of cylindrical coordinates.
"""
t1 = self.get_signal_travel_time(d_hor[0], z, self.__current_ray_types[0], self.__channel_pair[0])
t2 = self.get_signal_travel_time(d_hor[1], z, self.__current_ray_types[1], self.__channel_pair[1])
delta_t = t1 - t2
delta_t = delta_t.astype(float)
corr_index = self.__correlation.shape[0] / 2 + np.round(delta_t * self.__sampling_rate)
corr_index[np.isnan(delta_t)] = 0
mask = (corr_index > 0) & (corr_index < self.__correlation.shape[0]) & (~np.isinf(delta_t))
corr_index[~mask] = 0
res = np.take(self.__correlation, corr_index.astype(int))
res[~mask] = 0
return res
[docs] def get_signal_travel_time(self, d_hor, z, ray_type, channel_id):
"""
Calculate the signal travel time between a position and the
channel
Parameters
----------
d_hor, z: numbers or arrays of numbers
Coordinates of the point from which to calculate the
signal travel times. Correspond to (r, z) coordinates
in cylindrical coordinates.
ray_type: string
Ray type for which to calculate the travel times. Options
are direct, reflected and refracted
channel_id: int
ID of the channel to which the travel time shall be calculated
"""
channel_pos = self.__detector.get_relative_position(self.__station_id, channel_id)
channel_type = int(abs(channel_pos[2]))
travel_times = np.zeros_like(d_hor)
mask = np.ones_like(travel_times).astype(bool)
i_z = np.array(np.round((z - self.__header[channel_type]['z_min']) / self.__header[channel_type]['d_z'])).astype(int)
i_x_1 = np.array(np.floor((d_hor - self.__header[channel_type]['x_min']) / self.__header[channel_type]['d_x'])).astype(int)
cell_dist_1 = i_x_1 * self.__header[channel_type]['d_x'] + self.__header[channel_type]['x_min']
mask[i_x_1 > self.__lookup_table[channel_type][ray_type].shape[0] - 1] = False
mask[i_z > self.__lookup_table[channel_type][ray_type].shape[1] - 1] = False
i_x_1[~mask] = 0
i_z[~mask] = 0
travel_times_1 = self.__lookup_table[channel_type][ray_type][(i_x_1, i_z)]
i_x_2 = np.array(np.ceil((d_hor - self.__header[channel_type]['x_min']) / self.__header[channel_type]['d_x'])).astype(int)
cell_dist_2 = i_x_2 * self.__header[channel_type]['d_x'] + self.__header[channel_type]['x_min']
i_x_2[~mask] = 0
travel_times_2 = self.__lookup_table[channel_type][ray_type][(i_x_2, i_z)]
slopes = np.zeros_like(travel_times_1)
slopes[i_x_2 > i_x_1] = (travel_times_1 - travel_times_2)[i_x_2 > i_x_1] / (cell_dist_1 - cell_dist_2)[i_x_2 > i_x_1]
travel_times = (d_hor - cell_dist_1) * slopes + travel_times_1
travel_times[~mask] = np.nan
return travel_times
def __draw_2d_correlation_map(self, event, correlation_map):
fig4 = plt.figure(figsize=(4, 12))
ax4_1 = fig4.add_subplot(311)
d_0, z_0 = np.meshgrid(self.__distances_2d, self.__z_coordinates_2d)
ax4_1.pcolor(
d_0,
z_0,
np.max(correlation_map, axis=2).T
)
ax4_1.grid()
theta_0, d_0 = np.meshgrid(self.__azimuths_2d, self.__distances_2d)
ax4_2 = fig4.add_subplot(312, projection='polar')
ax4_2.pcolor(
theta_0,
d_0,
np.max(correlation_map, axis=1)
)
ax4_2.grid()
ax4_3 = fig4.add_subplot(313)
theta_0, z_0 = np.meshgrid(self.__azimuths_2d, self.__z_coordinates_2d)
ax4_3.pcolor(
theta_0 / units.deg,
z_0,
np.max(correlation_map, axis=0)
)
ax4_3.grid()
sim_vertex = None
for sim_shower in event.get_sim_showers():
sim_vertex = sim_shower.get_parameter(shp.vertex)
break
if sim_vertex is not None:
ax4_1.scatter(
[np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2)],
[sim_vertex[2]],
c='r',
alpha=.5,
marker='+'
)
ax4_2.scatter(
[hp.cartesian_to_spherical(sim_vertex[0], sim_vertex[1], sim_vertex[2])[1]],
[np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2)],
c='r',
alpha=.5,
marker='+'
)
ax4_3.scatter(
[hp.get_normalized_angle(
hp.cartesian_to_spherical(sim_vertex[0], sim_vertex[1], sim_vertex[2])[1]) / units.deg],
[(sim_vertex[2])],
c='r',
alpha=.5,
marker='+'
)
ax4_1.set_xlabel('d [m]')
ax4_1.set_ylabel('z [m]')
ax4_3.set_xlabel(r'$\phi [^\circ]$')
ax4_3.set_ylabel('z [m]')
fig4.tight_layout()
fig4.savefig('{}/{}_{}_2D_correlation_maps.png'.format(self.__debug_folder, event.get_run_number(), event.get_id()))
def __draw_search_zones(
self,
event,
slope,
offset,
line_fit_d,
line_fit_z,
min_z_offset,
max_z_offset,
i_max_d,
i_max_z,
corr_mask_d,
corr_mask_z,
z_fit,
i_max_theta,
theta_corr_mask,
median_theta
):
fig5 = plt.figure(figsize=(8, 8))
ax5_1_1 = fig5.add_subplot(221)
ax5_1_1.grid()
ax5_1_2 = fig5.add_subplot(222)
ax5_1_2.grid()
ax5_1_1.fill_between(
self.__distances_2d,
self.__distances_2d * slope + offset + 1.1 * max_z_offset,
self.__distances_2d * slope + offset - 1.1 * min_z_offset,
color='k',
alpha=.2
)
ax5_1_1.scatter(
self.__distances_2d[corr_mask_d],
self.__z_coordinates_2d[i_max_d][corr_mask_d]
)
ax5_1_1.scatter(
self.__distances_2d[~corr_mask_d],
self.__z_coordinates_2d[i_max_d][~corr_mask_d],
c='k',
alpha=.5
)
ax5_1_1.plot(
self.__distances_2d,
self.__distances_2d * slope + offset,
color='k',
linestyle=':'
)
ax5_1_1.fill_between(
self.__distances_2d,
self.__distances_2d * slope + offset + 1.1 * max_z_offset,
self.__distances_2d * slope + offset - 1.1 * min_z_offset,
color='k',
alpha=.2
)
ax5_1_1.plot(
self.__distances_2d,
self.__distances_2d * line_fit_d[0] + line_fit_d[1],
color='r',
linestyle=':'
)
ax5_1_1.fill_between(
self.__distances_2d,
self.__distances_2d * line_fit_d[0] + line_fit_d[1] + 1.1 * max_z_offset,
self.__distances_2d * line_fit_d[0] + line_fit_d[1] - 1.1 * min_z_offset,
color='r',
alpha=.2
)
ax5_1_2.scatter(
self.__distances_2d[i_max_z][corr_mask_z],
self.__z_coordinates_2d[corr_mask_z]
)
ax5_1_2.scatter(
self.__distances_2d[i_max_z][~corr_mask_z],
self.__z_coordinates_2d[~corr_mask_z],
c='k',
alpha=.5
)
ax5_1_2.plot(
self.__distances_2d,
self.__distances_2d * slope + offset,
color='k',
linestyle=':'
)
ax5_1_2.plot(
self.__distances_2d,
self.__distances_2d * line_fit_z[0] + line_fit_z[1],
color='r',
linestyle=':'
)
ax5_1_2.fill_between(
self.__distances_2d,
self.__distances_2d * line_fit_z[0] + line_fit_z[1] + 1.1 * max_z_offset,
self.__distances_2d * line_fit_z[0] + line_fit_z[1] - 1.1 * min_z_offset,
color='r',
alpha=.2
)
ax5_1_2.plot(
self.__distances_2d,
self.__distances_2d * slope + offset,
color='k',
linestyle=':'
)
ax5_1_2.fill_between(
self.__distances_2d,
self.__distances_2d * slope + offset + 1.1 * max_z_offset,
self.__distances_2d * slope + offset - 1.1 * min_z_offset,
color='k',
alpha=.2
)
if z_fit:
ax5_2 = fig5.add_subplot(2, 2, (3, 4))
ax5_2.grid()
else:
ax5_2 = fig5.add_subplot(2, 2, (3, 4), projection='polar')
if z_fit:
ax5_2.scatter(
self.__azimuths_2d[i_max_theta][theta_corr_mask],
np.abs(self.__z_coordinates_2d)[theta_corr_mask]
)
ax5_2.scatter(
self.__azimuths_2d[i_max_theta][~theta_corr_mask],
np.abs(self.__z_coordinates_2d)[~theta_corr_mask],
c='k',
alpha=.5
)
ax5_2.plot(
[median_theta, median_theta],
[np.abs(self.__z_coordinates_2d[0]), np.abs(self.__z_coordinates_2d[-1])],
color='k',
linestyle=':'
)
else:
ax5_2.scatter(
self.__azimuths_2d[i_max_theta][theta_corr_mask],
self.__distances_2d[theta_corr_mask]
)
ax5_2.scatter(
self.__azimuths_2d[i_max_theta][~theta_corr_mask],
self.__distances_2d[~theta_corr_mask],
c='k',
alpha=.5
)
ax5_2.plot(
[median_theta, median_theta],
[self.__distances_2d[0], self.__distances_2d[-1]],
color='k',
linestyle=':'
)
sim_vertex = None
for sim_shower in event.get_sim_showers():
sim_vertex = sim_shower.get_parameter(shp.vertex)
break
if sim_vertex is not None:
ax5_1_1.axhline(sim_vertex[2], color='r', linestyle='--', alpha=.5)
ax5_1_1.axvline(np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2), color='r', linestyle='--', alpha=.5)
ax5_1_2.axhline(sim_vertex[2], color='r', linestyle='--', alpha=.5)
ax5_1_2.axvline(np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2), color='r', linestyle='--', alpha=.5)
ax5_2.scatter(
[hp.cartesian_to_spherical(sim_vertex[0], sim_vertex[1], sim_vertex[2])[1]],
[np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2)],
c='r',
alpha=.5,
marker='+'
)
ax5_1_1.set_xlabel('d [m]')
ax5_1_2.set_xlabel('d [m]')
ax5_1_1.set_ylabel('z [m]')
ax5_1_2.set_ylabel('z [m]')
if z_fit:
ax5_2.set_xlabel(r'$\phi [^\circ]$')
ax5_2.set_ylabel('|z| [m]')
fig5.tight_layout()
fig5.savefig('{}/{}_{}_search_zones.png'.format(self.__debug_folder, event.get_run_number(), event.get_id()))
def __draw_vertex_reco(
self,
event,
correlation_sum,
x_0,
y_0,
z_0,
x_coords,
y_coords,
z_coords,
slope,
offset,
median_theta,
i_max
):
fig6 = plt.figure(figsize=(8, 16))
ax6_1 = fig6.add_subplot(411)
vmin = .6
vmax = 1.
colormap = cm.get_cmap('viridis', 16)
sim_vertex = None
for sim_shower in event.get_sim_showers():
sim_vertex = sim_shower.get_parameter(shp.vertex)
break
cplot1 = ax6_1.pcolor(
x_0[0],
z_0[0],
np.max(correlation_sum, axis=0),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot1, ax=ax6_1)
ax6_2 = fig6.add_subplot(412)
cplot2 = ax6_2.pcolor(
x_0[:, :, 0],
y_0[:, :, 0],
np.max(correlation_sum, axis=2),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot2, ax=ax6_2)
ax6_3 = fig6.add_subplot(413)
cplot3 = ax6_3.pcolor(
x_0[0],
z_coords[0],
np.max(correlation_sum, axis=0),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot3, ax=ax6_3)
ax6_3.set_aspect('equal')
ax6_4 = fig6.add_subplot(414)
cplot4 = ax6_4.pcolor(
x_coords[:, :, 0],
y_coords[:, :, 0],
np.max(correlation_sum, axis=2),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot4, ax=ax6_4)
ax6_4.set_aspect('equal')
if sim_vertex is not None:
sim_vertex_dhor = np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2)
ax6_1.scatter(
[sim_vertex_dhor],
[sim_vertex[2] - sim_vertex_dhor * slope - offset],
c='r',
marker='+'
)
ax6_2.scatter(
[np.cos(median_theta) * sim_vertex[0] + np.sin(median_theta) * sim_vertex[1]],
[-np.sin(median_theta) * sim_vertex[0] + np.cos(median_theta) * sim_vertex[1]],
c='r',
marker='+'
)
ax6_3.scatter(
[sim_vertex_dhor],
[sim_vertex[2]],
c='r',
marker='+'
)
ax6_4.scatter(
[sim_vertex[0]],
[sim_vertex[1]],
c='r',
marker='+'
)
ax6_1.scatter(
[x_0[i_max]],
[z_0[i_max]],
c='k',
marker='+'
)
ax6_2.scatter(
[x_0[i_max]],
[y_0[i_max]],
c='k',
marker='+'
)
ax6_3.scatter(
[x_0[i_max]],
[z_coords[i_max]],
c='k',
marker='+'
)
ax6_4.scatter(
[x_coords[i_max]],
[y_coords[i_max]],
c='k',
marker='+'
)
ax6_1.set_xlabel('r [m]')
ax6_1.set_ylabel(r'$\Delta z$ [m]')
ax6_2.set_xlabel(r'$\Delta x$ [m]')
ax6_2.set_ylabel(r'$\Delta y$ [m]')
ax6_3.set_xlabel('r [m]')
ax6_3.set_ylabel('z [m]')
ax6_4.set_xlabel('x [m]')
ax6_4.set_ylabel('y [m]')
fig6.tight_layout()
fig6.savefig('{}/{}_{}_slices.png'.format(self.__debug_folder, event.get_run_number(), event.get_id()))
def __draw_dnr_reco(
self,
event,
correlation_sum,
self_correlation_sum,
combined_correlations,
x_0,
y_0,
z_0,
slope,
offset,
median_theta,
i_max,
i_max_dnr
):
fig8 = plt.figure(figsize=(12, 8))
ax8_1 = fig8.add_subplot(232)
ax8_2 = fig8.add_subplot(235)
ax8_3 = fig8.add_subplot(233)
ax8_4 = fig8.add_subplot(236)
ax8_5 = fig8.add_subplot(231)
ax8_6 = fig8.add_subplot(234)
colormap = cm.get_cmap('viridis', 16)
vmin = .6
vmax = 1.
cplot1 = ax8_1.pcolor(
x_0[0],
z_0[0],
np.max(self_correlation_sum, axis=0),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot1, ax=ax8_1)
cplot2 = ax8_2.pcolor(
x_0[:, :, 0],
y_0[:, :, 0],
np.max(self_correlation_sum, axis=2),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot2, ax=ax8_2)
cplot3 = ax8_3.pcolor(
x_0[0],
z_0[0],
np.max(combined_correlations, axis=0),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot3, ax=ax8_3)
cplot4 = ax8_4.pcolor(
x_0[:, :, 0],
y_0[:, :, 0],
np.max(combined_correlations, axis=2),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot4, ax=ax8_4)
cplot5 = ax8_5.pcolor(
x_0[0],
z_0[0],
np.max(correlation_sum, axis=0),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
plt.colorbar(cplot5, ax=ax8_5)
cplot6 = ax8_6.pcolor(
x_0[:, :, 0],
y_0[:, :, 0],
np.max(correlation_sum, axis=2),
cmap=colormap,
vmin=vmin,
vmax=vmax
)
ax8_5.scatter(
[x_0[i_max]],
[z_0[i_max]],
c='k',
marker='+'
)
ax8_3.scatter(
[x_0[i_max_dnr]],
[z_0[i_max_dnr]],
c='k',
marker='+'
)
plt.colorbar(cplot6, ax=ax8_6)
sim_vertex = None
for sim_shower in event.get_sim_showers():
sim_vertex = sim_shower.get_parameter(shp.vertex)
break
if sim_vertex is not None:
sim_vertex_dhor = np.sqrt(sim_vertex[0] ** 2 + sim_vertex[1] ** 2)
ax8_1.scatter(
[sim_vertex_dhor],
[sim_vertex[2] - sim_vertex_dhor * slope - offset],
c='r',
marker='+'
)
ax8_2.scatter(
[np.cos(median_theta) * sim_vertex[0] + np.sin(median_theta) * sim_vertex[1]],
[-np.sin(median_theta) * sim_vertex[0] + np.cos(median_theta) * sim_vertex[1]],
c='r',
marker='+'
)
ax8_3.scatter(
[sim_vertex_dhor],
[sim_vertex[2] - sim_vertex_dhor * slope - offset],
c='r',
marker='+'
)
ax8_4.scatter(
[np.cos(median_theta) * sim_vertex[0] + np.sin(median_theta) * sim_vertex[1]],
[-np.sin(median_theta) * sim_vertex[0] + np.cos(median_theta) * sim_vertex[1]],
c='r',
marker='+'
)
ax8_5.scatter(
[sim_vertex_dhor],
[sim_vertex[2] - sim_vertex_dhor * slope - offset],
c='r',
marker='+'
)
ax8_6.scatter(
[np.cos(median_theta) * sim_vertex[0] + np.sin(median_theta) * sim_vertex[1]],
[-np.sin(median_theta) * sim_vertex[0] + np.cos(median_theta) * sim_vertex[1]],
c='r',
marker='+'
)
ax8_1.grid()
ax8_2.grid()
fig8.tight_layout()
fig8.savefig('{}/{}_{}_dnr_reco.png'.format(self.__debug_folder, event.get_run_number(), event.get_id()))