pyskyline/demo/mosse_viz.py

134 lines
No EOL
3.7 KiB
Python

import numpy as np
import scipy.signal as signal
from scipy.fftpack import fft, fftshift, ifft
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
N = 256
def gaussian_filter1d(size,sigma):
filter_range = np.linspace(-int(size/2),int(size/2),size)
gaussian_filter = [1 / (sigma * np.sqrt(2*np.pi)) * np.exp(-x**2/(2*sigma**2)) for x in filter_range]
return gaussian_filter
def generate_signal(N:int) -> np.ndarray:
x = np.arange(1, N)
y = np.zeros((N))
for i in range(len(x)):
y[i] = np.random.normal(scale=1) + (y[i-1] if i > 1 else 0)
return np.convolve(y,gaussian_filter1d(N,1),'same')
if __name__ == '__main__':
np.random.seed(42)
f_full = generate_signal(2048)
x = np.arange(-180,180,360/N)
f = f_full[1024:1024+N]
h = f_full[1024:1024+N]
g = signal.gaussian(N, std=10,sym=True)
F = fft(f)
F_ = np.conjugate(F)
G = fft(g)
K_ = (G*F_)/(F*F_)
H = fft(h)
r = ifft(H*K_)
# ==========================================
fig, (ax1, ax2, ax3, ax4, ax5, ax6, ax7) = plt.subplots(7, 1, gridspec_kw={'height_ratios':[4,4,4,1,1,1,1]})
ax1.set_title('Terrain')
plt_f, = ax1.plot(x,f)
plt_h, = ax1.plot(x,h)
ax2.set_title('MOSSE response signal')
ax2.set_ylim([0, 1.2])
line_r = ax2.axvline(x=-N//2+np.argmax(abs(r)), color='r')
plt_r, = ax2.plot(x,abs(r))
plt_r2, = ax2.plot(x,abs(r))
ax3.set_title('Gaussian')
plt_g, = ax3.plot(x,g)
ax1.set_xlim([-180,180])
ax2.set_xlim([-180,180])
ax3.set_xlim([-180,180])
slider1 = Slider(ax4, 'sigma', 0.3, 10, valinit=0.1)
slider2 = Slider(ax5, 'shift', -N//2 , N//2, valinit=0, valstep=1)
slider3 = Slider(ax6, 'seed', 0 , 50, valinit=0, valstep=1)
slider4 = Slider(ax7, 'N', 128 , 1024, valinit=256, valstep=8)
sigma = 0.3
shift = 0
def update():
# K_ = (G*F_)/(F*F_)
window = np.ones((N)) #signal.windows.hamming(N)
H = fft(h*window)
F = fft(f*window)
R = H*G/F
r = ifft(R)
s = np.argmax(abs(r))
r2 = np.copy(r)
r2[s-5:s+5] = 0
plt_g.set_data(x,g)
plt_r.set_data(x,abs(r))
plt_r2.set_data(x,abs(r2))
plt_h.set_data(x,h)
plt_f.set_data(x,f)
ax1.set_ylim([min(np.min(h),np.min(f))-1, max(np.max(h),np.max(f))+1])
ax2.set_ylim([0, np.max(r)+0.2])
line_r.set_xdata(round((np.argmax(abs(r))/N-0.5)*360))
fig.canvas.draw_idle()
def update_sigma(val):
global g, G, sigma, N
sigma = val
g = signal.gaussian(N, std=sigma,sym=True)
G = fft(g)
update()
def update_shift(val):
global shift, H, h
shift = -val
h = f_full[1024+round(shift*(N/360)):1024+N+round(shift*(N/360))]
noise = np.random.normal(0,0.5, N)
h = h+noise
update()
def update_seed(val):
global f_full, f, h, F, H, F_, shift
np.random.seed(val)
f_full = generate_signal(2048)
f = f_full[1024:1024+N]
h = f_full[1024+round(shift*(N/360)):1024+N+round(shift*(N/360))]
noise = np.random.normal(0,0.5, N)
h = h+noise
update()
def update_n(val):
global g, G, N, f_full, f, h, F, H, F_, shift, x, sigma, slider2
N = val
x = np.arange(-180,180,360/N)
g = signal.gaussian(N, std=sigma,sym=True)
G = fft(g)
f = f_full[1024:1024+N]
h = f_full[1024+shift:1024+N+shift]
noise = np.random.normal(0,0.5, N)
h = h+noise
update()
slider1.on_changed(update_sigma)
slider2.on_changed(update_shift)
slider3.on_changed(update_seed)
slider4.on_changed(update_n)
plt.subplots_adjust(hspace=0.5)
plt.show()