ericteunis
/
m-thesis-introduction
mirror of https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

ZH: ca_period_from_shower as used for Thesis Figures

This commit is contained in:
Eric Teunis de Boone 2023-10-31 16:47:45 +01:00
parent f50b5cb308
commit 42236b03a8
1 changed files with 166 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

226
airshower_beacon_simulation/ca_period_from_shower.py
View File

@ -25,7 +25,25 @@ try:
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):
try:
from joblib import Parallel, delayed
except:
Parallel = None
delayed = lambda x: x
def find_best_period_shifts_at_location(*args, algo=None, **kwargs):
"""
This is a placeholder function.
For args and kwargs see find_best_period_shifts_summing_at_location.
"""
if algo is None:
algo = 'sum'
return find_best_period_shifts_summing_at_location(*args, **kwargs, algo=algo)
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, algo='sum', verbose=False):
"""
Find the best sample_shift for each antenna by summing the antenna traces
and seeing how to get the best alignment.
@ -80,9 +98,9 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
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
a_first = f(t_sum - period_shift_first_trace*period)
a_sum += a_first
continue
# init figure
@ -92,47 +110,67 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
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_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)
#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_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)
ax2.plot(0, np.max(a_first), marker='*', label='first trace', ls='none', ms=20)
ax3 = ax2.twinx()
ax3.set_ylabel("Correlation")
# find the maxima for each period shift k
shift_maxima = np.zeros( len(allowed_ks) )
shift_corrs = np.zeros_like(shift_maxima)
for j, k in enumerate(allowed_ks):
augmented_a = f(t_sum + k*period)
augmented_a = f(t_sum - k*period)
shift_maxima[j] = np.max(augmented_a + a_first)
shift_corrs[j] = np.dot(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}')
l = ax.plot(t_sum, a_first + augmented_a, alpha=0.7, ls='dashed', label=f'{k:g}')
ax.axhline(shift_maxima[j], ls='dashdot', color=l[0].get_color(), alpha=0.7)
ax2.plot(k, shift_maxima[j], marker='o', ls='none', ms=20)
ax3.plot(k, shift_corrs[j], marker='3', ls='none', ms=20)
# transform maximum into best_sample_shift
best_idx = np.argmax(shift_maxima)
best_corr_idx = np.argmax(shift_corrs)
best_period_shifts[i] = allowed_ks[best_idx]
best_augmented_a = f(t_sum + k*period)
if verbose and (best_corr_idx != best_idx):
print("Correlation idx not equal to maximum idx")
print(best_corr_idx, best_idx)
if algo == 'sum':
best_period_shifts[i] = allowed_ks[best_idx]
elif algo == 'corr':
best_period_shifts[i] = allowed_ks[best_corr_idx]
k = best_period_shifts[i]
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)
l = ax.plot(t_sum, a_first + best_augmented_a, alpha=0.8, label=f'best $k$={best_period_shifts[i]:g}', lw=2)
ax.axhline(shift_maxima[j], ls='dashdot', color=l[0].get_color(), alpha=0.7)
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}')
ax3.plot(allowed_ks[best_corr_idx], shift_maxima[best_idx], marker='*', ls='none', ms=20, label=f'best corr $k$={allowed_ks[best_corr_idx]:g}')
ax.legend(title='period shift $k$; '+snr_str, ncol=5 )
ax.legend(title='period shift $k$; '+snr_str, ncol=5, loc='lower center')
ax2.legend(title=snr_str)
ax3.legend()
if fig_dir:
fig.tight_layout()
fig2.tight_layout()
@ -147,14 +185,15 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
old_xlim = ax.get_xlim()
if True: # zoomed on part without peak of this trace
wx = 100
wx = 200
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
x = t_sum[np.argmax(a_first)]
x = t_sum[np.argmax(f(t_sum))]
wx = 50 + (max(best_period_shifts) - min(best_period_shifts) )*dt + 1*period
ax.set_xlim(x-wx, x+wx)
fig.savefig(fname + ".zoomed.peak.pdf")
@ -165,13 +204,30 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
plt.close(fig)
plt.close(fig2)
if True: # final summed waveform
fig, ax = plt.subplots()
ax.set_title("Summed Traces with best k's at {}".format(title_location))
ax.set_xlabel("Time [ns]")
ax.set_ylabel("Amplitude")
ax.plot(t_sum, a_sum)
if fig_dir:
fig.tight_layout()
fname = path.join(fig_dir, path.basename(__file__) + f'.{fig_distinguish}i{i}' + fname_location)
fig.savefig(fname + ".sum.pdf")
plt.close(fig)
# 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)
return best_period_shifts, np.max(a_sum), sort_idx
if __name__ == "__main__":
import sys
@ -204,7 +260,16 @@ if __name__ == "__main__":
if path.isdir(args.input_fname):
args.input_fname = path.join(args.input_fname, "mysim.sry")
figsize = (12,8)
figsize = (6,4)
if True:
from matplotlib import rcParams
#rcParams["text.usetex"] = True
rcParams["font.family"] = "serif"
rcParams["font.size"] = "14"
rcParams["grid.linestyle"] = 'dotted'
rcParams["figure.figsize"] = figsize
figsize = rcParams['figure.figsize']
fig_dir = args.fig_dir
fig_subdir = path.join(fig_dir, 'shifts/')
@ -248,7 +313,7 @@ if __name__ == "__main__":
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}"
snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .2g}"
# For now only implement using one freq_name
freq_names = antennas[0].beacon_info.keys()
@ -349,8 +414,8 @@ if __name__ == "__main__":
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)
if not True: # zoomed on part without peak of this trace
wx, x = 100, 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'))
@ -358,7 +423,7 @@ if __name__ == "__main__":
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
wx = 150
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"))
@ -380,22 +445,23 @@ if __name__ == "__main__":
## Determine grid positions
##
dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0)
zgr = 0 + ev.core[2]
dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,zgr)
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_coarse = np.linspace(-scale4d, scale4d, 16)
y_coarse = np.linspace(-scale4d, scale4d, 16)
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_coarse = np.linspace(-scale4d, scale4d, 14)
y_coarse = np.linspace(-scale4d, scale4d, 14)
x_fine = np.linspace(-scale2d, scale2d, 40)
y_fine = np.linspace(-scale2d, scale2d, 40)
x_fine = np.linspace(-scale2d, scale2d, 18)
y_fine = np.linspace(-scale2d, scale2d, 18)
## Remove run_break_fname if it exists
try:
@ -435,9 +501,29 @@ if __name__ == "__main__":
yy = []
N_loc = len(maxima_per_loc)
for i, (x_, y_) in tqdm(enumerate(product(x,y)), total=N_loc, ):
# Make the grid a list of locations
xx = []
yy = []
for i, (x_, y_) in enumerate(product(x,y)):
xx.append( x_+xoff )
yy.append( y_+yoff )
xx = np.array(xx)
yy = np.array(yy)
locs = list(zip(xx, yy))
# Shift the reference trace
if r == 0:
first_trace_period_shift = 0
else:
first_trace_period_shift = 0#first_trace_period_shift + np.rint(np.mean(old_ks_per_loc))
print("New first trace period:", first_trace_period_shift)
# define loop func for joblib
def loop_func(loc, dXref=dXref, i=1):
tmp_fig_subdir = None
if i % 10 ==0:
if i == 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
@ -445,33 +531,41 @@ if __name__ == "__main__":
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}.",
return find_best_period_shifts_at_location(test_loc, ev.antennas, allowed_ks, period=1/f_beacon, dt=dt, period_shift_first_trace=first_trace_period_shift,
plot_iteration_with_shifted_trace=[ 1, 2, 3, 4, 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))
res = ( delayed(loop_func)(loc, i=i) for i, loc in enumerate(locs) )
if Parallel:
res = Parallel(n_jobs=None)(tqdm(res, total=len(locs)))
else:
res = tqdm(res, total=len(locs))
# unpack loop results
ks_per_loc, maxima_per_loc, sort_idx = zip(*res)
## 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)) )
scatter_kwargs = dict(cmap='Spectral_r', s=64*4, alpha=0.6)
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_title(f"Optimizing signal strength by varying $k$ per antenna,\n Grid Run {r}")
axs.set_ylabel(" vxvxB [km]")
axs.set_xlabel("-v x B [km]")
axs.set_aspect('equal', 'datalim')
sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6)
sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, **scatter_kwargs)
fig.colorbar(sc, ax=axs, label='Max Amplitude [$\\mu V/m$]')
axs.legend(title=snr_str)
#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)
axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) # max value
axs.plot(0,0, 'r+', ms=30) # true axis
if fig_dir:
old_xlims = axs.get_xlim()
@ -490,8 +584,10 @@ if __name__ == "__main__":
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)
print("Max at location: ", locs[best_idx], ": ", maxima_per_loc[best_idx])
print('Best k:', best_k[sort_idx[best_idx]]) # Sorted by DESC amplitude
print('Mean best k:', np.rint(np.mean(best_k)))
print('first k:', first_trace_period_shift)
## Save best ks to file
np.savetxt(path.join(fig_dir, path.basename(__file__)+f'.bestk.X{Xref}.run{r}.txt'), best_k )
@ -509,33 +605,43 @@ if __name__ == "__main__":
# 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
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)
xx, yy, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=(x[-1]-x[0])/2, wy=(y[-1]-y[0])/2, 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):
for i, loc in tqdm(enumerate(locs), total=len(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.set_title(f"Shower slice for best k, Grid Run {r}")
axs.set_ylabel(" vxvxB [km]")
axs.set_xlabel("-v x B [km]")
axs.legend(title=snr_str)
if power_reconstruction:
sc_c = p
sc_label = 'Power [$(\\mu V/m)^2$]'
else:
sc_c = maxima
sc_label='Max Amplitude [$\\mu V/m$]'
sc = axs.scatter(xx/1e3, yy/1e3, c=sc_c, **scatter_kwargs)
fig.colorbar(sc, ax=axs, label=sc_label)
# indicate maximum value
idx = np.argmax(p if power_reconstruction else maxima)
axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) # max value
axs.plot(0,0, 'r+', ms=30) # true axis
# make square figure
axs.set_aspect('equal', 'datalim')
#axs.legend(title=snr_str)
if fig_dir:
if power_reconstruction:
@ -547,7 +653,7 @@ if __name__ == "__main__":
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() ):
if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx] - first_trace_period_shift).all() ):
print(f"No changes from previous grid, breaking at iteration {r} out of {N_runs}")
try: