ZH: find beacon multiple by reconstructing shower amplitudes

simulations/airshower_beacon_simulation/bc_period_shower_plane.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+# vim: fdm=indent ts=4
+
+"""
+Find the best integer multiple to shift
+antennas to be able to resolve 
+"""
+
+import h5py
+from itertools import combinations, zip_longest, product
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.interpolate import interp1d
+
+from earsim import REvent
+from atmocal import AtmoCal
+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=1):
+    """
+    Find the best sample_shift for each antenna by summing the antenna traces
+    and seeing how to get the best alignment.
+    """
+    a_ = []
+    t_ = []
+    t_min = 1e9
+    t_max = -1e9
+
+    # propagate to test location
+    for i, ant in enumerate(antennas):
+        aloc = [ant.x, ant.y, ant.z]
+        delta, dist = atm.light_travel_time(test_loc, aloc)
+        delta = delta*1e9
+
+        t__ = np.subtract(ant.t_AxB, delta)
+        t_.append(t__)
+        a_.append(ant.E_AxB)
+
+        if t__[0] < t_min:
+            t_min = t__[0]
+        if t__[-1] > t_max:
+            t_max = t__[-1]
+
+    # Interpolate and find best sample shift
+    max_neg_shift = np.min(allowed_sample_shifts) * dt
+    max_pos_shift = np.max(allowed_sample_shifts) * dt
+
+    t_sum = np.arange(t_min+1+max_neg_shift, t_max-1+max_pos_shift, dt)
+    a_sum = np.zeros(len(t_sum))
+
+    best_sample_shifts = np.zeros( (len(antennas)) ,dtype=int)
+    for i, (t_r, E_) in enumerate(zip(t_, a_)):
+        t_min = t_r[0]
+        f = interp1d(t_r, E_, assume_sorted=True, bounds_error=False, fill_value=0)
+        a_int = f(t_sum)
+
+        if i == 0:
+            a_sum += a_int
+            best_sample_shifts[i] = 0
+            continue
+
+        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)
+
+        # transform maximum into best_sample_shift
+        best_idx = np.argmax(shift_maxima)
+        best_sample_shifts[i] = allowed_sample_shifts[best_idx]
+        a_sum += np.roll(a_int, best_sample_shifts[i])
+
+    # Return ks
+    return best_sample_shifts, np.max(a_sum)
+
+if __name__ == "__main__":
+    from os import path
+    import sys
+
+    atm = AtmoCal()
+
+    fname = "ZH_airshower/mysim.sry"
+
+    allowed_ks = np.arange(-2, 3, 1)
+    Xref = 450
+
+    x_coarse = np.linspace(-2e3, 2e3, 11)
+    y_coarse = np.linspace(-2e3, 2e3, 11)
+
+    ####
+    fname_dir = path.dirname(fname)
+    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
+    time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_fname
+
+    # Read in antennas from file
+    _, __, antennas = beacon.read_beacon_hdf5(antennas_fname)
+    # Read original REvent
+    ev = REvent(fname)
+    # .. patch in our antennas
+    ev.antennas = antennas
+
+    # For now only implement using one freq_name
+    freq_names = antennas[0].beacon_info.keys()
+    if len(freq_names) > 1:
+        raise NotImplementedError
+
+    freq_name = next(iter(freq_names))
+    f_beacon = ev.antennas[0].beacon_info[freq_name]['freq']
+
+    # determine best ks per location
+    dt = ev.antennas[0].t[1] - ev.antennas[0].t[0]
+    allowed_sample_shifts = np.ceil(allowed_ks/f_beacon /dt).astype(int)
+
+
+    # Remove time due to true phase
+    for i, ant in enumerate(ev.antennas):
+        true_phase = ant.beacon_info[freq_name]['phase']
+        true_phase_time = true_phase/(2*np.pi*f_beacon)
+
+        ev.antennas[i].t -= true_phase_time
+
+    # Prepare polarisation and passbands
+    rit.set_pol_and_bp(ev, low=0.03, high=0.08)
+
+    dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),Xref,0)
+    scale2d = dXref*np.tan(np.deg2rad(2.))
+
+    # Setup Plane grid to test
+    for r in range(6):
+        xoff, yoff = 0,0
+        if r == 0:
+            x = x_coarse
+            y = y_coarse
+        else:
+            best_idx = np.argmax(maxima_per_loc)
+            xoff = xx[best_idx]
+            yoff = yy[best_idx]
+            x /= 4
+            y /= 4
+
+        print(f"Testing grid {r} centered on {xoff}, {yoff}")
+
+        ks_per_loc = np.zeros( (len(x)*len(y), len(ev.antennas)) )
+        maxima_per_loc = np.zeros( (len(x)*len(y)))
+    
+        ## Check each location on grid
+        xx = []
+        yy = []
+        N_loc = len(maxima_per_loc)
+        for i, (x_, y_) in enumerate(product(x,y)):
+            if i % 10 ==0:
+                print(f"Testing location {i} out of {N_loc}")
+            test_loc = (x_+xoff)* ev.uAxB + (y_+yoff)*ev.uAxAxB + dXref *ev.uA
+            xx.append(x_+xoff)
+            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)
+    
+            # Translate sample shifts back into period multiple k
+            ks = shifts*f_beacon*dt
+    
+            ks_per_loc[i] = ks
+            maxima_per_loc[i] = maximum
+   
+        xx = np.array(xx)
+        yy = np.array(yy)
+
+        best_idx = np.argmax(maxima_per_loc)
+        np.savetxt(__file__+f'.run{r}.txt')
+        print(ks_per_loc[best_idx])
+    
+        if True: #plot maximum at test locations
+            fig, axs = plt.subplots()
+            axs.set_title(f"Grid Run {r}")
+            axs.set_ylabel("vxvxB [km]")
+            axs.set_xlabel("  vxB [km]")
+            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')
+    plt.show()