mirror of
https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git
synced 2024-11-14 02:23:32 +01:00
Eric Teunis de Boone
daf2209c73
This is moved out from the scripts that already have to modify timing
470 lines
17 KiB
Python
Executable file
470 lines
17 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
|
|
|
|
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 = []
|
|
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))
|
|
|
|
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 in plot_iteration_with_shifted_trace:
|
|
fig, ax = plt.subplots()
|
|
ax.set_title("Traces at ({:.1f},{:.1f},{:.1f}) i={i}/{tot}".format(*test_loc, i=i, tot=N_ant))
|
|
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 in plot_iteration_with_shifted_trace:
|
|
ax.plot(t_sum, augmented_a, alpha=0.7, ls='dashed', label=f'{shift}')
|
|
|
|
# transform maximum into best_sample_shift
|
|
best_idx = np.argmax(shift_maxima)
|
|
|
|
best_sample_shifts[i] = allowed_sample_shifts[best_idx]
|
|
best_augmented_a = np.roll(a_int, best_sample_shifts[i])
|
|
a_sum += best_augmented_a
|
|
|
|
# cleanup figure
|
|
if i in plot_iteration_with_shifted_trace:
|
|
if True: # plot best k again
|
|
ax.plot(t_sum, augmented_a, alpha=0.8, label=f'best k={best_sample_shifts[i]}', lw=2)
|
|
|
|
ax.legend( ncol=5 )
|
|
if fig_dir:
|
|
fig.tight_layout()
|
|
fname = path.join(fig_dir, path.basename(__file__) + f'.{fig_distinguish}i{i}' + '.loc{:.1f}-{:.1f}-{:.1f}'.format(*test_loc))
|
|
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_sample_shifts) - min(best_sample_shifts))*dt
|
|
ax.set_xlim(x-wx, x+wx)
|
|
fig.savefig(fname + ".zoomed.peak.pdf")
|
|
|
|
ax.set_xlim(*old_xlim)
|
|
|
|
fig.savefig(fname + ".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()
|
|
|
|
from scriptlib import MyArgumentParser
|
|
parser = MyArgumentParser(default_fig_dir="./figures/periods_from_shower_figures/")
|
|
args = parser.parse_args()
|
|
|
|
fname = "ZH_airshower/mysim.sry"
|
|
|
|
fig_dir = args.fig_dir
|
|
fig_subdir = path.join(fig_dir, 'shifts/')
|
|
show_plots = args.show_plots
|
|
|
|
allowed_ks = [ -2, -1, 0, 1, 2]
|
|
Xref = 400
|
|
|
|
N_runs = 3
|
|
remove_beacon_from_trace = True
|
|
|
|
####
|
|
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
|
|
|
|
## 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)
|
|
# 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']
|
|
|
|
# Prepare polarisation and passbands
|
|
rit.set_pol_and_bp(ev, low=0.03, high=0.08)
|
|
|
|
# Remove time due to true phase
|
|
# and optionally remove the beacon
|
|
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_AxB
|
|
ev.antennas[i].t_AxB += measured_repair_offsets[i]
|
|
|
|
if remove_beacon_from_trace:
|
|
clock_phase = measured_repair_offsets[i]*2*np.pi*f_beacon
|
|
meas_phase = ant.beacon_info[freq_name]['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_AxB, amplitude=ampl, phase=meas_phase-clock_phase)
|
|
|
|
ev.antennas[i].E_AxB -= calc_beacon
|
|
|
|
# Make a figure of the manipulated traces
|
|
if i == 2:
|
|
orig_beacon_amplifier = ampl/max(ant.beacon)
|
|
|
|
fig, ax = plt.subplots()
|
|
ax.set_title(f"Signal and Beacon traces Antenna {i}")
|
|
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()
|
|
|
|
# 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, 0#ant.t_AxB[np.argmax(ant.E_AxB)]
|
|
ax.set_xlim(x-wx, x+wx)
|
|
|
|
fig.savefig(path.join(fig_dir, __file__+f'.traces.A{i}.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 = 100
|
|
ax.set_xlim(x-wx, x+wx)
|
|
fig.savefig(path.join(fig_dir, __file__+f".traces.A{i}.zoomed.peak.pdf"))
|
|
|
|
ax.set_xlim(*old_xlim)
|
|
|
|
fig.savefig(path.join(fig_dir, __file__+f'.traces.A{i}.pdf'))
|
|
|
|
if show_plots:
|
|
plt.show()
|
|
|
|
|
|
# determine allowable ks per location
|
|
dt = ev.antennas[0].t_AxB[1] - ev.antennas[0].t_AxB[0]
|
|
allowed_sample_shifts = np.rint(allowed_ks/f_beacon /dt).astype(int)
|
|
print("Checking:", allowed_ks, ": shifts :", allowed_sample_shifts)
|
|
|
|
##
|
|
## 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 False: #quicky
|
|
x_coarse = np.linspace(-scale2d, scale2d, 4)
|
|
y_coarse = np.linspace(-scale2d, scale2d, 4)
|
|
|
|
x_fine = x_coarse/4
|
|
y_fine = y_coarse/4
|
|
else: # long
|
|
N_runs = 5
|
|
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})")
|
|
|
|
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=[ 5, 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(path.join(fig_dir, path.basename(__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"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$]')
|
|
|
|
# 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.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.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.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()
|
|
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$]')
|
|
|
|
|
|
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.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]*dt/f_beacon
|
|
|
|
if show_plots:
|
|
plt.show()
|