pyskyline/example/mosse_viz.py

134 lines
No EOL
3.7 KiB
Python
Executable file

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()