ZH: move ac_* function definitions into lib/snr.py

simulations/airshower_beacon_simulation/ac_show_signal_to_noise.py

@@ -10,67 +10,12 @@ 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
@@ -99,8 +44,8 @@ if __name__ == "__main__":
 
     # general properties
     dt = antennas[0].t[1] - antennas[0].t[0] # ns
-    pb = passband(30e-3, 80e-3) # GHz
-    beacon_pb = passband(f_beacon-1e-3, f_beacon+1e-3) # GHz
+    pb = lib.passband(30e-3, 80e-3) # GHz
+    beacon_pb = lib.passband(f_beacon-1e-3, f_beacon+1e-3) # GHz
 
     beacon_amp = np.max(txdata['amplitudes'])# mu V/m
 
@@ -108,7 +53,7 @@ if __name__ == "__main__":
     ## Beacon vs Noise SNR
     ##
     if True:
-        beacon_snrs = [ signal_to_noise(beacon_amp*ant.beacon, ant.noise, samplerate=1/dt, signal_band=beacon_pb) for ant in antennas ]
+        beacon_snrs = [ lib.signal_to_noise(beacon_amp*ant.beacon, ant.noise, samplerate=1/dt, signal_band=beacon_pb) for ant in antennas ]
 
         fig, ax = plt.subplots()
         ax.set_title("Maximum Beacon SNR")
@@ -123,7 +68,7 @@ if __name__ == "__main__":
     ## 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 ]
+        shower_snrs = [ lib.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("Total (Signal+Beacon+Noise) SNR")

simulations/airshower_beacon_simulation/lib/__init__.py

@@ -1,2 +1,3 @@
 from .lib import *
 from . import rit
+from .snr import *

simulations/airshower_beacon_simulation/lib/snr.py

@@ -0,0 +1,56 @@
+import numpy as np
+from collections import namedtuple
+
+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
+