ZH:lib/snr optional debugging plot in function

simulations/airshower_beacon_simulation/ac_show_signal_to_noise.py

@@ -49,12 +49,12 @@ if __name__ == "__main__":
         os.makedirs(fig_dir, exist_ok=True)
 
     # Read in antennas from file
-    f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)
+    f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname, traces_key='filtered_traces')
     _, __, txdata = beacon.read_tx_file(tx_fname)
 
     # general properties
     dt = antennas[0].t[1] - antennas[0].t[0] # ns
-    beacon_pb = lib.passband(f_beacon-1e-3, f_beacon+1e-3) # GHz
+    beacon_pb = lib.passband(f_beacon, None) # GHz
     beacon_amp = np.max(txdata['amplitudes'])# mu V/m
 
     # General Bandpass
@@ -62,19 +62,39 @@ if __name__ == "__main__":
     high_bp = args.passband_high if args.passband_high >= 0 else np.inf # GHz
     pb = lib.passband(low_bp, high_bp) # GHz
 
+    noise_pb = pb
+
     if args.use_passband: # Apply filter to raw beacon/noise to compare with Filtered Traces
         myfilter = lambda x: block_filter(x, dt, pb[0], pb[1])
     else: # Compare raw beacon/noise with Filtered Traces
         myfilter = lambda x: x
 
     ##
+    ## Debug plot of Beacon vs Noise SNR
+    ##
+    if True:
+        ant = antennas[0]
+
+        fig, ax = plt.subplots(figsize=figsize)
+        _ = lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), myfilter(ant.noise), samplerate=1/dt, signal_band=beacon_pb, noise_band=noise_pb, debug_ax=ax)
+
+        ax.set_title("Spectra and passband")
+        ax.set_xlabel("Frequency")
+        ax.set_ylabel("Amplitude")
+        low_x, high_x = min(beacon_pb[0], noise_pb[0]), max(beacon_pb[1] or 0, noise_pb[1])
+        ax.set_xlim(low_x, high_x)
+
+        if fig_dir:
+            fig.savefig(path.join(fig_dir, path.basename(__file__) + f".debug_plot.pdf"))
+    ##
     ## Beacon vs Noise SNR
     ##
     if True:
-        beacon_snrs = [ lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), myfilter(ant.noise), samplerate=1/dt, signal_band=beacon_pb, noise_band=pb) for ant in antennas ]
+        N_samples = len(antennas[0].beacon)
+        beacon_snrs = [ lib.signal_to_noise(myfilter(beacon_amp*ant.beacon), myfilter(ant.noise), samplerate=1/dt, signal_band=beacon_pb, noise_band=noise_pb) for i, ant in enumerate(antennas) ]
 
         fig, ax = plt.subplots(figsize=figsize)
-        ax.set_title("Maximum Beacon/Noise SNR")
+        ax.set_title(f"Maximum Beacon/Noise SNR (N_samples:{N_samples:.1e})")
         ax.set_xlabel("Antenna no.")
         ax.set_ylabel("SNR")
         ax.plot([ int(ant.name) for ant in antennas], beacon_snrs, 'o', ls='none')

simulations/airshower_beacon_simulation/lib/snr.py

@@ -1,6 +1,10 @@
 import numpy as np
 from collections import namedtuple
 
+from lib import direct_fourier_transform as dtft
+
+import matplotlib.pyplot as plt # for debug plotting
+
 passband = namedtuple("passband", ['low', 'high'], defaults=[0, np.inf])
 
 def get_freq_spec(val,dt):
@@ -27,30 +31,67 @@ def bandpass_mask(freqs, band=passband()):
 
     return low_pass & high_pass
 
-def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True):
-    fft, freqs = get_freq_spec(samples, samplerate)
+def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True, debug_ax=False):
+    bins = 0
+    fft, freqs = get_freq_spec(samples, 1/samplerate)
+    bandmask = [False]*len(freqs)
 
-    bandmask = bandpass_mask(freqs, band=band)
+    if band[1] is None:
+        # Only a single frequency given
+        # use a DTFT for finding the power
+        time = np.arange(0, len(samples), 1/samplerate)
 
-    if normalise_bandsize:
-        bins = np.count_nonzero(bandmask, axis=-1)
+        real, imag = dtft(band[0], time, samples)
+        power = np.sum(np.abs(real**2 + imag**2))
     else:
-        bins = 1
+        bandmask = bandpass_mask(freqs, band=band)
 
-    power = np.sum(np.abs(fft[bandmask])**2)
+        if normalise_bandsize:
+            bins = np.count_nonzero(bandmask, axis=-1)
+        else:
+            bins = 1
 
-    return power/bins
+        bins = max(1, bins)
 
-def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None):
+        power = 1/bins * np.sum(np.abs(fft[bandmask])**2)
+
+    # Prepare plotting variables if an Axes is supplied
+    if debug_ax:
+        if any(bandmask):
+            min_f, max_f = min(freqs[bandmask]), max(freqs[bandmask])
+        else:
+            min_f, max_f = 0, 0
+
+        if band[1] is None:
+            min_f, max_f = band[0], band[0]
+
+        if debug_ax is True:
+            debug_ax = plt.gca()
+
+        l = debug_ax.plot(freqs, np.abs(fft), alpha=0.9)
+
+        amp = np.sqrt(power)
+
+        if min_f != max_f:
+            debug_ax.plot( [min_f, max_f], [amp, amp], alpha=0.7, color=l[0].get_color(), ls='dashed')
+            debug_ax.axvspan(min_f, max_f, color=l[0].get_color(), alpha=0.2)
+        else:
+            debug_ax.plot( min_f, amp, '4', alpha=0.7, color=l[0].get_color(), ms=10)
+
+    return power
+
+def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None, debug_ax=False):
     if noise_band is None:
         noise_band = signal_band
 
     if noise is None:
         noise = samples
 
-    noise_power = bandpower(noise, samplerate, noise_band)
+    if debug_ax is True:
+        debug_ax = plt.gca()
 
-    signal_power = bandpower(samples, samplerate, signal_band)
+    noise_power = bandpower(noise, samplerate, noise_band, debug_ax=debug_ax)
+
+    signal_power = bandpower(samples, samplerate, signal_band, debug_ax=debug_ax)
 
     return (signal_power/noise_power)**0.5
-