m-thesis-introduction/airshower_beacon_simulation/ca_period_from_shower.py

#!/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 == '<lambda>':
                    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()