220 lines
8.1 KiB
Python
Executable File
220 lines
8.1 KiB
Python
Executable File
#!/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)')
|
|
parser.add_argument('--clock-repair', type=str, default='full', choices=['orig', 'ks', 'phases', 'full'], help='How to repair the clock offsets. (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
|
|
clock_repair_mode = args.clock_repair
|
|
fig_subdir = path.join(fig_dir, 'reconstruction', clock_repair_mode)
|
|
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-'+clock_repair_mode+'.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=clock_repair_mode, 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'))
|
|
|
|
|
|
if False: # very fine grid
|
|
N_X, Xlow, Xhigh = 23, 100, 1200
|
|
else:
|
|
N_X, Xlow, Xhigh = 6, 200, 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()
|