ZH: ba_beacon_phase move functions to ./lib.py

This commit is contained in:
Eric Teunis de Boone 2022-11-18 16:27:20 +01:00
parent 2240d67c1c
commit 726c506816
2 changed files with 127 additions and 96 deletions

View file

@ -1,104 +1,12 @@
#!/usr/bin/env python3
# vim: fdm=indent ts=4
import numpy as np
import h5py
import lib
import matplotlib.pyplot as plt
import numpy as np
import aa_generate_beacon as beacon
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
import lib
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")

View file

@ -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