m-thesis-introduction/airshower_beacon_simulation/da_reconstruction.py

#!/usr/bin/env python3
# vim: fdm=indent ts=4

"""
Do a reconstruction of airshower after correcting for the
clock offsets.
"""

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # required for projection='3d' on old matplotliblib versions
import numpy as np
from os import path
import pickle
import joblib

from earsim import REvent
from atmocal import AtmoCal
import aa_generate_beacon as beacon
import lib
from lib import rit


if __name__ == "__main__":
    import sys
    import os
    import matplotlib
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        matplotlib.use('Agg')

    atm = AtmoCal()

    from scriptlib import MyArgumentParser
    parser = MyArgumentParser()
    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 = (12,8)

    fig_dir = args.fig_dir
    fig_subdir = path.join(fig_dir, 'reconstruction')
    show_plots = args.show_plots

    apply_signal_window_from_max = True
    remove_beacon_from_traces = True

    ####
    fname_dir = args.data_dir
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    pickle_fname = path.join(fname_dir, 'res.pkl')
    tx_fname = path.join(fname_dir, beacon.tx_fname)
    beacon_snr_fname = path.join(fname_dir, beacon.beacon_snr_fname)

    # 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 info
    beacon_snrs = beacon.read_snr_file(beacon_snr_fname)
    snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .1e}"

    # 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']

    # Repair clock offsets with the measured offsets
    measured_repair_offsets = beacon.read_antenna_clock_repair_offsets(ev.antennas, mode='phases', freq_name=freq_name)
    for i, ant in enumerate(ev.antennas):
        # t_AxB will be set by the rit.set_pol_and_bp function
        ev.antennas[i].orig_t = ev.antennas[i].t
        ev.antennas[i].t += measured_repair_offsets[i]
        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_traces:
            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_AxB = ant.beacon_info[freq_name]['amplitude']
            calc_beacon = lib.sine_beacon(f, ev.antennas[i].t, amplitude=ampl_AxB, 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_AxB/max(ant.beacon)

                for k in range(2):
                    if k == 0:
                        time = ant.t_AxB
                        trace = ant.E_AxB
                        tmp_beacon = calc_beacon
                        fname_extra = ""
                    else:
                        time = ant.t
                        trace = ant.Ex
                        tmp_beacon = tx_amps[0]/tx_amps_sum * calc_beacon
                        fname_extra = ".Ex"

                    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(time, trace + tmp_beacon, alpha=0.6, ls='dashed', label='Signal') # calc_beacon was already removed
                    ax.plot(time, tmp_beacon,        alpha=0.6, ls='dashed', label='Calc Beacon')
                    ax.plot(time, trace,             alpha=0.6,              label="Signal - Calc Beacon")

                    if k == 0:
                        ax.legend(title=snr_str)
                    else:
                        ax.legend(title="Ex " + snr_str)

                    # save
                    if fig_dir:
                        fig.tight_layout()

                        if True: # zoom
                            old_xlim = ax.get_xlim()

                            if True: # zoomed on part without peak of this trace
                                wx, x = 100, 100
                                ax.set_xlim(x-wx, x+wx)

                                fig.savefig(path.join(fig_dir, __file__+f'.traces.A{ant.name}.zoomed.beacon{fname_extra}.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 = 100
                                ax.set_xlim(x-wx, x+wx)
                                fig.savefig(path.join(fig_dir, __file__+f".traces.A{ant.name}.zoomed.peak{fname_extra}.pdf"))

                            ax.set_xlim(*old_xlim)

                        fig.savefig(path.join(fig_dir, __file__+f'.traces.A{i}.pdf'))


    N_X, Xlow, Xhigh = 23, 100, 1200
    with joblib.parallel_backend("loky"):
        res = rit.reconstruction(ev, outfile=fig_subdir+'/fig.pdf', slice_outdir=fig_subdir+'/', Xlow=Xlow, N_X=N_X, Xhigh=Xhigh, disable_pol_and_bp=True)

    ## Save a pickle
    with open(pickle_fname, 'wb') as fp:
        pickle.dump(res,fp)

    if show_plots:
        plt.show()