m-thesis-introduction/fourier/06_direct_fourier.py

182 lines
6.2 KiB
Python
Raw Permalink Normal View History

#!/usr/bin/env python3
# vim: fdm=indent ts=4
__doc__ = \
"""
Show how the fourier transform can be calculated
in a continuous fashion
"""
if __name__ == "__main__":
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mylib.fft import direct_fourier_transform
n_samples = 1000
nphi = 100
f_beacon = 2.3434
noise_A = 5
t = np.linspace(0, 100, n_samples+1)
phi_in = np.linspace(0, 2*np.pi, nphi)
2022-11-14 11:41:35 +01:00
test_freqs = f_beacon + 0.1 * np.linspace(-1, 1, 100+1)
phi_out = np.zeros( (len(phi_in), len(test_freqs)) )
amp_out = np.zeros( (len(phi_in), len(test_freqs)) )
if not True:
# Same length samples
# means we can precalculate the c_k and s_k terms
c_k, s_k = ft_corr_vectors(test_freqs, t)
for i, phi in enumerate(phi_in):
s = np.sin(2*np.pi*t*f_beacon + phi) + noise_A * np.random.normal(size=len(t))
real = np.dot(c_k, s)
imag = np.dot(s_k, s)
phi_out[i] = (np.arctan2(real, imag))
amp_out[i] = (2/len(t) * (real**2 + imag**2)**0.5)
else:
sampleset_gen = ( np.sin(2*np.pi*t*f_beacon + phi) + noise_A*np.random.normal(size=len(t)) for phi in phi_in )
ft_amp_gen = direct_fourier_transform(test_freqs, t, sampleset_gen)
for i, ft_amp in enumerate(ft_amp_gen):
real = ft_amp[0]
imag = ft_amp[1]
phi_out[i] = np.arctan2(real, imag)
amp_out[i] = 2/len(t) * (real**2 + imag**2)**0.5
######
# Figures
######
import matplotlib.colors as colors
cmap = plt.cm.plasma
freq_norm = colors.Normalize(vmin=np.amin(test_freqs), vmax=np.amax(test_freqs))
freq_cmap = cmap(freq_norm(test_freqs))
try:
# Amplitudes Histogram
if True:
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.set_xlabel("Amplitude")
ax.set_ylabel("Frequency")
ax.set_zlabel("Count")
for j, amp in enumerate(amp_out.T):
# per test_freq
counts, edges = np.histogram(amp, bins='auto', )
l = ax.plot(edges[:-1], counts, zs=test_freqs[j], zdir='y', color=freq_cmap[j])
ax.add_collection3d(plt.fill_between(edges[:-1], 0, counts, color=l[0].get_color(), alpha=0.3), zs=test_freqs[j], zdir='y')
ax.view_init(elev=20., azim=-35)
elif False:
fig, ax = plt.subplots()
ax.set_xlabel("Amplitude")
ax.set_ylabel("Count")
#
for j, amp in enumerate(amp_out.T):
# per test_freq
ax.hist(amp, histtype='step', bins='auto', color=freq_cmap[j])
# Single Amplitude / Frequency plot showing frequency fitting
freq_out = None
if True:
2022-11-14 11:41:35 +01:00
from numpy.polynomial import Polynomial as P
freq_out = np.zeros(len(phi_in))
amp_cut = 0.5
fig, ax = plt.subplots()
2022-11-14 11:41:35 +01:00
ax.set_title("Frequency estimation by parabola fitting.\nStars are used for the parabola fit, vertical line is where $\\partial_f = 0 $")
ax.set_xlabel("Frequency")
ax.set_ylabel("Amplitude")
ax.axvline(f_beacon, lw=5, ls=(0,(5,5)))
for j, amp in enumerate(amp_out):
if j > 2:
continue
max_amp_idx = np.argmax(amp)
max_amp = amp[max_amp_idx]
# filter amplitudes below amp_cut*max_amp
valid_mask = amp >= 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:]
lower_end = np.argmin(lower_mask[::-1])
upper_end = np.argmin(upper_mask)
valid_mask[0:(max_amp_idx - lower_end)] = False
valid_mask[(max_amp_idx + upper_end):] = False
p_fit = P.fit(test_freqs[valid_mask], amp[valid_mask], 2)
func = p_fit.convert()
2022-11-14 11:41:35 +01:00
# Find frequency of derivative == 0
deriv = func.deriv(1)
freq = deriv.roots()[0]
freq_out[j] = freq
2022-11-14 11:41:35 +01:00
l = ax.plot(test_freqs, amp, marker='.')
ax.plot(test_freqs[valid_mask], amp[valid_mask], marker='*', color=l[0].get_color())
ax.axvline(freq_out[j], color=l[0].get_color())
2022-11-14 11:41:35 +01:00
if True: # plot the fit
tmp_test_freqs = test_freqs[max_amp_idx] + 0.05*np.linspace(-1,1,101, endpoint=True)
func_amps = func(tmp_test_freqs)
func_amps_idx = func_amps > 0
func_amps = func_amps[func_amps_idx]
tmp_test_freqs = tmp_test_freqs[func_amps_idx]
ax.plot(tmp_test_freqs, func_amps, ls='dotted', color=l[0].get_color())
# Amplitudes figure
if True:
fig, ax = plt.subplots()
ax.set_ylabel("Amplitude")
ax.set_xlabel("Frequency")
if True:
for j, amp in enumerate(amp_out.T):
#per test_freq
ax.plot(np.tile(test_freqs[j], len(amp)), amp, marker='.', color=freq_cmap[j], alpha=max(0.05, 0))
ax.errorbar(test_freqs, np.mean(amp_out, axis=0), yerr=np.std(amp_out, axis=0))
ax.plot(test_freqs, np.mean(amp_out, axis=0), marker='*', color='red', zorder=6)
# Phase in vs. Phase out figure
if True:
fig, ax = plt.subplots()
amp_cut = 0.8
ax.set_title(f"Measured phases passing amplitude > {amp_cut}")
ax.set_xlabel("Phase in")
ax.set_ylabel("Phase out")
for j, phi in enumerate(phi_out.T):
if np.mean(amp_out[:,j]) < amp_cut:
# ignore when the amplitudes are not close to 1
continue
# per test_freq
ax.plot(phi_in, np.unwrap(phi), label='f-test_f:{:.2e}'.format(f_beacon-test_freqs[j]))#, color=freq_cmap[j])
ax.legend()
finally:
plt.show()