m-thesis-introduction/airshower_beacon_simulation/ca_period_from_shower.py

582 lines
22 KiB
Python
Executable file

#!/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
import os
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
try:
from tqdm import tqdm
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):
"""
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 = []
N_ant = len(antennas)
if dt is None:
dt = antennas[0].t_AxB[1] - antennas[0].t_AxB[0]
if not hasattr(plot_iteration_with_shifted_trace, '__len__'):
if plot_iteration_with_shifted_trace:
plot_iteration_with_shifted_trace = [ plot_iteration_with_shifted_trace ]
else:
plot_iteration_with_shifted_trace = []
# 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))
a_first = np.zeros(len(t_sum))
best_period_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)
if i == 0:
a_sum += f(t_sum)
a_first = a_sum
best_period_shifts[i] = period_shift_first_trace
continue
# init figure
if i in plot_iteration_with_shifted_trace:
fig, ax = plt.subplots(figsize=figsize)
if shower_plane_loc is not None:
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_xlabel("Time [ns]")
ax.set_ylabel("Amplitude")
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_xlabel("$k$th Period")
ax2.set_ylabel("Summed Amplitude")
ax2.plot(0, np.max(a_first), marker='*', label='strongest trace', ls='none', ms=20)
# find the maxima for each period shift k
shift_maxima = np.zeros( len(allowed_ks) )
for j, k in enumerate(allowed_ks):
augmented_a = f(t_sum + k*period)
shift_maxima[j] = np.max(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}')
ax2.plot(k, shift_maxima[j], marker='o', ls='none', ms=20)
# transform maximum into best_sample_shift
best_idx = np.argmax(shift_maxima)
best_period_shifts[i] = allowed_ks[best_idx]
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)
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}')
ax.legend(title='period shift $k$; '+snr_str, ncol=5 )
ax2.legend(title=snr_str)
if fig_dir:
fig.tight_layout()
fig2.tight_layout()
if shower_plane_loc is not None:
fname_location = '.sloc{:.1g}-{:.1g}-{:.1g}'.format(*shower_plane_loc)
else:
fname_location = '.loc{:.1f}-{:.1f}-{:.1f}'.format(*test_loc)
fname = path.join(fig_dir, path.basename(__file__) + f'.{fig_distinguish}i{i}' + fname_location)
if True:
old_xlim = ax.get_xlim()
if True: # zoomed on part without peak of this trace
wx = 100
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
ax.set_xlim(x-wx, x+wx)
fig.savefig(fname + ".zoomed.peak.pdf")
ax.set_xlim(*old_xlim)
fig.savefig(fname + ".pdf")
fig2.savefig(fname + ".maxima.pdf")
plt.close(fig)
plt.close(fig2)
# 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)
if __name__ == "__main__":
import sys
import os
import matplotlib
if os.name == 'posix' and "DISPLAY" not in os.environ:
matplotlib.use('Agg')
plt.rcParams.update({'figure.max_open_warning': 0})
atm = AtmoCal()
from scriptlib import MyArgumentParser
parser = MyArgumentParser(default_fig_dir="./figures/periods_from_shower_figures/")
parser.add_argument('--quick_run', action='store_true', help='Use a very coarse grid (6x6)')
parser.add_argument('-X', type=int, default=400, help='Atmospheric depth to start the initial grid.')
parser.add_argument('--max-k', type=float, default=2, help='Maximum abs(k) allowed to be shifted. (Default: %(default)d)')
parser.add_argument('-N', '--N_runs', type=int, default=5, help='Maximum amount of iterations to grid search. (Default: %(default)d)')
parser.add_argument('-l', '--passband-low', type=float, default=30e-3, help='Lower frequency [GHz] of the passband filter. (set -1 for np.inf) (Default: %(default)g)')
parser.add_argument('-u', '--passband-high', type=float, default=80e-3, help='Upper frequency [GHz] of the passband filter. (set -1 for np.inf) (Default: %(default)g)')
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, 'shifts/')
show_plots = args.show_plots
max_k = int(args.max_k)
allowed_ks = np.arange(-max_k, max_k+1, dtype=int)
Xref = args.X
N_runs = args.N_runs
remove_beacon_from_trace = True
apply_signal_window_from_max = True
low_bp = args.passband_low if args.passband_low >= 0 else np.inf # GHz
high_bp = args.passband_high if args.passband_high >= 0 else np.inf # GHz
####
fname_dir = args.data_dir
antennas_fname = path.join(fname_dir, beacon.antennas_fname)
time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_fname
tx_fname = path.join(fname_dir, beacon.tx_fname)
beacon_snr_fname = path.join(fname_dir, beacon.beacon_snr_fname)
## This is a file indicating whether the k-finding algorithm was
## stopped early. This happens when the ks do not change between
## two consecutive iterations.
run_break_fname = path.join(fname_dir, 'ca_breaked_run')
# 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 information
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']
##
## Manipulate time and traces of each antenna
##
### Remove time due to true phase
### and optionally remove the beacon
### Note: there is no use in changing *_AxB variables here (except for plotting),
### they're recomputed by the upcoming rit.set_pol_and_bp call.
measured_repair_offsets = beacon.read_antenna_clock_repair_offsets(ev.antennas, mode='phases', freq_name=freq_name)
for i, ant in enumerate(ev.antennas):
ev.antennas[i].orig_t = ev.antennas[i].t
ev.antennas[i].t += measured_repair_offsets[i]
# t_AxB will be set by the rit.set_pol_and_bp function
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_trace:
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 = ant.beacon_info[freq_name]['amplitude']
calc_beacon = lib.sine_beacon(f, ev.antennas[i].t, amplitude=ampl, 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/max(ant.beacon)
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(ant.t_AxB, ant.E_AxB + calc_beacon, alpha=0.6, ls='dashed', label='Signal') # calc_beacon was already removed
ax.plot(ant.t_AxB, calc_beacon, alpha=0.6, ls='dashed', label='Calc Beacon')
ax.plot(ant.t_AxB, ant.E_AxB, alpha=0.6, label="Signal - Calc Beacon")
ax.legend(title=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 = 200, 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'))
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
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"))
ax.set_xlim(*old_xlim)
fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.traces.A{ant.name}.pdf'))
if show_plots:
plt.show()
# Prepare polarisation and passbands
rit.set_pol_and_bp(ev, low=low_bp, high=high_bp)
# determine allowable ks per location
dt = ev.antennas[0].t_AxB[1] - ev.antennas[0].t_AxB[0]
print("Checking k:", allowed_ks)
##
## Determine grid positions
##
dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0)
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_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_fine = np.linspace(-scale2d, scale2d, 40)
y_fine = np.linspace(-scale2d, scale2d, 40)
## Remove run_break_fname if it exists
try:
os.remove(run_break_fname)
except OSError:
pass
##
## Do calculations on the grid
##
for r in range(N_runs):
# Setup Plane grid to test
if r == 0:
xoff, yoff = 0,0
x = x_coarse
y = y_coarse
else:
# zooming in
# best_idx is defined at the end of the loop
old_ks_per_loc = ks_per_loc[best_idx]
xoff, yoff = locs[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}) at X={Xref}.")
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 tqdm(enumerate(product(x,y)), total=N_loc, ):
tmp_fig_subdir = None
if i % 10 ==0:
if hasattr(tqdm, '__code__') and tqdm.__code__.co-name == '<lambda>':
print(f"Testing location {i} out of {N_loc}")
tmp_fig_subdir = fig_subdir
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}.",
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))
## 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)) )
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_aspect('equal', 'datalim')
sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6)
fig.colorbar(sc, ax=axs, label='Max Amplitude [$\\mu V/m$]')
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)
if fig_dir:
old_xlims = axs.get_xlim()
old_ylims = axs.get_ylim()
fig.tight_layout()
fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.maxima.X{Xref}.run{r}.pdf'))
if False:
axs.plot(tx.x/1e3, tx.y/1e3, marker='X', color='k')
fig.tight_layout()
fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.maxima.X{Xref}.run{r}.with_tx.pdf'))
axs.set_xlim(*old_xlims)
axs.set_ylim(*old_ylims)
fig.tight_layout()
##
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)
## Save best ks to file
np.savetxt(path.join(fig_dir, path.basename(__file__)+f'.bestk.X{Xref}.run{r}.txt'), best_k )
## Do a small reconstruction of the shower for best ks
if True:
print("Reconstructing for best k")
for j in range(2):
power_reconstruction = j==1
if power_reconstruction: # Do power reconstruction
# backup antenna times
backup_times = [ ant.t_AxB for ant in ev.antennas ]
# 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
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)
# 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):
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.legend(title=snr_str)
if fig_dir:
if power_reconstruction:
fname_extra = "power"
else:
fname_extra = "max_amp"
fig.tight_layout()
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() ):
print(f"No changes from previous grid, breaking at iteration {r} out of {N_runs}")
try:
with open(run_break_fname, 'wt', encoding='utf-8') as fp:
fp.write(f"Breaked at grid iteration {r} out of {N_runs}")
except:
pass
break
old_ks_per_loc = ks_per_loc[best_idx]
# 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]
h5beacon_info = h5ant['beacon_info']
# find out freq_name
if freq_name is None:
freq_name = [ k for k in h5beacon_info.keys() if np.isclose(h5beacon_info[k].attrs['freq'], f_beacon)][0]
h5attrs = h5beacon_info[freq_name].attrs
h5attrs['best_k'] = old_ks_per_loc[i]
h5attrs['best_k_time'] = old_ks_per_loc[i]/f_beacon
if show_plots:
plt.show()