mirror of
https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git
synced 2024-11-14 02:23:32 +01:00
422 lines
14 KiB
Python
Executable file
422 lines
14 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()
|
|
|
|
fname = "ZH_airshower/mysim.sry"
|
|
|
|
fig_dir = "./figures/periods_from_shower_figures/"
|
|
fig_subdir = path.join(fig_dir, 'shifts/')
|
|
show_plots = False
|
|
|
|
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
|
|
for i, ant in enumerate(ev.antennas):
|
|
clock_phase = ant.beacon_info[freq_name]['sigma_phase_mean']
|
|
clock_phase_time = clock_phase/(2*np.pi*f_beacon)
|
|
|
|
ev.antennas[i].orig_t = ev.antennas[i].t_AxB
|
|
ev.antennas[i].t_AxB += clock_phase_time
|
|
|
|
if remove_beacon_from_trace:
|
|
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()
|
|
|
|
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.zoomed.A{i}.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 not True: #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
|
|
xoff, yoff = 0,0
|
|
if r == 0:
|
|
x = x_coarse
|
|
y = y_coarse
|
|
else:
|
|
# zooming in
|
|
old_ks_per_loc = ks_per_loc[best_idx]
|
|
xoff = xx[best_idx]
|
|
yoff = yy[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(fig_dir + __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, 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)
|
|
|
|
if fig_dir:
|
|
old_xlims = axs.get_xlim()
|
|
old_ylims = axs.get_ylim()
|
|
fig.tight_layout()
|
|
fig.savefig(path.join(fig_dir, __file__+f'.maxima.run{r}.pdf'))
|
|
if True:
|
|
axs.plot(tx.x/1e3, tx.y/1e3, marker='X', color='k')
|
|
fig.tight_layout()
|
|
fig.savefig(path.join(fig_dir, __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(fig_dir + __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")
|
|
_, __, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=scale2d, wy=scale2d, xoff=xoff, yoff=yoff, zgr=0)
|
|
|
|
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')
|
|
sc = axs.scatter(xx/1e3, yy/1e3, c=p, cmap='Spectral_r', alpha=0.6)
|
|
fig.colorbar(sc, ax=axs)
|
|
|
|
if fig_dir:
|
|
fig.tight_layout()
|
|
fig.savefig(path.join(fig_dir, __file__+f'.reconstruction.run{r}.pdf'))
|
|
|
|
# Abort if no improvement
|
|
if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ):
|
|
print("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()
|