From 7008e9e8ce1b5378eb09c45ca4dd3880867a0798 Mon Sep 17 00:00:00 2001
From: Eric Teunis de Boone <ericteunis@deboone.nl>
Date: Fri, 20 Jan 2023 14:33:24 +0100
Subject: [PATCH] ZH: remove beacon from traces before da_reconstruction

---
 .../da_reconstruction.py                      | 62 +++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/simulations/airshower_beacon_simulation/da_reconstruction.py b/simulations/airshower_beacon_simulation/da_reconstruction.py
index e41ff14..f1d40b1 100755
--- a/simulations/airshower_beacon_simulation/da_reconstruction.py
+++ b/simulations/airshower_beacon_simulation/da_reconstruction.py
@@ -74,8 +74,10 @@ if __name__ == "__main__":
     for i, ant in enumerate(ev.antennas):
         # t_AxB will be set by the rit.set_pol_and_bp function
         ev.antennas[i].t += measured_repair_offsets[i]
+        ev.antennas[i].t_AxB += measured_repair_offsets[i]
 
         # .. and remove the beacon from the traces
+        # Note: ant.E_AxB is recalculated by rit.set_pol_and_bp
         if remove_beacon_from_traces:
             clock_phase = measured_repair_offsets[i]*2*np.pi*f_beacon
 
@@ -96,6 +98,66 @@ if __name__ == "__main__":
                 elif j == 2:
                     ev.antennas[i].Ez -= amp*(1/tx_amps_sum)*calc_beacon
 
+            # Subtract the beacon from E_AxB
+            ev.antennas[i].E_AxB -= calc_beacon
+
+            ##
+            ## Make a figure of the manipulated traces
+            ##
+            if i == 2:
+                orig_beacon_amplifier = ampl_AxB/max(ant.beacon)
+
+                for k in range(2):
+                    if k == 0:
+                        time = ant.t_AxB
+                        trace = ant.E_AxB
+                        tmp_beacon = calc_beacon
+                        fname_extra = ""
+                    else:
+                        time = ant.t
+                        trace = ant.Ex
+                        tmp_beacon = tx_amps[0]/tx_amps_sum * calc_beacon
+                        fname_extra = ".Ex"
+
+                    fig, ax = plt.subplots()
+                    ax.set_title(f"Signal and Beacon traces Antenna {i}")
+                    ax.set_xlabel("Time [ns]")
+                    ax.set_ylabel("Amplitude [$\\mu V/m$]")
+
+                    ax.plot(time, trace + tmp_beacon, alpha=0.6, ls='dashed', label='Signal') # calc_beacon was already removed
+                    ax.plot(time, tmp_beacon,        alpha=0.6, ls='dashed', label='Calc Beacon')
+                    ax.plot(time, trace,             alpha=0.6,              label="Signal - Calc Beacon")
+
+                    if k == 0:
+                        ax.legend()
+                    else:
+                        ax.legend(title="Ex")
+
+                    # save
+                    if fig_dir:
+                        fig.tight_layout()
+
+                        if True: # zoom
+                            old_xlim = ax.get_xlim()
+
+                            if True: # zoomed on part without peak of this trace
+                                wx, x = 100, 100
+                                ax.set_xlim(x-wx, x+wx)
+
+                                fig.savefig(path.join(fig_dir, __file__+f'.traces.A{i}.zoomed.beacon{fname_extra}.pdf'))
+
+                            if True: # zoomed on peak of this trace
+                                idx = np.argmax(ev.antennas[i].E_AxB)
+                                x = ev.antennas[i].t_AxB[idx]
+                                wx = 100
+                                ax.set_xlim(x-wx, x+wx)
+                                fig.savefig(path.join(fig_dir, __file__+f".traces.A{i}.zoomed.peak{fname_extra}.pdf"))
+
+                            ax.set_xlim(*old_xlim)
+
+                        fig.savefig(path.join(fig_dir, __file__+f'.traces.A{i}.pdf'))
+
+
     N_X, Xlow, Xhigh = 23, 100, 1200
     with joblib.parallel_backend("loky"):
         res = rit.reconstruction(ev, outfile=fig_subdir+'/fig.pdf', slice_outdir=fig_subdir+'/', Xlow=Xlow, N_X=N_X, Xhigh=Xhigh)