From 42236b03a82f25032c950a58943cf1537e973333 Mon Sep 17 00:00:00 2001 From: Eric Teunis de Boone Date: Tue, 31 Oct 2023 16:47:45 +0100 Subject: [PATCH] ZH: ca_period_from_shower as used for Thesis Figures --- .../ca_period_from_shower.py | 226 +++++++++++++----- 1 file changed, 166 insertions(+), 60 deletions(-) diff --git a/airshower_beacon_simulation/ca_period_from_shower.py b/airshower_beacon_simulation/ca_period_from_shower.py index b533106..46386f9 100755 --- a/airshower_beacon_simulation/ca_period_from_shower.py +++ b/airshower_beacon_simulation/ca_period_from_shower.py @@ -25,7 +25,25 @@ try: except: tqdm = lambda x: x -def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, period=1, dt=None, period_shift_first_trace=0, plot_iteration_with_shifted_trace=None, fig_dir=None, fig_distinguish=None,snr_str=None, shower_plane_loc=None): +try: + from joblib import Parallel, delayed +except: + Parallel = None + delayed = lambda x: x + +def find_best_period_shifts_at_location(*args, algo=None, **kwargs): + """ + This is a placeholder function. + + For args and kwargs see find_best_period_shifts_summing_at_location. + """ + + if algo is None: + algo = 'sum' + + return find_best_period_shifts_summing_at_location(*args, **kwargs, algo=algo) + +def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, period=1, dt=None, period_shift_first_trace=0, plot_iteration_with_shifted_trace=None, fig_dir=None, fig_distinguish=None,snr_str=None, shower_plane_loc=None, algo='sum', verbose=False): """ Find the best sample_shift for each antenna by summing the antenna traces and seeing how to get the best alignment. @@ -80,9 +98,9 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0) if i == 0: - a_sum += f(t_sum) - a_first = a_sum best_period_shifts[i] = period_shift_first_trace + a_first = f(t_sum - period_shift_first_trace*period) + a_sum += a_first continue # init figure @@ -92,47 +110,67 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, title_location = "s({:.1g},{:.1g},{:.1g})".format(*shower_plane_loc) else: title_location = "({.1g},{:.1g},{:.1g})".format(*test_loc) - ax.set_title("Traces at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant)) + #ax.set_title("Traces at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant)) ax.set_xlabel("Time [ns]") ax.set_ylabel("Amplitude") - ax.plot(t_sum, a_sum) + #ax.plot(t_sum, a_sum) fig2, ax2 = plt.subplots(figsize=figsize) - ax2.set_title("Maxima at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant)) + #ax2.set_title("Maxima at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant)) ax2.set_xlabel("$k$th Period") ax2.set_ylabel("Summed Amplitude") - ax2.plot(0, np.max(a_first), marker='*', label='strongest trace', ls='none', ms=20) + ax2.plot(0, np.max(a_first), marker='*', label='first trace', ls='none', ms=20) + ax3 = ax2.twinx() + ax3.set_ylabel("Correlation") # find the maxima for each period shift k shift_maxima = np.zeros( len(allowed_ks) ) + shift_corrs = np.zeros_like(shift_maxima) for j, k in enumerate(allowed_ks): - augmented_a = f(t_sum + k*period) + augmented_a = f(t_sum - k*period) shift_maxima[j] = np.max(augmented_a + a_first) + shift_corrs[j] = np.dot(augmented_a, a_first) if i in plot_iteration_with_shifted_trace and abs(k) <= 3: - ax.plot(t_sum, augmented_a, alpha=0.7, ls='dashed', label=f'{k:g}') + l = ax.plot(t_sum, a_first + augmented_a, alpha=0.7, ls='dashed', label=f'{k:g}') + ax.axhline(shift_maxima[j], ls='dashdot', color=l[0].get_color(), alpha=0.7) + ax2.plot(k, shift_maxima[j], marker='o', ls='none', ms=20) + ax3.plot(k, shift_corrs[j], marker='3', ls='none', ms=20) # transform maximum into best_sample_shift best_idx = np.argmax(shift_maxima) + best_corr_idx = np.argmax(shift_corrs) - best_period_shifts[i] = allowed_ks[best_idx] - best_augmented_a = f(t_sum + k*period) + if verbose and (best_corr_idx != best_idx): + print("Correlation idx not equal to maximum idx") + print(best_corr_idx, best_idx) + + if algo == 'sum': + best_period_shifts[i] = allowed_ks[best_idx] + elif algo == 'corr': + best_period_shifts[i] = allowed_ks[best_corr_idx] + + k = best_period_shifts[i] + best_augmented_a = f(t_sum - k*period) a_sum += best_augmented_a # cleanup figure if i in plot_iteration_with_shifted_trace: # plot the traces if True: # plot best k again - ax.plot(t_sum, best_augmented_a, alpha=0.8, label=f'best $k$={best_period_shifts[i]:g}', lw=2) + l = ax.plot(t_sum, a_first + best_augmented_a, alpha=0.8, label=f'best $k$={best_period_shifts[i]:g}', lw=2) + ax.axhline(shift_maxima[j], ls='dashdot', color=l[0].get_color(), alpha=0.7) if True: # plot best shift ax2.plot(allowed_ks[best_idx], shift_maxima[best_idx], marker='*', ls='none', ms=20, label=f'best $k$={best_period_shifts[i]:g}') + ax3.plot(allowed_ks[best_corr_idx], shift_maxima[best_idx], marker='*', ls='none', ms=20, label=f'best corr $k$={allowed_ks[best_corr_idx]:g}') - ax.legend(title='period shift $k$; '+snr_str, ncol=5 ) + ax.legend(title='period shift $k$; '+snr_str, ncol=5, loc='lower center') ax2.legend(title=snr_str) + ax3.legend() if fig_dir: fig.tight_layout() fig2.tight_layout() @@ -147,14 +185,15 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, old_xlim = ax.get_xlim() if True: # zoomed on part without peak of this trace - wx = 100 + wx = 200 x = max(t_r) - wx ax.set_xlim(x-wx, x+wx) fig.savefig(fname + ".zoomed.beacon.pdf") if True: # zoomed on peak of this trace - x = t_r[np.argmax(E_)] - wx = 50 + (max(best_period_shifts) - min(best_period_shifts))*dt + x = t_sum[np.argmax(a_first)] + x = t_sum[np.argmax(f(t_sum))] + wx = 50 + (max(best_period_shifts) - min(best_period_shifts) )*dt + 1*period ax.set_xlim(x-wx, x+wx) fig.savefig(fname + ".zoomed.peak.pdf") @@ -165,13 +204,30 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, plt.close(fig) plt.close(fig2) + if True: # final summed waveform + fig, ax = plt.subplots() + + ax.set_title("Summed Traces with best k's at {}".format(title_location)) + ax.set_xlabel("Time [ns]") + ax.set_ylabel("Amplitude") + + ax.plot(t_sum, a_sum) + + if fig_dir: + fig.tight_layout() + + fname = path.join(fig_dir, path.basename(__file__) + f'.{fig_distinguish}i{i}' + fname_location) + fig.savefig(fname + ".sum.pdf") + + plt.close(fig) + # sort by antenna (undo sorting by maximum) undo_sort_idx = np.argsort(sort_idx) best_period_shifts = best_period_shifts[undo_sort_idx] # Return ks - return best_period_shifts, np.max(a_sum) + return best_period_shifts, np.max(a_sum), sort_idx if __name__ == "__main__": import sys @@ -204,7 +260,16 @@ if __name__ == "__main__": if path.isdir(args.input_fname): args.input_fname = path.join(args.input_fname, "mysim.sry") - figsize = (12,8) + figsize = (6,4) + if True: + from matplotlib import rcParams + #rcParams["text.usetex"] = True + rcParams["font.family"] = "serif" + rcParams["font.size"] = "14" + rcParams["grid.linestyle"] = 'dotted' + rcParams["figure.figsize"] = figsize + figsize = rcParams['figure.figsize'] + fig_dir = args.fig_dir fig_subdir = path.join(fig_dir, 'shifts/') @@ -248,7 +313,7 @@ if __name__ == "__main__": ev.antennas = antennas # read in snr information beacon_snrs = beacon.read_snr_file(beacon_snr_fname) - snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .1g}" + snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .2g}" # For now only implement using one freq_name freq_names = antennas[0].beacon_info.keys() @@ -349,8 +414,8 @@ if __name__ == "__main__": if True: # zoom old_xlim = ax.get_xlim() - if True: # zoomed on part without peak of this trace - wx, x = 200, min(ant.t_AxB) + if not True: # zoomed on part without peak of this trace + wx, x = 100, min(ant.t_AxB) ax.set_xlim(x-5, x+wx) fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.traces.A{ant.name}.zoomed.beacon.pdf')) @@ -358,7 +423,7 @@ if __name__ == "__main__": if True: # zoomed on peak of this trace idx = np.argmax(ev.antennas[i].E_AxB) x = ev.antennas[i].t_AxB[idx] - wx = 300 + wx = 150 ax.set_xlim(x-wx//2, x+wx//2) fig.savefig(path.join(fig_dir, path.basename(__file__)+f".traces.A{ant.name}.zoomed.peak.pdf")) @@ -380,22 +445,23 @@ if __name__ == "__main__": ## Determine grid positions ## - dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0) + zgr = 0 + ev.core[2] + dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,zgr) scale2d = dXref*np.tan(np.deg2rad(2.)) scale4d = dXref*np.tan(np.deg2rad(4.)) if args.quick_run: #quicky - x_coarse = np.linspace(-scale2d, scale2d, 6) - y_coarse = np.linspace(-scale2d, scale2d, 6) + x_coarse = np.linspace(-scale4d, scale4d, 16) + y_coarse = np.linspace(-scale4d, scale4d, 16) x_fine = x_coarse/4 y_fine = y_coarse/4 else: # long - x_coarse = np.linspace(-scale4d, scale4d, 40) - y_coarse = np.linspace(-scale4d, scale4d, 40) + x_coarse = np.linspace(-scale4d, scale4d, 14) + y_coarse = np.linspace(-scale4d, scale4d, 14) - x_fine = np.linspace(-scale2d, scale2d, 40) - y_fine = np.linspace(-scale2d, scale2d, 40) + x_fine = np.linspace(-scale2d, scale2d, 18) + y_fine = np.linspace(-scale2d, scale2d, 18) ## Remove run_break_fname if it exists try: @@ -435,9 +501,29 @@ if __name__ == "__main__": yy = [] N_loc = len(maxima_per_loc) - for i, (x_, y_) in tqdm(enumerate(product(x,y)), total=N_loc, ): + # Make the grid a list of locations + xx = [] + yy = [] + for i, (x_, y_) in enumerate(product(x,y)): + xx.append( x_+xoff ) + yy.append( y_+yoff ) + + xx = np.array(xx) + yy = np.array(yy) + locs = list(zip(xx, yy)) + + # Shift the reference trace + if r == 0: + first_trace_period_shift = 0 + else: + first_trace_period_shift = 0#first_trace_period_shift + np.rint(np.mean(old_ks_per_loc)) + + print("New first trace period:", first_trace_period_shift) + + # define loop func for joblib + def loop_func(loc, dXref=dXref, i=1): tmp_fig_subdir = None - if i % 10 ==0: + if i == 0: if hasattr(tqdm, '__code__') and tqdm.__code__.co-name == '': print(f"Testing location {i} out of {N_loc}") tmp_fig_subdir = fig_subdir @@ -445,33 +531,41 @@ if __name__ == "__main__": test_loc = loc[0]* ev.uAxB + loc[1]*ev.uAxAxB + dXref *ev.uA # Find best k for each antenna - ks_per_loc[i], maxima_per_loc[i] = find_best_period_shifts_summing_at_location(test_loc, ev.antennas, allowed_ks, period=1/f_beacon, dt=dt, - plot_iteration_with_shifted_trace=[ 5, len(ev.antennas)-1], - fig_dir=tmp_fig_subdir, fig_distinguish=f"X{Xref}. run{r}.", + return find_best_period_shifts_at_location(test_loc, ev.antennas, allowed_ks, period=1/f_beacon, dt=dt, period_shift_first_trace=first_trace_period_shift, + plot_iteration_with_shifted_trace=[ 1, 2, 3, 4, 5, len(ev.antennas)-1], + fig_dir=tmp_fig_subdir, fig_distinguish=f"X{Xref}.run{r}.", snr_str=snr_str,shower_plane_loc=(loc[0]/1e3, loc[1]/1e3, dXref), ) - xx = np.array(xx) - yy = np.array(yy) - locs = list(zip(xx, yy)) + res = ( delayed(loop_func)(loc, i=i) for i, loc in enumerate(locs) ) + + if Parallel: + res = Parallel(n_jobs=None)(tqdm(res, total=len(locs))) + else: + res = tqdm(res, total=len(locs)) + + # unpack loop results + ks_per_loc, maxima_per_loc, sort_idx = zip(*res) ## Save maxima to file np.savetxt(path.join(fig_dir, path.basename(__file__)+f'.maxima.X{Xref}.run{r}.txt'), np.column_stack((locs, maxima_per_loc, ks_per_loc)) ) + scatter_kwargs = dict(cmap='Spectral_r', s=64*4, alpha=0.6) if True: #plot maximum at test locations fig, axs = plt.subplots(figsize=figsize) - axs.set_title(f"Optimizing signal strength by varying $k$ per antenna,\n Grid Run {r}") - axs.set_ylabel("vxvxB [km]") - axs.set_xlabel(" vxB [km]") + #axs.set_title(f"Optimizing signal strength by varying $k$ per antenna,\n Grid Run {r}") + axs.set_ylabel(" vxvxB [km]") + axs.set_xlabel("-v x B [km]") axs.set_aspect('equal', 'datalim') - sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6) + sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, **scatter_kwargs) fig.colorbar(sc, ax=axs, label='Max Amplitude [$\\mu V/m$]') - axs.legend(title=snr_str) + #axs.legend(title=snr_str) # indicate maximum value idx = np.argmax(maxima_per_loc) - axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) + axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) # max value + axs.plot(0,0, 'r+', ms=30) # true axis if fig_dir: old_xlims = axs.get_xlim() @@ -490,8 +584,10 @@ if __name__ == "__main__": best_idx = np.argmax(maxima_per_loc) best_k = ks_per_loc[best_idx] - print("Max at location: ", locs[best_idx]) - print('Best k:', best_k) + print("Max at location: ", locs[best_idx], ": ", maxima_per_loc[best_idx]) + print('Best k:', best_k[sort_idx[best_idx]]) # Sorted by DESC amplitude + print('Mean best k:', np.rint(np.mean(best_k))) + print('first k:', first_trace_period_shift) ## Save best ks to file np.savetxt(path.join(fig_dir, path.basename(__file__)+f'.bestk.X{Xref}.run{r}.txt'), best_k ) @@ -509,33 +605,43 @@ if __name__ == "__main__": # incorporate ks into timing for i, ant in enumerate(ev.antennas): - ev.antennas[i].t_AxB = ant.t_AxB - best_k[i] * 1/f_beacon + ev.antennas[i].t_AxB = ant.t_AxB + best_k[i] * 1/f_beacon - xx, yy, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=x[-1]-x[0], wy=y[-1]-y[0], xoff=xoff, yoff=yoff, zgr=0) + xx, yy, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=(x[-1]-x[0])/2, wy=(y[-1]-y[0])/2, xoff=xoff, yoff=yoff, zgr=0) # repair antenna times for i, backup_t_AxB in enumerate(backup_times): ev.antennas[i].t_AxB = backup_t_AxB else: # get maximum amplitude at each location maxima = np.empty( len(locs) ) - for i, loc in enumerate(locs): + for i, loc in tqdm(enumerate(locs), total=len(locs)): test_loc = loc[0]* ev.uAxB + loc[1]*ev.uAxAxB + dXref *ev.uA P, t_, a_, a_sum, t_sum = rit.pow_and_time(test_loc, ev, dt=dt) maxima[i] = np.max(a_sum) fig, axs = plt.subplots(figsize=figsize) - axs.set_title(f"Shower slice for best k, Grid Run {r}") - axs.set_ylabel("vxvxB [km]") - axs.set_xlabel(" vxB [km]") - axs.set_aspect('equal', 'datalim') - if power_reconstruction: - sc = axs.scatter(xx/1e3, yy/1e3, c=p, cmap='Spectral_r', alpha=0.6) - fig.colorbar(sc, ax=axs, label='Power') - else: - sc = axs.scatter(xx/1e3, yy/1e3, c=maxima, cmap='Spectral_r', alpha=0.6) - fig.colorbar(sc, ax=axs, label='Max Amplitude [$\\mu V/m$]') + #axs.set_title(f"Shower slice for best k, Grid Run {r}") + axs.set_ylabel(" vxvxB [km]") + axs.set_xlabel("-v x B [km]") - axs.legend(title=snr_str) + if power_reconstruction: + sc_c = p + sc_label = 'Power [$(\\mu V/m)^2$]' + else: + sc_c = maxima + sc_label='Max Amplitude [$\\mu V/m$]' + sc = axs.scatter(xx/1e3, yy/1e3, c=sc_c, **scatter_kwargs) + fig.colorbar(sc, ax=axs, label=sc_label) + + # indicate maximum value + idx = np.argmax(p if power_reconstruction else maxima) + axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) # max value + axs.plot(0,0, 'r+', ms=30) # true axis + + # make square figure + axs.set_aspect('equal', 'datalim') + + #axs.legend(title=snr_str) if fig_dir: if power_reconstruction: @@ -547,7 +653,7 @@ if __name__ == "__main__": fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.reconstruction.X{Xref}.run{r}.{fname_extra}.pdf')) # Abort if no improvement - if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ): + if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx] - first_trace_period_shift).all() ): print(f"No changes from previous grid, breaking at iteration {r} out of {N_runs}") try: