#!/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 try: from tqdm import tqdm except: tqdm = lambda x: x 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. """ 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)) a_first = np.zeros(len(t_sum)) best_period_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) if i == 0: 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 if i in plot_iteration_with_shifted_trace: fig, ax = plt.subplots(figsize=figsize) if shower_plane_loc is not None: 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_xlabel("Time [ns]") ax.set_ylabel("Amplitude") #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_xlabel("$k$th Period") ax2.set_ylabel("Summed Amplitude") 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) 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: 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] 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 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, loc='lower center') ax2.legend(title=snr_str) ax3.legend() if fig_dir: fig.tight_layout() fig2.tight_layout() if shower_plane_loc is not None: fname_location = '.sloc{:.1g}-{:.1g}-{:.1g}'.format(*shower_plane_loc) else: fname_location = '.loc{:.1f}-{:.1f}-{:.1f}'.format(*test_loc) fname = path.join(fig_dir, path.basename(__file__) + f'.{fig_distinguish}i{i}' + fname_location) if True: old_xlim = ax.get_xlim() if True: # zoomed on part without peak of this trace 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_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") ax.set_xlim(*old_xlim) fig.savefig(fname + ".pdf") fig2.savefig(fname + ".maxima.pdf") 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), sort_idx if __name__ == "__main__": import sys import os import matplotlib if os.name == 'posix' and "DISPLAY" not in os.environ: matplotlib.use('Agg') plt.rcParams.update({'figure.max_open_warning': 0}) atm = AtmoCal() from scriptlib import MyArgumentParser parser = MyArgumentParser(default_fig_dir="./figures/periods_from_shower_figures/") parser.add_argument('--quick_run', action='store_true', help='Use a very coarse grid (6x6)') parser.add_argument('-X', type=int, default=400, help='Atmospheric depth to start the initial grid.') parser.add_argument('--max-k', type=float, default=2, help='Maximum abs(k) allowed to be shifted. (Default: %(default)d)') parser.add_argument('-N', '--N_runs', type=int, default=5, help='Maximum amount of iterations to grid search. (Default: %(default)d)') parser.add_argument('-l', '--passband-low', type=float, default=30e-3, help='Lower frequency [GHz] of the passband filter. (set -1 for np.inf) (Default: %(default)g)') parser.add_argument('-u', '--passband-high', type=float, default=80e-3, help='Upper frequency [GHz] of the passband filter. (set -1 for np.inf) (Default: %(default)g)') parser.add_argument('--input-fname', type=str, default=None, help='Path to mysim.sry, either directory or path. If empty it takes DATA_DIR and appends mysim.sry. (Default: %(default)s)') args = parser.parse_args() if not args.input_fname: args.input_fname = args.data_dir if path.isdir(args.input_fname): args.input_fname = path.join(args.input_fname, "mysim.sry") 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/') show_plots = args.show_plots max_k = int(args.max_k) allowed_ks = np.arange(-max_k, max_k+1, dtype=int) Xref = args.X N_runs = args.N_runs remove_beacon_from_trace = True apply_signal_window_from_max = True low_bp = args.passband_low if args.passband_low >= 0 else np.inf # GHz high_bp = args.passband_high if args.passband_high >= 0 else np.inf # GHz #### fname_dir = args.data_dir antennas_fname = path.join(fname_dir, beacon.antennas_fname) time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_fname tx_fname = path.join(fname_dir, beacon.tx_fname) beacon_snr_fname = path.join(fname_dir, beacon.beacon_snr_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) _, __, txdata = beacon.read_tx_file(tx_fname) # Read original REvent ev = REvent(args.input_fname) # .. patch in our antennas 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']: .2g}" # 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'] ## ## Manipulate time and traces of each antenna ## ### Remove time due to true phase ### and optionally remove the beacon ### Note: there is no use in changing *_AxB variables here (except for plotting), ### they're recomputed by the upcoming rit.set_pol_and_bp call. 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 ev.antennas[i].t += measured_repair_offsets[i] # t_AxB will be set by the rit.set_pol_and_bp function ev.antennas[i].t_AxB += measured_repair_offsets[i] if apply_signal_window_from_max: N_pre, N_post = 250, 250 # TODO: make this configurable # Get max idx from all the traces # and select the strongest max_idx = [] maxs = [] for trace in [ant.Ex, ant.Ey, ant.Ez]: idx = np.argmax(np.abs(trace)) max_idx.append(idx) maxs.append( np.abs(trace[idx]) ) idx = np.argmax(maxs) max_idx = max_idx[idx] # Create window around max_idx low_idx = max(0, max_idx-N_pre) high_idx = min(len(ant.t), max_idx+N_post) ev.antennas[i].orig_t = ant.orig_t[low_idx:high_idx] ev.antennas[i].t = ant.t[low_idx:high_idx] ev.antennas[i].Ex = ant.Ex[low_idx:high_idx] ev.antennas[i].Ey = ant.Ey[low_idx:high_idx] ev.antennas[i].Ez = ant.Ez[low_idx:high_idx] ev.antennas[i].t_AxB = ant.t_AxB[low_idx:high_idx] ev.antennas[i].E_AxB = ant.E_AxB[low_idx:high_idx] # .. and remove the beacon from the traces # Note: ant.E_AxB is recalculated by rit.set_pol_and_bp if remove_beacon_from_trace: clock_phase = measured_repair_offsets[i]*2*np.pi*f_beacon beacon_phase = ant.beacon_info[freq_name]['beacon_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, amplitude=ampl, phase=beacon_phase-clock_phase) tx_amps = txdata['amplitudes'] tx_amps_sum = np.sum(tx_amps) # Split up contribution to the various polarisations for j, amp in enumerate(tx_amps): if j == 0: ev.antennas[i].Ex -= amp*(1/tx_amps_sum)*calc_beacon elif j == 1: ev.antennas[i].Ey -= amp*(1/tx_amps_sum)*calc_beacon elif j == 2: ev.antennas[i].Ez -= amp*(1/tx_amps_sum)*calc_beacon # Subtract the beacon from E_AxB ev.antennas[i].E_AxB -= calc_beacon # Make a figure of the manipulated traces if i == 72: orig_beacon_amplifier = ampl/max(ant.beacon) fig, ax = plt.subplots(figsize=figsize) ax.set_title(f"Signal and Beacon traces Antenna {ant.name}") 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(title=snr_str) # save if fig_dir: fig.tight_layout() if True: # zoom old_xlim = ax.get_xlim() 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')) if True: # zoomed on peak of this trace idx = np.argmax(ev.antennas[i].E_AxB) x = ev.antennas[i].t_AxB[idx] 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")) ax.set_xlim(*old_xlim) fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.traces.A{ant.name}.pdf')) if show_plots: plt.show() # Prepare polarisation and passbands rit.set_pol_and_bp(ev, low=low_bp, high=high_bp) # determine allowable ks per location dt = ev.antennas[0].t_AxB[1] - ev.antennas[0].t_AxB[0] print("Checking k:", allowed_ks) ## ## Determine grid positions ## 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(-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, 14) y_coarse = np.linspace(-scale4d, scale4d, 14) x_fine = np.linspace(-scale2d, scale2d, 18) y_fine = np.linspace(-scale2d, scale2d, 18) ## 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}) at X={Xref}.") 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) # 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 == 0: if hasattr(tqdm, '__code__') and tqdm.__code__.co-name == '': print(f"Testing location {i} out of {N_loc}") tmp_fig_subdir = fig_subdir test_loc = loc[0]* ev.uAxB + loc[1]*ev.uAxAxB + dXref *ev.uA # Find best k for each antenna 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), ) 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("-v x B [km]") axs.set_aspect('equal', 'datalim') 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) # indicate maximum value idx = np.argmax(maxima_per_loc) 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() old_ylims = axs.get_ylim() fig.tight_layout() fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.maxima.X{Xref}.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.X{Xref}.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], ": ", 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 ) ## 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])/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 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("-v x B [km]") 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: fname_extra = "power" else: fname_extra = "max_amp" fig.tight_layout() 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] - first_trace_period_shift).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]/f_beacon if show_plots: plt.show()