#!/usr/bin/env python3 # vim: fdm=indent ts=4 """ Find the best integer multiple to shift antennas to be able to resolve """ import h5py from itertools import combinations, zip_longest, product import matplotlib.pyplot as plt import numpy as np from os import path from scipy.interpolate import interp1d from earsim import REvent from atmocal import AtmoCal import aa_generate_beacon as beacon import lib from lib import rit def find_best_sample_shifts_summing_at_location(test_loc, antennas, allowed_sample_shifts, dt=None, sample_shift_first_trace=0, plot_iteration_with_shifted_trace=None, fig_dir=None, fig_distinguish=None): """ Find the best sample_shift for each antenna by summing the antenna traces and seeing how to get the best alignment. """ a_ = [] t_ = [] t_min = 1e9 t_max = -1e9 a_maxima = [] if dt is None: dt = antennas[0].t_AxB[1] - antennas[0].t_AxB[0] # propagate to test location for i, ant in enumerate(antennas): aloc = [ant.x, ant.y, ant.z] delta, dist = atm.light_travel_time(test_loc, aloc) delta = delta*1e9 t__ = np.subtract(ant.t_AxB, delta) t_.append(t__) a_.append(ant.E_AxB) a_maxima.append(max(ant.E_AxB)) if t__[0] < t_min: t_min = t__[0] if t__[-1] > t_max: t_max = t__[-1] # sort traces with descending maxima sort_idx = np.argsort(a_maxima)[::-1] t_ = [ t_[i] for i in sort_idx ] a_ = [ a_[i] for i in sort_idx ] # Interpolate and find best sample shift max_neg_shift = 0 #np.min(allowed_sample_shifts) * dt max_pos_shift = 0 #np.max(allowed_sample_shifts) * dt t_sum = np.arange(t_min+max_neg_shift, t_max+max_pos_shift, dt) a_sum = np.zeros(len(t_sum)) best_sample_shifts = np.zeros( (len(antennas)) ,dtype=int) for i, (t_r, E_) in enumerate(zip(t_, a_)): f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0) a_int = f(t_sum) if i == 0: a_sum += a_int best_sample_shifts[i] = sample_shift_first_trace continue # init figure if i == plot_iteration_with_shifted_trace: fig, ax = plt.subplots() ax.set_title("Traces at ({:.1f},{:.1f},{:.1f})".format(*test_loc)) ax.set_xlabel("Time [ns]") ax.set_ylabel("Amplitude") ax.plot(t_sum, a_sum) shift_maxima = np.zeros( len(allowed_sample_shifts) ) for j, shift in enumerate(allowed_sample_shifts): augmented_a = np.roll(a_int, shift) shift_maxima[j] = np.max(augmented_a + a_sum) if i == plot_iteration_with_shifted_trace: ax.plot(t_sum, augmented_a, label=f'{shift} shifted') # transform maximum into best_sample_shift best_idx = np.argmax(shift_maxima) best_sample_shifts[i] = allowed_sample_shifts[best_idx] a_sum += np.roll(a_int, best_sample_shifts[i]) # cleanup figure if i == plot_iteration_with_shifted_trace: if fig_dir: # note this is a global variable here fig.savefig(path.join(fig_dir, __file__ + '.loc{:.1f}-{:.1f}-{:.1f}'.format(*test_loc) + f'.i{i}{fig_distinguish}.pdf')) plt.close(fig) # sort by antenna (undo sorting by maximum) undo_sort_idx = np.argsort(sort_idx) best_sample_shifts = best_sample_shifts[undo_sort_idx] # Return ks return best_sample_shifts, np.max(a_sum) if __name__ == "__main__": import sys import os import matplotlib if os.name == 'posix' and "DISPLAY" not in os.environ: matplotlib.use('Agg') atm = AtmoCal() fname = "ZH_airshower/mysim.sry" fig_dir = "./figures/periods_from_shower_figures/" fig_subdir = path.join(fig_dir, 'shifts/') allowed_ks = np.arange(-2, 3, 1) Xref = 400 x_coarse = np.linspace(-20e3, 20e3, 10) y_coarse = np.linspace(-20e3, 20e3, 10) x_fine = np.linspace(-2e3, 2e3, 30) y_fine = np.linspace(-2e3, 2e3, 30) N_runs = 3 #### fname_dir = path.dirname(fname) antennas_fname = path.join(fname_dir, beacon.antennas_fname) time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_fname # Read in antennas from file _, __, antennas = beacon.read_beacon_hdf5(antennas_fname) # Read original REvent ev = REvent(fname) # .. patch in our antennas ev.antennas = antennas # For now only implement using one freq_name freq_names = antennas[0].beacon_info.keys() if len(freq_names) > 1: raise NotImplementedError freq_name = next(iter(freq_names)) f_beacon = ev.antennas[0].beacon_info[freq_name]['freq'] # determine allowable ks per location dt = ev.antennas[0].t[1] - ev.antennas[0].t[0] allowed_sample_shifts = np.rint(allowed_ks/f_beacon /dt).astype(int) print("Checking:", allowed_ks, ": shifts :", allowed_sample_shifts) # Remove time due to true phase for i, ant in enumerate(ev.antennas): true_phase = ant.beacon_info[freq_name]['phase'] true_phase_time = true_phase/(2*np.pi*f_beacon) ev.antennas[i].t -= true_phase_time # Prepare polarisation and passbands rit.set_pol_and_bp(ev, low=0.03, high=0.08) dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0) scale2d = dXref*np.tan(np.deg2rad(2.)) # Setup Plane grid to test for r in range(N_runs): xoff, yoff = 0,0 if r == 0: x = x_coarse y = y_coarse else: old_ks_per_loc = ks_per_loc[best_idx] xoff = xx[best_idx] yoff = yy[best_idx] if r == 1: x = x_fine y = y_fine else: x /= 4 y /= 4 print(f"Testing grid run {r} centered on {xoff}, {yoff}") ks_per_loc = np.zeros( (len(x)*len(y), len(ev.antennas)) , dtype=int) maxima_per_loc = np.zeros( (len(x)*len(y))) ## Check each location on grid xx = [] yy = [] N_loc = len(maxima_per_loc) for i, (x_, y_) in enumerate(product(x,y)): tmp_fig_subdir = None if i % 10 ==0: print(f"Testing location {i} out of {N_loc}") tmp_fig_subdir = fig_subdir test_loc = (x_+xoff)* ev.uAxB + (y_+yoff)*ev.uAxAxB + dXref *ev.uA xx.append(x_+xoff) yy.append(y_+yoff) # Find best k for each antenna shifts, maximum = find_best_sample_shifts_summing_at_location(test_loc, ev.antennas, allowed_sample_shifts, dt=dt, fig_dir=tmp_fig_subdir, plot_iteration_with_shifted_trace=len(ev.antennas)-1, fig_distinguish=f".run{r}") # Translate sample shifts back into period multiple k ks = np.rint(shifts*f_beacon*dt) ks_per_loc[i] = ks maxima_per_loc[i] = maximum xx = np.array(xx) yy = np.array(yy) locs = list(zip(xx, yy)) ## Save maxima to file np.savetxt(fig_dir + __file__+f'.maxima.run{r}.txt', np.column_stack((locs, maxima_per_loc, ks_per_loc)) ) if True: #plot maximum at test locations fig, axs = plt.subplots() axs.set_title(f"Grid Run {r}") axs.set_ylabel("vxvxB [km]") axs.set_xlabel(" vxB [km]") sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6) fig.colorbar(sc, ax=axs) if fig_dir: fig.savefig(path.join(fig_dir, __file__+f'.maxima.run{r}.pdf')) ## Save ks to file best_idx = np.argmax(maxima_per_loc) np.savetxt(fig_dir + __file__+f'.bestk.run{r}.txt', ks_per_loc[best_idx] ) print('Best k:', ks_per_loc[best_idx]) # Abort if no improvement if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ): print("No changes from previous grid, breaking") break # Save best ks to hdf5 antenna file with h5py.File(antennas_fname, 'a') as fp: group = fp.require_group('antennas') for i, ant in enumerate(antennas): h5ant = group[ant.name] h5ant.attrs['best_k'] = old_ks_per_loc[i] plt.show()