m-thesis-introduction/airshower_beacon_simulation/db_longitudinal_figure.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 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')

    from scriptlib import MyArgumentParser
    parser = MyArgumentParser()
    args = parser.parse_args()

    fig_dir = args.fig_dir
    show_plots = args.show_plots

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

    # create fig_dir
    if fig_dir:
        os.makedirs(fig_dir, exist_ok=True)

    # Load res from pickle
    with open(pickle_fname, 'rb') as f:
        res = pickle.load(f)

    ##
    # Longitudinal figures
    ##
    for i in range(2):
        mode = ['grammage', 'distance'][i]

        fig = rit.longitudinal_figure(res.dl[0], res.dX[0], res.profile_rit[0], mode=mode)

        if fig_dir:
            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".{mode}.pdf"))