mirror of
https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git
synced 2024-12-22 03:23:34 +01:00
ZH: ca_period_from_shower as used for Thesis Figures
This commit is contained in:
parent
f50b5cb308
commit
42236b03a8
1 changed files with 166 additions and 60 deletions
|
|
@ -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)
|
||||
|
||||
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]
|
||||
best_augmented_a = f(t_sum + k*period)
|
||||
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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue