ZH: ca_period_from_shower as used for Thesis Figures

airshower_beacon_simulation/ca_period_from_shower.py

@@ -25,7 +25,25 @@ try:
 except:
     tqdm = lambda x: x
 
-def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, period=1, dt=None, period_shift_first_trace=0, plot_iteration_with_shifted_trace=None, fig_dir=None, fig_distinguish=None,snr_str=None, shower_plane_loc=None):
+try:
+    from joblib import Parallel, delayed
+except:
+    Parallel = None
+    delayed = lambda x: x
+
+def find_best_period_shifts_at_location(*args, algo=None, **kwargs):
+    """
+    This is a placeholder function.
+
+    For args and kwargs see find_best_period_shifts_summing_at_location.
+    """
+
+    if algo is None:
+        algo = 'sum'
+
+    return find_best_period_shifts_summing_at_location(*args, **kwargs, algo=algo)
+
+def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks, period=1, dt=None, period_shift_first_trace=0, plot_iteration_with_shifted_trace=None, fig_dir=None, fig_distinguish=None,snr_str=None, shower_plane_loc=None, algo='sum', verbose=False):
     """
     Find the best sample_shift for each antenna by summing the antenna traces
     and seeing how to get the best alignment.
@@ -80,9 +98,9 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
         f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0)
 
         if i == 0:
-            a_sum += f(t_sum)
-            a_first = a_sum
             best_period_shifts[i] = period_shift_first_trace
+            a_first = f(t_sum - period_shift_first_trace*period)
+            a_sum += a_first
             continue
 
         # init figure
@@ -92,47 +110,67 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
                 title_location = "s({:.1g},{:.1g},{:.1g})".format(*shower_plane_loc)
             else:
                 title_location = "({.1g},{:.1g},{:.1g})".format(*test_loc)
-            ax.set_title("Traces at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant))
+            #ax.set_title("Traces at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant))
             ax.set_xlabel("Time [ns]")
             ax.set_ylabel("Amplitude")
-            ax.plot(t_sum, a_sum)
+            #ax.plot(t_sum, a_sum)
 
             fig2, ax2 = plt.subplots(figsize=figsize)
-            ax2.set_title("Maxima at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant))
+            #ax2.set_title("Maxima at {}; i={i}/{tot}".format(title_location, i=i, tot=N_ant))
             ax2.set_xlabel("$k$th Period")
             ax2.set_ylabel("Summed Amplitude")
-            ax2.plot(0, np.max(a_first), marker='*', label='strongest trace', ls='none', ms=20)
+            ax2.plot(0, np.max(a_first), marker='*', label='first trace', ls='none', ms=20)
+            ax3 = ax2.twinx()
+            ax3.set_ylabel("Correlation")
 
         # find the maxima for each period shift k
         shift_maxima = np.zeros( len(allowed_ks) )
+        shift_corrs = np.zeros_like(shift_maxima)
         for j, k in enumerate(allowed_ks):
-            augmented_a = f(t_sum + k*period)
+            augmented_a = f(t_sum - k*period)
 
             shift_maxima[j] = np.max(augmented_a + a_first)
+            shift_corrs[j] = np.dot(augmented_a, a_first)
 
             if i in plot_iteration_with_shifted_trace and abs(k) <= 3:
-                ax.plot(t_sum, augmented_a, alpha=0.7, ls='dashed', label=f'{k:g}')
+                l = ax.plot(t_sum, a_first + augmented_a, alpha=0.7, ls='dashed', label=f'{k:g}')
+                ax.axhline(shift_maxima[j], ls='dashdot', color=l[0].get_color(), alpha=0.7)
+
                 ax2.plot(k, shift_maxima[j], marker='o', ls='none', ms=20)
+                ax3.plot(k, shift_corrs[j], marker='3', ls='none', ms=20)
 
         # transform maximum into best_sample_shift
         best_idx = np.argmax(shift_maxima)
+        best_corr_idx = np.argmax(shift_corrs)
 
+        if verbose and (best_corr_idx != best_idx):
+            print("Correlation idx not equal to maximum idx")
+            print(best_corr_idx, best_idx)
+
+        if algo == 'sum':
             best_period_shifts[i] = allowed_ks[best_idx]
-        best_augmented_a = f(t_sum + k*period)
+        elif algo == 'corr':
+            best_period_shifts[i] = allowed_ks[best_corr_idx]
+
+        k = best_period_shifts[i]
+        best_augmented_a = f(t_sum - k*period)
         a_sum += best_augmented_a
 
         # cleanup figure
         if i in plot_iteration_with_shifted_trace:
             # plot the traces
             if True: # plot best k again
-                ax.plot(t_sum, best_augmented_a, alpha=0.8, label=f'best $k$={best_period_shifts[i]:g}', lw=2)
+                l = ax.plot(t_sum, a_first + best_augmented_a, alpha=0.8, label=f'best $k$={best_period_shifts[i]:g}', lw=2)
+                ax.axhline(shift_maxima[j], ls='dashdot', color=l[0].get_color(), alpha=0.7)
 
             if True: # plot best shift
                 ax2.plot(allowed_ks[best_idx], shift_maxima[best_idx], marker='*', ls='none', ms=20, label=f'best $k$={best_period_shifts[i]:g}')
+                ax3.plot(allowed_ks[best_corr_idx], shift_maxima[best_idx], marker='*', ls='none', ms=20, label=f'best corr $k$={allowed_ks[best_corr_idx]:g}')
 
 
-            ax.legend(title='period shift $k$; '+snr_str, ncol=5 )
+            ax.legend(title='period shift $k$; '+snr_str, ncol=5, loc='lower center')
             ax2.legend(title=snr_str)
+            ax3.legend()
             if fig_dir:
                 fig.tight_layout()
                 fig2.tight_layout()
@@ -147,14 +185,15 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
                     old_xlim = ax.get_xlim()
 
                     if True: # zoomed on part without peak of this trace
-                        wx = 100
+                        wx = 200
                         x = max(t_r) - wx
                         ax.set_xlim(x-wx, x+wx)
                         fig.savefig(fname + ".zoomed.beacon.pdf")
 
                     if True: # zoomed on peak of this trace
-                        x = t_r[np.argmax(E_)]
+                        x = t_sum[np.argmax(a_first)]
-                        wx = 50 + (max(best_period_shifts) - min(best_period_shifts))*dt
+                        x = t_sum[np.argmax(f(t_sum))]
+                        wx = 50 + (max(best_period_shifts) - min(best_period_shifts) )*dt + 1*period
                         ax.set_xlim(x-wx, x+wx)
                         fig.savefig(fname + ".zoomed.peak.pdf")
 
@@ -165,13 +204,30 @@ def find_best_period_shifts_summing_at_location(test_loc, antennas, allowed_ks,
             plt.close(fig)
             plt.close(fig2)
 
+    if True: # final summed waveform
+        fig, ax = plt.subplots()
+
+        ax.set_title("Summed Traces with best k's at {}".format(title_location))
+        ax.set_xlabel("Time [ns]")
+        ax.set_ylabel("Amplitude")
+
+        ax.plot(t_sum, a_sum)
+
+        if fig_dir:
+            fig.tight_layout()
+
+            fname = path.join(fig_dir, path.basename(__file__) + f'.{fig_distinguish}i{i}' + fname_location)
+            fig.savefig(fname + ".sum.pdf")
+            
+            plt.close(fig)
+
     # sort by antenna (undo sorting by maximum)
     undo_sort_idx = np.argsort(sort_idx)
 
     best_period_shifts = best_period_shifts[undo_sort_idx]
 
     # Return ks
-    return best_period_shifts, np.max(a_sum)
+    return best_period_shifts, np.max(a_sum), sort_idx
 
 if __name__ == "__main__":
     import sys
@@ -204,7 +260,16 @@ if __name__ == "__main__":
     if path.isdir(args.input_fname):
         args.input_fname = path.join(args.input_fname, "mysim.sry")
 
-    figsize = (12,8)
+    figsize = (6,4)
+    if True:
+        from matplotlib import rcParams
+        #rcParams["text.usetex"] = True
+        rcParams["font.family"] = "serif"
+        rcParams["font.size"] = "14"
+        rcParams["grid.linestyle"] = 'dotted'
+        rcParams["figure.figsize"] = figsize
+        figsize = rcParams['figure.figsize']
+
 
     fig_dir = args.fig_dir
     fig_subdir = path.join(fig_dir, 'shifts/')
@@ -248,7 +313,7 @@ if __name__ == "__main__":
     ev.antennas = antennas
     # read in snr information
     beacon_snrs = beacon.read_snr_file(beacon_snr_fname)
-    snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .1g}"
+    snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .2g}"
 
     # For now only implement using one freq_name
     freq_names = antennas[0].beacon_info.keys()
@@ -349,8 +414,8 @@ if __name__ == "__main__":
                     if True: # zoom
                         old_xlim = ax.get_xlim()
 
-                        if True: # zoomed on part without peak of this trace
+                        if not True: # zoomed on part without peak of this trace
-                            wx, x = 200, min(ant.t_AxB)
+                            wx, x = 100, min(ant.t_AxB)
                             ax.set_xlim(x-5, x+wx)
 
                             fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.traces.A{ant.name}.zoomed.beacon.pdf'))
@@ -358,7 +423,7 @@ if __name__ == "__main__":
                         if True: # zoomed on peak of this trace
                             idx = np.argmax(ev.antennas[i].E_AxB)
                             x = ev.antennas[i].t_AxB[idx]
-                            wx = 300
+                            wx = 150
                             ax.set_xlim(x-wx//2, x+wx//2)
                             fig.savefig(path.join(fig_dir, path.basename(__file__)+f".traces.A{ant.name}.zoomed.peak.pdf"))
 
@@ -380,22 +445,23 @@ if __name__ == "__main__":
     ## Determine grid positions
     ##
 
-    dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0)
+    zgr = 0 + ev.core[2]
+    dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,zgr)
     scale2d = dXref*np.tan(np.deg2rad(2.))
     scale4d = dXref*np.tan(np.deg2rad(4.))
 
     if args.quick_run: #quicky
-        x_coarse = np.linspace(-scale2d, scale2d, 6)
+        x_coarse = np.linspace(-scale4d, scale4d, 16)
-        y_coarse = np.linspace(-scale2d, scale2d, 6)
+        y_coarse = np.linspace(-scale4d, scale4d, 16)
 
         x_fine = x_coarse/4
         y_fine = y_coarse/4
     else: # long
-        x_coarse = np.linspace(-scale4d, scale4d, 40)
+        x_coarse = np.linspace(-scale4d, scale4d, 14)
-        y_coarse = np.linspace(-scale4d, scale4d, 40)
+        y_coarse = np.linspace(-scale4d, scale4d, 14)
 
-        x_fine = np.linspace(-scale2d, scale2d, 40)
+        x_fine = np.linspace(-scale2d, scale2d, 18)
-        y_fine = np.linspace(-scale2d, scale2d, 40)
+        y_fine = np.linspace(-scale2d, scale2d, 18)
 
     ## Remove run_break_fname if it exists
     try:
@@ -435,9 +501,29 @@ if __name__ == "__main__":
         yy = []
         N_loc = len(maxima_per_loc)
 
-        for i, (x_, y_) in tqdm(enumerate(product(x,y)), total=N_loc, ):
+        # Make the grid a list of locations
+        xx = []
+        yy = []
+        for i, (x_, y_) in enumerate(product(x,y)):
+            xx.append( x_+xoff )
+            yy.append( y_+yoff )
+
+        xx = np.array(xx)
+        yy = np.array(yy)
+        locs = list(zip(xx, yy))
+
+        # Shift the reference trace
+        if r == 0:
+            first_trace_period_shift = 0
+        else:
+            first_trace_period_shift = 0#first_trace_period_shift + np.rint(np.mean(old_ks_per_loc))
+
+            print("New first trace period:", first_trace_period_shift)
+
+        # define loop func for joblib
+        def loop_func(loc, dXref=dXref, i=1):
             tmp_fig_subdir = None
-            if i % 10 ==0:
+            if i  == 0:
                 if hasattr(tqdm, '__code__') and tqdm.__code__.co-name == '<lambda>':
                     print(f"Testing location {i} out of {N_loc}")
                 tmp_fig_subdir = fig_subdir
@@ -445,33 +531,41 @@ if __name__ == "__main__":
             test_loc = loc[0]* ev.uAxB + loc[1]*ev.uAxAxB + dXref *ev.uA
 
             # Find best k for each antenna
-            ks_per_loc[i], maxima_per_loc[i] = find_best_period_shifts_summing_at_location(test_loc, ev.antennas, allowed_ks, period=1/f_beacon, dt=dt,
+            return find_best_period_shifts_at_location(test_loc, ev.antennas, allowed_ks, period=1/f_beacon, dt=dt, period_shift_first_trace=first_trace_period_shift,
-                    plot_iteration_with_shifted_trace=[ 5, len(ev.antennas)-1],
+                    plot_iteration_with_shifted_trace=[ 1, 2, 3, 4, 5, len(ev.antennas)-1],
-                    fig_dir=tmp_fig_subdir, fig_distinguish=f"X{Xref}. run{r}.",
+                    fig_dir=tmp_fig_subdir, fig_distinguish=f"X{Xref}.run{r}.",
                     snr_str=snr_str,shower_plane_loc=(loc[0]/1e3, loc[1]/1e3, dXref),
                     )
 
-        xx = np.array(xx)
+        res = ( delayed(loop_func)(loc, i=i) for i, loc in enumerate(locs) )
-        yy = np.array(yy)
+
-        locs = list(zip(xx, yy))
+        if Parallel:
+            res = Parallel(n_jobs=None)(tqdm(res, total=len(locs)))
+        else:
+            res = tqdm(res, total=len(locs))
+
+        # unpack loop results
+        ks_per_loc, maxima_per_loc, sort_idx = zip(*res)
 
         ## Save maxima to file
         np.savetxt(path.join(fig_dir, path.basename(__file__)+f'.maxima.X{Xref}.run{r}.txt'), np.column_stack((locs, maxima_per_loc, ks_per_loc)) )
 
+        scatter_kwargs = dict(cmap='Spectral_r', s=64*4, alpha=0.6)
         if True: #plot maximum at test locations
             fig, axs = plt.subplots(figsize=figsize)
-            axs.set_title(f"Optimizing signal strength by varying $k$ per antenna,\n Grid Run {r}")
+            #axs.set_title(f"Optimizing signal strength by varying $k$ per antenna,\n Grid Run {r}")
-            axs.set_ylabel("vxvxB [km]")
+            axs.set_ylabel(" vxvxB [km]")
-            axs.set_xlabel("  vxB [km]")
+            axs.set_xlabel("-v x B [km]")
             axs.set_aspect('equal', 'datalim')
-            sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, cmap='Spectral_r', alpha=0.6)
+            sc = axs.scatter(xx/1e3, yy/1e3, c=maxima_per_loc, **scatter_kwargs)
             fig.colorbar(sc, ax=axs, label='Max Amplitude [$\\mu V/m$]')
 
-            axs.legend(title=snr_str)
+            #axs.legend(title=snr_str)
 
             # indicate maximum value
             idx = np.argmax(maxima_per_loc)
-            axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30)
+            axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) # max value
+            axs.plot(0,0, 'r+', ms=30) # true axis
 
             if fig_dir:
                 old_xlims = axs.get_xlim()
@@ -490,8 +584,10 @@ if __name__ == "__main__":
         best_idx = np.argmax(maxima_per_loc)
         best_k = ks_per_loc[best_idx]
 
-        print("Max at location: ", locs[best_idx])
+        print("Max at location: ", locs[best_idx], ": ", maxima_per_loc[best_idx])
-        print('Best k:', best_k)
+        print('Best k:', best_k[sort_idx[best_idx]]) # Sorted by DESC amplitude
+        print('Mean best k:', np.rint(np.mean(best_k)))
+        print('first k:', first_trace_period_shift)
 
         ## Save best ks to file
         np.savetxt(path.join(fig_dir, path.basename(__file__)+f'.bestk.X{Xref}.run{r}.txt'), best_k )
@@ -509,33 +605,43 @@ if __name__ == "__main__":
 
                     # incorporate ks into timing
                     for i, ant in enumerate(ev.antennas):
-                        ev.antennas[i].t_AxB = ant.t_AxB - best_k[i] * 1/f_beacon
+                        ev.antennas[i].t_AxB = ant.t_AxB + best_k[i] * 1/f_beacon
 
-                    xx, yy, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=x[-1]-x[0], wy=y[-1]-y[0], xoff=xoff, yoff=yoff, zgr=0)
+                    xx, yy, p, ___ = rit.shower_plane_slice(ev, X=Xref, Nx=len(x), Ny=len(y), wx=(x[-1]-x[0])/2, wy=(y[-1]-y[0])/2, xoff=xoff, yoff=yoff, zgr=0)
 
                     # repair antenna times
                     for i, backup_t_AxB in enumerate(backup_times):
                         ev.antennas[i].t_AxB = backup_t_AxB
                 else: # get maximum amplitude at each location
                     maxima = np.empty( len(locs) )
-                    for i, loc in enumerate(locs):
+                    for i, loc in tqdm(enumerate(locs), total=len(locs)):
                         test_loc = loc[0]* ev.uAxB + loc[1]*ev.uAxAxB + dXref *ev.uA
                         P, t_, a_, a_sum, t_sum = rit.pow_and_time(test_loc, ev, dt=dt)
                         maxima[i] = np.max(a_sum)
 
                 fig, axs = plt.subplots(figsize=figsize)
-                axs.set_title(f"Shower slice for best k, Grid Run {r}")
+                #axs.set_title(f"Shower slice for best k, Grid Run {r}")
-                axs.set_ylabel("vxvxB [km]")
+                axs.set_ylabel(" vxvxB [km]")
-                axs.set_xlabel("  vxB [km]")
+                axs.set_xlabel("-v x B [km]")
-                axs.set_aspect('equal', 'datalim')
-                if power_reconstruction:
-                    sc = axs.scatter(xx/1e3, yy/1e3, c=p, cmap='Spectral_r', alpha=0.6)
-                    fig.colorbar(sc, ax=axs, label='Power')
-                else:
-                    sc = axs.scatter(xx/1e3, yy/1e3, c=maxima, cmap='Spectral_r', alpha=0.6)
-                    fig.colorbar(sc, ax=axs, label='Max Amplitude [$\\mu V/m$]')
 
-                axs.legend(title=snr_str)
+                if power_reconstruction:
+                    sc_c = p
+                    sc_label = 'Power [$(\\mu V/m)^2$]'
+                else:
+                    sc_c = maxima
+                    sc_label='Max Amplitude [$\\mu V/m$]'
+                sc = axs.scatter(xx/1e3, yy/1e3, c=sc_c, **scatter_kwargs)
+                fig.colorbar(sc, ax=axs, label=sc_label)
+            
+                # indicate maximum value
+                idx = np.argmax(p if power_reconstruction else maxima)
+                axs.plot(xx[idx]/1e3, yy[idx]/1e3, 'bx', ms=30) # max value
+                axs.plot(0,0, 'r+', ms=30) # true axis
+
+                # make square figure
+                axs.set_aspect('equal', 'datalim')
+
+                #axs.legend(title=snr_str)
 
                 if fig_dir:
                     if power_reconstruction:
@@ -547,7 +653,7 @@ if __name__ == "__main__":
                     fig.savefig(path.join(fig_dir, path.basename(__file__)+f'.reconstruction.X{Xref}.run{r}.{fname_extra}.pdf'))
 
         # Abort if no improvement
-        if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx]).all() ):
+        if ( r!= 0 and (old_ks_per_loc == ks_per_loc[best_idx] - first_trace_period_shift).all() ):
             print(f"No changes from previous grid, breaking at iteration {r} out of {N_runs}")
 
             try: