ZH: improve cc reconstruction script from ffb1fa8

simulations/airshower_beacon_simulation/cc_power_along_axis_reconstruction.py

@@ -7,8 +7,10 @@ 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
 
 from earsim import REvent
 from atmocal import AtmoCal
@@ -31,12 +33,19 @@ if __name__ == "__main__":
     fig_dir = "./figures/"
     fig_subdir = path.join(fig_dir, 'reconstruction')
     show_plots = False
+    pickle_fname = path.join(fig_subdir, 'res.pkl')
 
     ####
     fname_dir = path.dirname(fname)
     antennas_fname = path.join(fname_dir, beacon.antennas_fname)
     time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_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)
     # Read original REvent
@@ -44,6 +53,14 @@ if __name__ == "__main__":
     # .. patch in our antennas
     ev.antennas = antennas
 
+    # 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
     for i, ant in enumerate(ev.antennas):
         clock_phase_time = ant.beacon_info[freq_name]['sigma_phase_mean']/(2*np.pi*f_beacon)
@@ -52,12 +69,12 @@ if __name__ == "__main__":
         total_clock_time = best_k_time + clock_phase_time
         ev.antennas[i].t += total_clock_time
 
-    ## Start a pickle
-    with open(pickle_fname, 'wb') as fp:
     N_X, Xlow, Xhigh = 23, 100, 1200
-        res = rit.reconstruction(ev, outfile=fig_subdir+'/fig.pdf', slice_outdir=fig_subdir, Xlow=Xlow, N_X=N_X, Xhigh=Xhigh)
+    res = rit.reconstruction(ev, outfile=fig_subdir+'/fig.pdf', slice_outdir=fig_subdir+'/', Xlow=Xlow, N_X=N_X, Xhigh=Xhigh)
 
+    ## Save a pickle
+    with open(pickle_fname, 'wb') as fp:
         pickle.dump(res,fp)
 
-    if show_plots():
+    if show_plots:
         plt.show()