ZH: do reconstruction for best_k for each grid run iteration

simulations/airshower_beacon_simulation/ca_period_from_shower.py

@@ -154,15 +154,9 @@ if __name__ == "__main__":
     fig_subdir = path.join(fig_dir, 'shifts/')
     show_plots = False
 
-    allowed_ks = np.arange(-2, 3, 1)
+    allowed_ks = [ -2, -1, 0, 1, 2]
     Xref = 400
 
-    x_coarse = np.linspace(-20e3, 20e3, 10)
-    y_coarse = np.linspace(-20e3, 20e3, 10)
-
-    x_fine = np.linspace(-2e3, 2e3, 30)
-    y_fine = np.linspace(-2e3, 2e3, 30)
-
     N_runs = 3
 
     ####
@@ -250,16 +244,42 @@ if __name__ == "__main__":
         plt.show()
 
 
+    ##
+    ## Determine grid positions
+    ##
+
     dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0)
     scale2d = dXref*np.tan(np.deg2rad(2.))
+    scale4d = dXref*np.tan(np.deg2rad(4.))
+
+    if not True: #quicky
+        N_runs = 2
+        x_coarse = np.linspace(-scale2d, scale2d, 4)
+        y_coarse = np.linspace(-scale2d, scale2d, 4)
+
+        x_fine = x_coarse/4
+        y_fine = y_coarse/4
+    else: # long
+        N_runs = 5
+        x_coarse = np.linspace(-scale4d, scale4d, 40)
+        y_coarse = np.linspace(-scale4d, scale4d, 40)
+
+        x_fine = np.linspace(-scale2d, scale2d, 40)
+        y_fine = np.linspace(-scale2d, scale2d, 40)
+
+
+    ##
+    ## Do calculations on the grid
+    ##
 
-    # Setup Plane grid to test
     for r in range(N_runs):
+        # Setup Plane grid to test
         xoff, yoff = 0,0
         if r == 0:
             x = x_coarse
             y = y_coarse
         else:
+            # zooming in
             old_ks_per_loc = ks_per_loc[best_idx]
             xoff = xx[best_idx]
             yoff = yy[best_idx]
@@ -308,7 +328,7 @@ if __name__ == "__main__":
 
         if True: #plot maximum at test locations
             fig, axs = plt.subplots()
-            axs.set_title(f"Grid Run {r}")
+            axs.set_title(f"Optimizing signal strength, Grid Run {r}")
             axs.set_ylabel("vxvxB [km]")
             axs.set_xlabel("  vxB [km]")
             axs.set_aspect('equal', 'datalim')
@@ -328,15 +348,37 @@ if __name__ == "__main__":
                 axs.set_ylim(*old_ylims)
                 fig.tight_layout()
 
-        ## Save ks to file
+        ##
         best_idx = np.argmax(maxima_per_loc)
-        np.savetxt(fig_dir + __file__+f'.bestk.run{r}.txt', ks_per_loc[best_idx] )
+        best_k = ks_per_loc[best_idx]
+
         print("Max at location: ", locs[best_idx])
-        print('Best k:', ks_per_loc[best_idx])
+        print('Best k:', best_k)
+
+        ## Save best ks to file
+        np.savetxt(fig_dir + __file__+f'.bestk.run{r}.txt', best_k )
+
+        ## Do a small reconstruction of the shower for best ks
+        if True:
+            print("Reconstructing for best k")
+            _, __, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=scale2d, wy=scale2d, xoff=xoff, yoff=yoff, zgr=0)
+
+            fig, axs = plt.subplots()
+            axs.set_title(f"Shower reconstruction with best k, Grid Run {r}")
+            axs.set_ylabel("vxvxB [km]")
+            axs.set_xlabel("  vxB [km]")
+            axs.set_aspect('equal', 'datalim')
+            sc = axs.scatter(xx/1e3, yy/1e3, c=p, cmap='Spectral_r', alpha=0.6)
+            fig.colorbar(sc, ax=axs)
+
+            if fig_dir:
+                fig.tight_layout()
+                fig.savefig(path.join(fig_dir, __file__+f'.reconstruction.run{r}.pdf'))
 
         # Abort if no improvement
         if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ):
             print("No changes from previous grid, breaking")
+            # TODO: notate this case somewhere
             break
 
     old_ks_per_loc = ks_per_loc[best_idx]
@@ -359,4 +401,5 @@ if __name__ == "__main__":
             h5attrs['best_k'] = old_ks_per_loc[i]
             h5attrs['best_k_time'] = old_ks_per_loc[i]*dt/f_beacon
 
-    plt.show()
+    if show_plots:
+        plt.show()