#!/usr/bin/env python3 # vim: fdm=indent ts=4 """ Do a reconstruction of airshower after correcting for the clock offsets. """ import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D # required for projection='3d' on old matplotliblib versions import numpy as np from os import path import pickle import joblib from earsim import REvent from atmocal import AtmoCal import aa_generate_beacon as beacon import lib from lib import rit if __name__ == "__main__": import sys import os import matplotlib if os.name == 'posix' and "DISPLAY" not in os.environ: matplotlib.use('Agg') atm = AtmoCal() from scriptlib import MyArgumentParser parser = MyArgumentParser() parser.add_argument('--input-fname', type=str, default=None, help='Path to mysim.sry, either directory or path. If empty it takes DATA_DIR and appends mysim.sry. (Default: %(default)s)') args = parser.parse_args() if not args.input_fname: args.input_fname = args.data_dir if path.isdir(args.input_fname): args.input_fname = path.join(args.input_fname, "mysim.sry") figsize = (12,8) fig_dir = args.fig_dir fig_subdir = path.join(fig_dir, 'reconstruction') show_plots = args.show_plots apply_signal_window_from_max = True remove_beacon_from_traces = True #### fname_dir = args.data_dir antennas_fname = path.join(fname_dir, beacon.antennas_fname) pickle_fname = path.join(fname_dir, 'res.pkl') tx_fname = path.join(fname_dir, beacon.tx_fname) beacon_snr_fname = path.join(fname_dir, beacon.beacon_snr_fname) # create fig_dir if fig_dir: os.makedirs(fig_dir, exist_ok=True) if fig_subdir: os.makedirs(fig_subdir, exist_ok=True) # Read in antennas from file _, tx, antennas = beacon.read_beacon_hdf5(antennas_fname) _, __, txdata = beacon.read_tx_file(tx_fname) # Read original REvent ev = REvent(args.input_fname) # .. patch in our antennas ev.antennas = antennas # Read in snr info beacon_snrs = beacon.read_snr_file(beacon_snr_fname) snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .1g}" # 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'] # Repair clock offsets with the measured offsets measured_repair_offsets = beacon.read_antenna_clock_repair_offsets(ev.antennas, mode='phases', freq_name=freq_name) for i, ant in enumerate(ev.antennas): # t_AxB will be set by the rit.set_pol_and_bp function ev.antennas[i].orig_t = ev.antennas[i].t ev.antennas[i].t += measured_repair_offsets[i] ev.antennas[i].t_AxB += measured_repair_offsets[i] if apply_signal_window_from_max: N_pre, N_post = 250, 250 # TODO: make this configurable # Get max idx from all the traces # and select the strongest max_idx = [] maxs = [] for trace in [ant.Ex, ant.Ey, ant.Ez]: idx = np.argmax(np.abs(trace)) max_idx.append(idx) maxs.append( np.abs(trace[idx]) ) idx = np.argmax(maxs) max_idx = max_idx[idx] # Create window around max_idx low_idx = max(0, max_idx-N_pre) high_idx = min(len(ant.t), max_idx+N_post) ev.antennas[i].orig_t = ant.orig_t[low_idx:high_idx] ev.antennas[i].t = ant.t[low_idx:high_idx] ev.antennas[i].Ex = ant.Ex[low_idx:high_idx] ev.antennas[i].Ey = ant.Ey[low_idx:high_idx] ev.antennas[i].Ez = ant.Ez[low_idx:high_idx] ev.antennas[i].t_AxB = ant.t_AxB[low_idx:high_idx] ev.antennas[i].E_AxB = ant.E_AxB[low_idx:high_idx] # .. and remove the beacon from the traces # Note: ant.E_AxB is recalculated by rit.set_pol_and_bp if remove_beacon_from_traces: clock_phase = measured_repair_offsets[i]*2*np.pi*f_beacon beacon_phase = ant.beacon_info[freq_name]['beacon_phase'] f = ant.beacon_info[freq_name]['freq'] ampl_AxB = ant.beacon_info[freq_name]['amplitude'] calc_beacon = lib.sine_beacon(f, ev.antennas[i].t, amplitude=ampl_AxB, phase=beacon_phase-clock_phase) tx_amps = txdata['amplitudes'] tx_amps_sum = np.sum(tx_amps) # Split up contribution to the various polarisations for j, amp in enumerate(tx_amps): if j == 0: ev.antennas[i].Ex -= amp*(1/tx_amps_sum)*calc_beacon elif j == 1: ev.antennas[i].Ey -= amp*(1/tx_amps_sum)*calc_beacon elif j == 2: ev.antennas[i].Ez -= amp*(1/tx_amps_sum)*calc_beacon # Subtract the beacon from E_AxB ev.antennas[i].E_AxB -= calc_beacon ## ## Make a figure of the manipulated traces ## if i == 72: orig_beacon_amplifier = ampl_AxB/max(ant.beacon) for k in range(2): if k == 0: time = ant.t_AxB trace = ant.E_AxB tmp_beacon = calc_beacon fname_extra = "" else: time = ant.t trace = ant.Ex tmp_beacon = tx_amps[0]/tx_amps_sum * calc_beacon fname_extra = ".Ex" fig, ax = plt.subplots(figsize=figsize) ax.set_title(f"Signal and Beacon traces Antenna {ant.name}") ax.set_xlabel("Time [ns]") ax.set_ylabel("Amplitude [$\\mu V/m$]") ax.plot(time, trace + tmp_beacon, alpha=0.6, ls='dashed', label='Signal') # calc_beacon was already removed ax.plot(time, tmp_beacon, alpha=0.6, ls='dashed', label='Calc Beacon') ax.plot(time, trace, alpha=0.6, label="Signal - Calc Beacon") if k == 0: ax.legend(title=snr_str) else: ax.legend(title="Ex " + snr_str) # save if fig_dir: fig.tight_layout() if True: # zoom old_xlim = ax.get_xlim() if True: # zoomed on part without peak of this trace wx, x = 100, 100 ax.set_xlim(x-wx, x+wx) fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.traces.A{ant.name}.zoomed.beacon{fname_extra}.pdf')) if True: # zoomed on peak of this trace idx = np.argmax(ev.antennas[i].E_AxB) x = ev.antennas[i].t_AxB[idx] wx = 100 ax.set_xlim(x-wx, x+wx) fig.savefig(path.join(fig_dir, path.basename(__file__)+f".traces.A{ant.name}.zoomed.peak{fname_extra}.pdf")) ax.set_xlim(*old_xlim) fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.traces.A{i}.pdf')) N_X, Xlow, Xhigh = 23, 100, 1200 with joblib.parallel_backend("loky"): res = rit.reconstruction(ev, outfile=fig_subdir+'/fig.pdf', slice_outdir=fig_subdir+'/', Xlow=Xlow, N_X=N_X, Xhigh=Xhigh, disable_pol_and_bp=True) ## Save a pickle with open(pickle_fname, 'wb') as fp: pickle.dump(res,fp) if show_plots: plt.show()