ZH: restructure reconstruction script

simulations/airshower_beacon_simulation/do_reconstruction.py

@@ -24,6 +24,7 @@ if __name__ == "__main__":
 
     ####
     fname_dir = path.dirname(fname)
+    max_clock_offset = None
     if True: # use the original files
         ev = REvent(fname)
 
@@ -33,47 +34,61 @@ if __name__ == "__main__":
         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)
 
-        # Full reconstruction?
+        if not path.isfile(antennas_fname):
-        if False:
+            print("Antenna file cannot be found, did you try generating a beacon?")
-            print("Full reconstruction")
+            sys.exit(1)
 
-            if max_clock_offset:
+    # Full reconstruction?
-                print(f"Randomising timing: {max_clock_offset}")
+    if False:
-                clock_offsets = max_clock_offset * (2*rng.uniform(size=len(ev.antennas)) - 1)
+        print("Full reconstruction")
-                for i, ant in enumerate(ev.antennas):
-                    ant.t -= clock_offsets[i]
-            else:
-                print(f"Not randomising timing")
 
-            with open(dirname + '/res.pkl', 'wb') as fp:
+        if max_clock_offset:
-                res = rit.reconstruction(ev, outfile=dirname+'/fig.pdf', slice_outdir=dirname+'/', Xlow=100, N_X=15, Xhigh=800)
+            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")
 
-                pickle.dump(res, fp)
+        with open(dirname + '/res.pkl', 'wb') as fp:
-            sys.exit(0)
+            res = rit.reconstruction(ev, outfile=dirname+'/fig.pdf', slice_outdir=dirname+'/', Xlow=100, N_X=23, Xhigh=1200)
 
-        # Else do single slice for original timing
+            pickle.dump(res, fp)
-        # and randomised timing
+        sys.exit(0)
-        print("Shower plane reconstruction + randomised timing")
 
-        atm = AtmoCal()
+    # Only a single slice
-        X = 750
+    else:
-        rit.set_pol_and_bp(ev)
+        print("Shower plane reconstruction")
-        dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),X,0)
-        scale2d = dXref*np.tan(np.deg2rad(2.))
 
-        ants_copy = deepcopy(ev.antennas)
+#        atm = AtmoCal()
-        clock_offsets = max_clock_offset * (2*rng.uniform(size=len(ants_copy)) - 1)
+#        X = 750
-        for i, ant in enumerate(ants_copy):
+#        rit.set_pol_and_bp(ev)
-            ant.t -= clock_offsets[i]
+#        dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),750,0)
+#        scale2d = dXref*np.tan(np.deg2rad(2.))
 
-        fig, axs = plt.subplots(1,2, sharex=True, sharey=True)
+        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)")
-        axs[0].set_title("Simulated timing")
+
-        axs[1].set_title(f"Randomised timing [0, {max_clock_offset}]")
+        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):
@@ -97,19 +112,11 @@ if __name__ == "__main__":
             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")
 
-
-    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)
-
-
     plt.show()