ZH: findks: u rewrite function from samples to periods

This makes sure to actually use interpolation to find the best period offset instead of rolling the samples.
The latter won't work when the period is not dividable by the sampling period.
This commit is contained in:
Eric Teunis de Boone 2023-04-28 17:08:53 +02:00
parent 9f610ab2e7
commit 3bc5504c2b

View file

@ -41,9 +41,6 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
else: else:
plot_iteration_with_shifted_trace = [] plot_iteration_with_shifted_trace = []
allowed_sample_shifts = np.rint(allowed_ks * period /dt).astype(int)
sample_shift_first_trace = np.rint(period_shift_first_trace * period/dt).astype(int)
# propagate to test location # propagate to test location
for i, ant in enumerate(antennas): for i, ant in enumerate(antennas):
aloc = [ant.x, ant.y, ant.z] aloc = [ant.x, ant.y, ant.z]
@ -72,14 +69,13 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
t_sum = np.arange(t_min+max_neg_shift, t_max+max_pos_shift, dt) t_sum = np.arange(t_min+max_neg_shift, t_max+max_pos_shift, dt)
a_sum = np.zeros(len(t_sum)) a_sum = np.zeros(len(t_sum))
best_sample_shifts = np.zeros( (len(antennas)) ,dtype=int) best_period_shifts = np.zeros( (len(antennas)) ,dtype=int)
for i, (t_r, E_) in enumerate(zip(t_, a_)): for i, (t_r, E_) in enumerate(zip(t_, a_)):
f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0) f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0)
a_int = f(t_sum)
if i == 0: if i == 0:
a_sum += a_int a_sum += f(t_sum)
best_sample_shifts[i] = sample_shift_first_trace best_period_shifts[i] = period_shift_first_trace
continue continue
# init figure # init figure
@ -90,26 +86,27 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
ax.set_ylabel("Amplitude") ax.set_ylabel("Amplitude")
ax.plot(t_sum, a_sum) ax.plot(t_sum, a_sum)
shift_maxima = np.zeros( len(allowed_sample_shifts) ) # find the maxima for each period shift k
for j, shift in enumerate(allowed_sample_shifts): shift_maxima = np.zeros( len(allowed_ks) )
augmented_a = np.roll(a_int, shift) for j, k in enumerate(allowed_ks):
augmented_a = f(t_sum + k*period)
shift_maxima[j] = np.max(augmented_a + a_sum) shift_maxima[j] = np.max(augmented_a + a_sum)
if i in plot_iteration_with_shifted_trace: 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'{shift}') ax.plot(t_sum, augmented_a, alpha=0.7, ls='dashed', label=f'{k:g}')
# transform maximum into best_sample_shift # transform maximum into best_sample_shift
best_idx = np.argmax(shift_maxima) best_idx = np.argmax(shift_maxima)
best_sample_shifts[i] = allowed_sample_shifts[best_idx] best_period_shifts[i] = allowed_ks[best_idx]
best_augmented_a = np.roll(a_int, best_sample_shifts[i]) best_augmented_a = f(t_sum + k*period)
a_sum += best_augmented_a a_sum += best_augmented_a
# cleanup figure # cleanup figure
if i in plot_iteration_with_shifted_trace: if i in plot_iteration_with_shifted_trace:
if True: # plot best k again 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.plot(t_sum, best_augmented_a, alpha=0.8, label=f'best $k$={best_period_shifts[i]:g}', lw=2)
ax.legend( ncol=5 ) ax.legend( ncol=5 )
if fig_dir: if fig_dir:
@ -126,7 +123,7 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
if True: # zoomed on peak of this trace if True: # zoomed on peak of this trace
x = t_r[np.argmax(E_)] x = t_r[np.argmax(E_)]
wx = 50 + (max(best_sample_shifts) - min(best_sample_shifts))*dt wx = 50 + (max(best_period_shifts) - min(best_period_shifts))*dt
ax.set_xlim(x-wx, x+wx) ax.set_xlim(x-wx, x+wx)
fig.savefig(fname + ".zoomed.peak.pdf") fig.savefig(fname + ".zoomed.peak.pdf")
@ -138,10 +135,10 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
# sort by antenna (undo sorting by maximum) # sort by antenna (undo sorting by maximum)
undo_sort_idx = np.argsort(sort_idx) undo_sort_idx = np.argsort(sort_idx)
best_sample_shifts = best_sample_shifts[undo_sort_idx] best_period_shifts = best_period_shifts[undo_sort_idx]
# Return ks # Return ks
return best_sample_shifts, np.max(a_sum) return best_period_shifts, np.max(a_sum)
if __name__ == "__main__": if __name__ == "__main__":
import sys import sys
@ -413,16 +410,10 @@ if __name__ == "__main__":
yy.append(y_+yoff) yy.append(y_+yoff)
# Find best k for each antenna # Find best k for each antenna
shifts, maximum = find_best_period_shifts_summing_at_location(test_loc, ev.antennas, allowed_ks, period=1/f_beacon, dt=dt, 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], plot_iteration_with_shifted_trace=[ 5, len(ev.antennas)-1],
fig_dir=tmp_fig_subdir, fig_distinguish=f"run{r}.") fig_dir=tmp_fig_subdir, 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) xx = np.array(xx)
yy = np.array(yy) yy = np.array(yy)
locs = list(zip(xx, yy)) locs = list(zip(xx, yy))