Source code for NuRadioReco.modules.neutrinoVertexReconstructor.neutrino3DVertexReconstructor

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()))