ZH: ba_beacon_phase move functions to ./lib.py

simulations/airshower_beacon_simulation/ba_beacon_phases.py

@@ -1,104 +1,12 @@
 #!/usr/bin/env python3
 # vim: fdm=indent ts=4
 
-import numpy as np
 import h5py
-
+import matplotlib.pyplot as plt
-import lib
+import numpy as np
 
 import aa_generate_beacon as beacon
-
+import lib
-from lib import direct_fourier_transform
-
-from numpy.polynomial import Polynomial
-
-def find_beacon_in_traces(
-        traces,
-        t_trace,
-        f_beacon_estimate = 50e6,
-        frequency_fit = False,
-        N_test_freqs = 5e2,
-        f_beacon_estimate_band = 0.01,
-        amp_cut = 0.8
-        ):
-    """
-    f_beacon_band is inclusive
-
-    traces is [trace, trace, trace, ..  ]
-    """
-
-    amplitudes = np.zeros(len(traces))
-    phases = np.zeros(len(traces))
-    frequencies = np.zeros(len(traces))
-
-    if frequency_fit: # fit frequency
-        test_freqs = f_beacon_estimate + f_beacon_estimate_band * np.linspace(-1, 1, int(N_test_freqs)+1)
-        ft_amp_gen = direct_fourier_transform(test_freqs, t_trace, (x for x in traces))
-
-        n_samples = len(t_trace)
-
-        for i, ft_amp in enumerate(ft_amp_gen):
-            real, imag = ft_amp
-            amps = 1/n_samples * ( real**2 + imag**2)**0.5
-
-            # find frequency peak and surrounding
-            # bins valid for parabola fitting
-            max_amp_idx = np.argmax(amps)
-            max_amp = amps[max_amp_idx]
-
-            if True:
-                frequencies[i] = test_freqs[max_amp_idx]
-                continue
-
-            valid_mask = amps >= amp_cut*max_amp
-
-            if True: # make sure not to use other peaks
-                lower_mask = valid_mask[0:max_amp_idx]
-                upper_mask = valid_mask[max_amp_idx:]
-
-                if any(lower_mask):
-                    lower_end = np.argmin(lower_mask[::-1])
-                else:
-                    lower_end = max_amp_idx
-
-                if any(upper_mask):
-                    upper_end = np.argmin(upper_mask)
-                else:
-                    upper_end = 0
-
-
-                valid_mask[0:(max_amp_idx - lower_end)] = False
-                valid_mask[(max_amp_idx + upper_end):] = False
-
-            if all(~valid_mask):
-                frequencies[i] = np.nan
-                continue
-
-            # fit Parabola
-            parafit = Polynomial.fit(test_freqs[valid_mask], amps[valid_mask], 2)
-            func = parafit.convert()
-
-            # find frequency where derivative is 0
-            deriv = func.deriv(1)
-            freq = deriv.roots()[0]
-
-            frequencies[i] = freq
-    else:
-        frequencies[:] = f_beacon_estimate
-
-    # evaluate fourier transform at freq for each trace
-    for i, freq in enumerate(frequencies):
-        if freq is np.nan:
-            phases[i] = np.nan
-            amplitudes[i] = np.nan
-            continue
-
-        real, imag = direct_fourier_transform(freq, t_trace, traces[i])
-
-        phases[i] = np.arctan2(real, imag)
-        amplitudes[i] = 1/len(t_trace) * (real**2 + imag**2)**0.5
-
-    return frequencies, phases, amplitudes
 
 if __name__ == "__main__":
     from os import path
@@ -181,7 +89,6 @@ if __name__ == "__main__":
 
     # show histogram of found frequencies
     if True:
-        import matplotlib.pyplot as plt
         if True or allow_frequency_fitting:
             fig, ax = plt.subplots()
             ax.set_xlabel("Frequency")

simulations/airshower_beacon_simulation/lib.py

@@ -3,6 +3,7 @@
 Library for this simulation
 """
 import numpy as np
+from numpy.polynomial import Polynomial
 
 from earsim import Antenna
 
@@ -38,6 +39,40 @@ def beacon_from(tx, rx, f, t=0, t0=0, c_light=3e8, radiate_rsq=True, amplitude=1
     return sine_beacon(f, t, t0=t0, amplitude=amplitude,**kwargs)
 
 
+def ft_corr_vectors(freqs, time):
+    """
+    Get the cosine and sine terms for freqs at time.
+
+    Takes the outer product of freqs and time.
+    """
+    freqtime = np.outer(freqs, time)
+
+    c_k = np.cos(2*np.pi*freqtime)
+    s_k = np.sin(2*np.pi*freqtime)
+
+    return c_k, s_k
+
+def direct_fourier_transform(freqs, time, samplesets_iterable):
+    """
+    Determine the fourier transform of each sampleset in samplesets_iterable at freqs.
+
+    The samplesets are expected to have the same time vector.
+
+    Returns either a generator to return the fourier transform for each sampleset
+    if samplesets_iterable is a generator
+    or a numpy array.
+    """
+
+    c_k, s_k = ft_corr_vectors(freqs, time)
+
+    if not hasattr(samplesets_iterable, '__len__') and hasattr(samplesets_iterable, '__iter__'):
+        # samplesets_iterable is an iterator
+        # return a generator containing (real, imag) amplitudes
+        return ( (np.dot(c_k, samples), np.dot(s_k, samples)) for samples in samplesets_iterable )
+
+    # Numpy array
+    return np.dot(c_k, samplesets_iterable), np.dot(s_k, samplesets_iterable)
+
 def phase_field_from_tx(x, y, tx, f_beacon, c_light=3e8, t0=0, wrap_phase=True, return_meshgrid=True):
     xs, ys = np.meshgrid(x, y, sparse=True)
 
@@ -62,3 +97,92 @@ def phase_mod(phase, low=np.pi):
     interval $[-low, 2\pi - low)$.
     """
     return (phase + low) % (2*np.pi) - low
+
+
+def find_beacon_in_traces(
+        traces,
+        t_trace,
+        f_beacon_estimate = 50e6,
+        frequency_fit = False,
+        N_test_freqs = 5e2,
+        f_beacon_estimate_band = 0.01,
+        amp_cut = 0.8
+        ):
+    """
+    f_beacon_band is inclusive
+
+    traces is [trace, trace, trace, ..  ]
+    """
+
+    amplitudes = np.zeros(len(traces))
+    phases = np.zeros(len(traces))
+    frequencies = np.zeros(len(traces))
+
+    if frequency_fit: # fit frequency
+        test_freqs = f_beacon_estimate + f_beacon_estimate_band * np.linspace(-1, 1, int(N_test_freqs)+1)
+        ft_amp_gen = direct_fourier_transform(test_freqs, t_trace, (x for x in traces))
+
+        n_samples = len(t_trace)
+
+        for i, ft_amp in enumerate(ft_amp_gen):
+            real, imag = ft_amp
+            amps = 1/n_samples * ( real**2 + imag**2)**0.5
+
+            # find frequency peak and surrounding
+            # bins valid for parabola fitting
+            max_amp_idx = np.argmax(amps)
+            max_amp = amps[max_amp_idx]
+
+            if True:
+                frequencies[i] = test_freqs[max_amp_idx]
+                continue
+
+            valid_mask = amps >= amp_cut*max_amp
+
+            if True: # make sure not to use other peaks
+                lower_mask = valid_mask[0:max_amp_idx]
+                upper_mask = valid_mask[max_amp_idx:]
+
+                if any(lower_mask):
+                    lower_end = np.argmin(lower_mask[::-1])
+                else:
+                    lower_end = max_amp_idx
+
+                if any(upper_mask):
+                    upper_end = np.argmin(upper_mask)
+                else:
+                    upper_end = 0
+
+
+                valid_mask[0:(max_amp_idx - lower_end)] = False
+                valid_mask[(max_amp_idx + upper_end):] = False
+
+            if all(~valid_mask):
+                frequencies[i] = np.nan
+                continue
+
+            # fit Parabola
+            parafit = Polynomial.fit(test_freqs[valid_mask], amps[valid_mask], 2)
+            func = parafit.convert()
+
+            # find frequency where derivative is 0
+            deriv = func.deriv(1)
+            freq = deriv.roots()[0]
+
+            frequencies[i] = freq
+    else:
+        frequencies[:] = f_beacon_estimate
+
+    # evaluate fourier transform at freq for each trace
+    for i, freq in enumerate(frequencies):
+        if freq is np.nan:
+            phases[i] = np.nan
+            amplitudes[i] = np.nan
+            continue
+
+        real, imag = direct_fourier_transform(freq, t_trace, traces[i])
+
+        phases[i] = np.arctan2(real, imag)
+        amplitudes[i] = 2/len(t_trace) * (real**2 + imag**2)**0.5
+
+    return frequencies, phases, amplitudes