m-thesis-introduction/simulations/airshower_beacon_simulation/reconstruction.py

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

import h5py
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # required for projection='3d' on old matplotliblib versions
import numpy as np
from copy import deepcopy

import aa_generate_beacon as beacon
import lib
from lib import rit
from earsim import REvent, Antenna
from atmocal import AtmoCal
import pickle

if __name__ == "__main__":
    from os import path
    import sys

    fname = "ZH_airshower/mysim.sry"
    dirname = '/scratch/etdeboone/reconstructions/'+  (sys.argv[2]+'/' if 2 in sys.argv else '')
    print("Dirname:", dirname)

    ####
    fname_dir = path.dirname(fname)
    max_clock_offset = None
    if True: # use the original files
        ev = REvent(fname)

        max_clock_offset = int(sys.argv[1]) if len(sys.argv) > 1 else 5# ns
        # randomise timing for antenna copies
        seed = 12345
        rng = np.random.default_rng(seed)

        dirname += f'rand{max_clock_offset}ns'
    else:
        raise NotImplementedError
        antennas_fname = path.join(fname_dir, beacon.antennas_fname)

        if not path.isfile(antennas_fname):
            print("Antenna file cannot be found, did you try generating a beacon?")
            sys.exit(1)

    # Full reconstruction?
    if False:
        print("Full reconstruction")

        if max_clock_offset:
            print(f"Randomising timing: {max_clock_offset}")
            clock_offsets = max_clock_offset * (2*rng.uniform(size=len(ev.antennas)) - 1)
            for i, ant in enumerate(ev.antennas):
                ev.antennas[i].t += clock_offsets[i]
        else:
            print(f"Not randomising timing")

        with open(dirname + '/res.pkl', 'wb') as fp:
            res = rit.reconstruction(ev, outfile=dirname+'/fig.pdf', slice_outdir=dirname+'/', Xlow=100, N_X=23, Xhigh=1200)

            pickle.dump(res, fp)
        sys.exit(0)

    # Only a single slice
    else:
        print("Shower plane reconstruction")

#        atm = AtmoCal()
#        X = 750
#        rit.set_pol_and_bp(ev)
#        dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),750,0)
#        scale2d = dXref*np.tan(np.deg2rad(2.))

        if max_clock_offset is None:
            print(" + randomised timing")

            ants_copy = deepcopy(ev.antennas)
            clock_offsets = max_clock_offset * (2*rng.uniform(size=len(ants_copy)) - 1)
            for i, ant in enumerate(ants_copy):
                ant.t -= clock_offsets[i]

        fig, axs = plt.subplots(1,1+(max_clock_offset is not None), sharex=True, sharey=True)
        fig.suptitle(f"Reconstructed Shower plane slice at X={X}")
        if max_clock_offset is None:
            axs = [ax]
        for ax in axs:
            ax.set_xlabel("-v x B (m)")
            ax.set_ylabel(" vxvxB (m)")

        if max_clock_offset is not None:
            axs[0].set_title("Simulated timing")
            axs[1].set_title(f"Randomised timing [0, {max_clock_offset}]")

        try:
            for i, ax in enumerate(axs):
                # modify timing
                if i == 1:
                    ev.antennas = ants_copy

                xx,yy,p,km= rit.shower_plane_slice(ev,X,21,21,scale2d,scale2d)

                if i == 0:
                    vmax, vmin = max(p), min(p)

                sc = ax.scatter(xx,yy,c=p, vmax=vmax, vmin=vmin)

                # Plot centers
                im = np.argmax(p)

                ax.plot(0, 0,'r+',ms=30)
                ax.plot(xx[im],yy[im],'bx',ms=30)

            fig.colorbar(sc, ax=axs[0])
        # Always make sure to show the plot
        except:
            if True:
                fig.savefig(dirname + f"/shower_plane_slice_X{X}_T{max_clock_offset}_exception.pdf")
            plt.show()

        if True:
            fig.savefig(dirname + f"/shower_plane_slice_X{X}_T{max_clock_offset}.pdf")

    plt.show()