Merge branch 'add-noise' into main

simulations/airshower_beacon_simulation/aa_generate_beacon.py

@@ -82,6 +82,8 @@ def Antenna_from_h5ant(h5ant, traces_key='traces', raise_exception=True, read_Ax
         ant.Ez = deepcopy(h5ant[traces_key][3])
         if len(h5ant[traces_key]) > 4:
             ant.beacon = deepcopy(h5ant[traces_key][4])
+        if len(h5ant[traces_key]) > 5:
+            ant.noise = deepcopy(h5ant[traces_key][5])
 
     # E_AxB
     if read_AxB and 'E_AxB' in h5ant:
@@ -240,6 +242,8 @@ def write_baseline_time_diffs_hdf5(fname, baselines, true_phase_diffs, k_periods
 if __name__ == "__main__":
     from os import path
 
+    rng = np.random.default_rng()
+
     fname = "ZH_airshower/mysim.sry"
 
     # Transmitter
@@ -276,6 +280,9 @@ if __name__ == "__main__":
 
         beacon_amplitudes *= ampl
 
+    # Noise properties
+    noise_sigma = 1e-4 # set to None to ignore
+
     # Disable block_filter
     if False:
         block_filter = lambda x, dt, low, high: x
@@ -294,6 +301,7 @@ if __name__ == "__main__":
 
     print("Beacon amplitude at tx [muV/m]:", beacon_amplitudes)
     print("Tx location:", [tx.x, tx.y, tx.z])
+    print("Noise sigma:", noise_sigma)
 
     # read in antennas
     ev = REvent(fname)
@@ -328,13 +336,20 @@ if __name__ == "__main__":
                 print(f"Modified trace lengths by {pre_N},{after_N-N_samples}")
 
         beacon = lib.beacon_from(tx, antenna, f_beacon, antenna.t, c_light=c_light, radiate_rsq=beacon_radiate_rsq)
-        traces = np.array([antenna.Ex, antenna.Ey, antenna.Ez, beacon])
+
+        # noise realisation
+        noise_realisation = 0
+        if noise_sigma is not None:
+            noise_realisation = rng.normal(0, noise_sigma, size=len(beacon))
+
+        # Collect all data to be saved (with the first 3 values the E fields)
+        traces = np.array([antenna.Ex, antenna.Ey, antenna.Ez, beacon, noise_realisation])
 
         # add to relevant polarisation
         # and apply block filter
         dt = antenna.t[1] - antenna.t[0]
         for j, amp in enumerate(beacon_amplitudes):
-            traces[j] = block_filter(traces[j] + amp*beacon, dt, low_bp, high_bp)
+            traces[j] = block_filter(traces[j] + amp*beacon + noise_realisation, dt, low_bp, high_bp)
 
         append_antenna_hdf5( antennas_fname, antenna, traces, name='traces', prepend_time=True)
 

simulations/airshower_beacon_simulation/ac_show_signal_to_noise.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+# vim: indent=fdm ts=4
+
+"""
+Show Signal to noise for the original simulation signal,
+the beacon signal and the combined signal for each antenna
+"""
+
+import numpy as np
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+from collections import namedtuple
+
+from earsim import REvent
+import aa_generate_beacon as beacon
+import lib
+
+
+passband = namedtuple("passband", ['low', 'high'], defaults=[0, np.inf])
+
+def get_freq_spec(val,dt):
+    """From earsim/tools.py"""
+    fval = np.fft.fft(val)[:len(val)//2]
+    freq =  np.fft.fftfreq(len(val),dt)[:len(val)//2]
+    return fval, freq
+
+
+def bandpass_samples(samples, samplerate, band=passband()):
+        """
+        Bandpass the samples with this passband.
+        This is a hard filter.
+        """
+        fft, freqs = get_freq_spec(samples, samplerate)
+
+        fft[ ~ self.freq_mask(freqs) ] = 0
+
+        return np.fft.irfft(fft)
+
+def bandpass_mask(freqs, band=passband()):
+    low_pass = abs(freqs) <= band[1]
+    high_pass = abs(freqs) >= band[0]
+
+    return low_pass & high_pass
+
+def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True):
+    fft, freqs = get_freq_spec(samples, samplerate)
+
+    bandmask = bandpass_mask(freqs, band=band)
+
+    if normalise_bandsize:
+        bins = np.count_nonzero(bandmask, axis=-1)
+    else:
+        bins = 1
+
+    power = np.sum(np.abs(fft[bandmask])**2)
+
+    return power/bins
+
+def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None):
+    if noise_band is None:
+        noise_band = signal_band
+
+    if noise is None:
+        noise = samples
+
+    noise_power = bandpower(noise, samplerate, noise_band)
+
+    signal_power = bandpower(samples, samplerate, signal_band)
+
+    return (signal_power/noise_power)**0.5
+
+
+if __name__ == "__main__":
+    from os import path
+    import sys
+    import matplotlib
+    import os
+    if os.name == 'posix' and "DISPLAY" not in os.environ:
+        matplotlib.use('Agg')
+
+    f_beacon_band = (49e-3,55e-3) #GHz
+
+    fname = "ZH_airshower/mysim.sry"
+
+    fig_dir = "./figures/"
+    show_plots = not False
+
+    ####
+    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
+
+    # create fig_dir
+    if fig_dir:
+        os.makedirs(fig_dir, exist_ok=True)
+
+    # Read in antennas from file
+    _, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)
+    # Read original REvent
+    ev = REvent(fname)
+
+    # general properties
+    dt = antennas[0].t[1] - antennas[0].t[0] # ns
+    pb = passband(30e-3, 80e-3) # GHz
+    beacon_pb = passband(50e-3, 55e-3) # GHz
+
+    ##
+    ## Beacon vs Noise SNR
+    ##
+    if True:
+        beacon_snrs = [ signal_to_noise(ant.beacon, ant.noise, samplerate=1/dt, signal_band=beacon_pb) for ant in antennas ]
+
+        fig, ax = plt.subplots()
+        ax.set_title("Beacon SNR")
+        ax.set_xlabel("Antenna")
+        ax.set_ylabel("SNR")
+        ax.plot([ int(ant.name) for ant in antennas], beacon_snrs, 'o', ls='none')
+        
+        if fig_dir:
+            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".beacon_snr.pdf"))
+
+    ##
+    ## Airshower signal vs Noise SNR
+    ##
+    if True:
+        shower_snrs = [ signal_to_noise(ant.E_AxB, ant.noise, samplerate=1/dt, signal_band=pb) for ant in antennas ]
+
+        fig, ax = plt.subplots()
+        ax.set_title("Shower SNR")
+        ax.set_xlabel("Antenna")
+        ax.set_ylabel("SNR")
+        ax.plot([ int(ant.name) for ant in antennas], shower_snrs, 'o', ls='none')
+        
+        if fig_dir:
+            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".shower_snr.pdf"))
+    
+    if show_plots:
+        plt.show()