ZH: include trace windowing for power on grid script

simulations/airshower_beacon_simulation/aa_generate_beacon.py

@@ -29,7 +29,7 @@ def read_antenna_clock_repair_offsets(antennas, mode='all', freq_name=None):
 
         # original timing
         if mode == 'orig':
-            _clock_delta = ant.attrs['clock_offset']
+            _clock_delta = -1*ant.attrs['clock_offset']
 
         # phase
         if mode in ['all', 'phases']:

simulations/airshower_beacon_simulation/dc_grid_power_time_fixes.py

@@ -50,11 +50,13 @@ if __name__ == "__main__":
     show_plots = args.show_plots
 
     remove_beacon_from_traces = True
+    apply_signal_window_from_max = True
 
     ####
     fname_dir = path.dirname(fname)
     antennas_fname = path.join(fname_dir, beacon.antennas_fname)
     pickle_fname = path.join(fname_dir, 'res.pkl')
+    tx_fname = path.join(fname_dir, beacon.tx_fname)
 
     # create fig_dir
     if fig_dir:
@@ -62,13 +64,13 @@ if __name__ == "__main__":
 
     # Read in antennas from file
     _, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)
+    _, __, txdata = beacon.read_tx_file(tx_fname)
     # Read original REvent
     ev = REvent(fname)
+    bak_ants = ev.antennas
     # .. patch in our antennas
     ev.antennas = antennas
 
-    rit.set_pol_and_bp(ev)
-
     ##
     ## Setup grid
     ##
@@ -86,10 +88,6 @@ if __name__ == "__main__":
             'scale02d': scale02d,
             }
 
-    # backup antenna times
-    backup_antenna_t = [ ant.t for ant in ev.antennas ]
-    backup_antenna_t_AxB = [ ant.t_AxB for ant in ev.antennas ]
-
     plot_titling = {
             'no_offset': "no clock offset",
             'repair_none': "unrepaired clock offset",
@@ -108,16 +106,56 @@ if __name__ == "__main__":
     #
     # We need to remove it here, so we do not shoot ourselves in
     # the foot when changing to the various clock offsets.
+    #
+    # Note that the bandpass filter is applied only after E_AxB is
+    # reconstructed so we have to manipulate the original traces.
     if remove_beacon_from_traces:
+        tx_amps = txdata['amplitudes']
+        tx_amps_sum = np.sum(tx_amps)
+
         for i, ant in enumerate(ev.antennas):
             beacon_phase = ant.beacon_info[freq_name]['beacon_phase']
             f = ant.beacon_info[freq_name]['freq']
             ampl_AxB = ant.beacon_info[freq_name]['amplitude']
             calc_beacon = lib.sine_beacon(f, ev.antennas[i].t, amplitude=ampl_AxB, phase=beacon_phase)
 
-            # Only need to manipulate E_AxB
+            # Split up contribution to the various polarisations
+            for j, amp in enumerate(tx_amps):
+                if j == 0:
+                    ev.antennas[i].Ex -= amp*(1/tx_amps_sum)*calc_beacon
+                elif j == 1:
+                    ev.antennas[i].Ey -= amp*(1/tx_amps_sum)*calc_beacon
+                elif j == 2:
+                    ev.antennas[i].Ez -= amp*(1/tx_amps_sum)*calc_beacon
+
+            # 
             ev.antennas[i].E_AxB -= calc_beacon
 
+    # Slice the traces to a small part around the peak
+    if apply_signal_window_from_max:
+        N_pre, N_post = 250, 250
+
+        for i, ant in enumerate(ev.antennas):
+            max_idx = np.argmax(ant.E_AxB)
+
+            low_idx = max(0, max_idx-N_pre)
+            high_idx = min(len(ant.t), max_idx+N_post)
+
+            ev.antennas[i].t = ant.t[low_idx:high_idx]
+            ev.antennas[i].t_AxB = ant.t_AxB[low_idx:high_idx]
+
+            ev.antennas[i].Ex = ant.Ex[low_idx:high_idx]
+            ev.antennas[i].Ey = ant.Ey[low_idx:high_idx]
+            ev.antennas[i].Ez = ant.Ez[low_idx:high_idx]
+            ev.antennas[i].E_AxB = ant.E_AxB[low_idx:high_idx]
+
+    # backup antenna times
+    backup_antenna_t = [ ant.t for ant in ev.antennas ]
+    backup_antenna_t_AxB = [ ant.t_AxB for ant in ev.antennas ]
+
+    ## Apply polarisation and bandpass filter
+    rit.set_pol_and_bp(ev)
+
     with joblib.parallel_backend("loky"):
         for case in wanted_cases:
             print(f"Starting {case} figure")
@@ -133,9 +171,41 @@ if __name__ == "__main__":
             for i, ant in enumerate(ev.antennas):
                 total_clock_offset = measured_offsets[i]
 
-                #ev.antennas[i].t =     backup_antenna_t[i] + total_clock_offset
+                ev.antennas[i].t =     backup_antenna_t[i] + total_clock_offset
                 ev.antennas[i].t_AxB = backup_antenna_t_AxB[i] + total_clock_offset
 
+                if i == 0:
+                    # Specifically compare the times
+                    print(bak_ants[i].t[0], ev.antennas[i].t[0], ev.antennas[i].t[0], ev.antennas[i].attrs['clock_offset'], measured_offsets[i])
+
+            #
+            # Plot overlapping traces at 0,0,0
+            #
+            if True:
+                P, t_, a_, a_sum, t_sum = rit.pow_and_time([0,0,0], ev, dt=1)
+
+                fig, axs = plt.subplots()
+                axs.set_title("Antenna traces" + "\n" + plot_titling[case])
+                axs.set_xlabel("Time [ns]")
+                axs.set_ylabel("Amplitude [$\\mu V/m$]")
+
+                a_max = [ np.amax(ant.E_AxB) for ant in ev.antennas ]
+                power_sort_idx = np.argsort(a_max)
+
+                N_plot = 30
+                for i, idx in enumerate(power_sort_idx):
+                    if i > N_plot:
+                        break
+                    
+                    alpha = max(0.7, 1/len(a_))
+
+                    axs.plot(t_[idx], a_[idx], color='r', alpha=alpha)
+
+                if fig_dir:
+                    fig.tight_layout()
+                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.trace_overlap.{case}.pdf'))
+                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.trace_overlap.{case}.png'), transparent=True)
+
             # Measure power on grid
             for scalename, scale in scales.items():
                 wx, wy = scale, scale