ZH: Periods from shower saves figures better

simulations/airshower_beacon_simulation/ca_period_from_shower.py

@@ -18,7 +18,7 @@ import aa_generate_beacon as beacon
 import lib
 from lib import rit
 
-def find_best_sample_shifts_summing_at_location(test_loc, antennas, allowed_sample_shifts, dt=None, sample_shift_first_trace=0):
+def find_best_sample_shifts_summing_at_location(test_loc, antennas, allowed_sample_shifts, dt=None, sample_shift_first_trace=0, plot_iteration_with_shifted_trace=None):
     """
     Find the best sample_shift for each antenna by summing the antenna traces
     and seeing how to get the best alignment.
@@ -63,12 +63,22 @@ def find_best_sample_shifts_summing_at_location(test_loc, antennas, allowed_samp
             best_sample_shifts[i] = sample_shift_first_trace
             continue
 
+        if i == plot_iteration_with_shifted_trace:
+            fig, ax = plt.subplots()
+            ax.set_title("Traces at ({:.1f},{:.1f},{:.1f})".format(*test_loc))
+            ax.set_xlabel("Time [ns]")
+            ax.set_ylabel("Amplitude")
+            ax.plot(t_sum, a_sum)
+
         shift_maxima = np.zeros( len(allowed_sample_shifts) )
         for j, shift in enumerate(allowed_sample_shifts):
             augmented_a = np.roll(a_int, shift)
 
             shift_maxima[j] = np.max(augmented_a + a_sum)
 
+            if i == plot_iteration_with_shifted_trace:
+                ax.plot(t_sum, augmented_a, label=f'{shift} shifted')
+
         # transform maximum into best_sample_shift
         best_idx = np.argmax(shift_maxima)
         best_sample_shifts[i] = allowed_sample_shifts[best_idx]
@@ -85,6 +95,8 @@ if __name__ == "__main__":
 
     fname = "ZH_airshower/mysim.sry"
 
+    savepath = "./periods_from_shower_figures/"
+
     allowed_ks = np.arange(-2, 3, 1)
     Xref = 400
 
@@ -94,6 +106,8 @@ if __name__ == "__main__":
     x_fine = np.linspace(-2e3, 2e3, 30)
     y_fine = np.linspace(-2e3, 2e3, 30)
 
+    N_runs = 3
+
     ####
     fname_dir = path.dirname(fname)
     antennas_fname = path.join(fname_dir, beacon.antennas_fname)
@@ -114,7 +128,7 @@ if __name__ == "__main__":
     freq_name = next(iter(freq_names))
     f_beacon = ev.antennas[0].beacon_info[freq_name]['freq']
 
-    # determine best ks per location
+    # determine allowable ks per location
     dt = ev.antennas[0].t[1] - ev.antennas[0].t[0]
     allowed_sample_shifts = np.rint(allowed_ks/f_beacon /dt).astype(int)
     print("Checking:", allowed_ks, ": shifts :", allowed_sample_shifts)
@@ -134,12 +148,15 @@ if __name__ == "__main__":
     scale2d = dXref*np.tan(np.deg2rad(2.))
 
     # Setup Plane grid to test
-    for r in range(6):
+    for r in range(N_runs):
         xoff, yoff = 0,0
         if r == 0:
             x = x_coarse
             y = y_coarse
         else:
+            old_ks_per_loc = ks_per_loc[best_idx]
+            xoff = xx[best_idx]
+            yoff = yy[best_idx]
             if r == 1:
                 x = x_fine
                 y = y_fine
@@ -164,7 +181,7 @@ if __name__ == "__main__":
             yy.append(y_+yoff)
 
             # Find best k for each antenna
-            shifts, maximum = find_best_sample_shifts_summing_at_location(test_loc, ev.antennas, allowed_sample_shifts, dt)
+            shifts, maximum = find_best_sample_shifts_summing_at_location(test_loc, ev.antennas, allowed_sample_shifts, dt=dt)
 
             # Translate sample shifts back into period multiple k
             ks = np.rint(shifts*f_beacon*dt)
@@ -174,6 +191,11 @@ if __name__ == "__main__":
 
         xx = np.array(xx)
         yy = np.array(yy)
+        locs = list(zip(xx, yy))
+
+        ## Save maxima to file
+        np.savetxt(savepath + __file__+f'.maxima.run{r}.txt', np.column_stack((locs, maxima_per_loc, ks_per_loc)) )
+
         if True: #plot maximum at test locations
             fig, axs = plt.subplots()
             axs.set_title(f"Grid Run {r}")
@@ -182,26 +204,18 @@ if __name__ == "__main__":
             sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6)
             fig.colorbar(sc, ax=axs)
 
-            fig.savefig(__file__+f'.run{r}.pdf')
+            fig.savefig(savepath + __file__+f'.maxima.run{r}.pdf')
 
         ## Save ks to file
         best_idx = np.argmax(maxima_per_loc)
-        np.savetxt(__file__+f'.bestk.run{r}.txt', ks_per_loc[best_idx] )
+        np.savetxt(savepath + __file__+f'.bestk.run{r}.txt', ks_per_loc[best_idx] )
-        print(ks_per_loc[best_idx])
+        print('Best k:', ks_per_loc[best_idx])
-
-        ## Save maxima to file
-        np.savetxt(__file__+f'.maxima.run{r}.txt', maxima_per_loc)
 
         # Abort if no improvement
         if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ):
-            print("No improvement, breaking")
+            print("No changes from previous grid, breaking")
             break
 
-        # Prepare variables for next loop
-        old_ks_per_loc = ks_per_loc[best_idx]
-        xoff = xx[best_idx]
-        yoff = yy[best_idx]
-
     # Save best ks to hdf5 antenna file
     with h5py.File(antennas_fname, 'a') as fp:
         group = fp.require_group('antennas')