pyskyline/example/mosse_viz.py
Alexandre Foucher f85f23776a 3/10/24
2024-10-03 15:47:52 +02:00

130 lines
No EOL
4.3 KiB
Python
Executable file

import numpy as np
import scipy.signal as signal
from scipy.fftpack import fft, ifft
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
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')
class UI:
def __init__(self) -> None:
self._size_F = 256
self._size_H = 256
self._sigma = 1.1
self._seed = 0
self._shift = 0
self._fig, self._axs = self.init_ui()
self.update()
def init_ui(self) -> None:
fig, (ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8) = plt.subplots(8, 1, gridspec_kw={'height_ratios':[4,4,4,1,1,1,1,1]})
ax1.set_title('Terrain')
self.plt_f, = ax1.plot([])
self.plt_h, = ax1.plot([])
ax2.set_title('MOSSE response signal')
ax2.set_ylim([0, 1.2])
self.line_r = ax2.axvline()
self.plt_r, = ax2.plot([])
self.plt_r2, = ax2.plot([])
ax3.set_title('Gaussian')
self.plt_g, = ax3.plot([])
ax1.set_xlim([-180,180])
ax2.set_xlim([-180,180])
ax3.set_xlim([-180,180])
self._slider1 = Slider(ax4, 'sigma', 0.1, 5, valinit=self._sigma)
self._slider2 = Slider(ax5, 'shift', -180, 180, valinit=self._shift, valstep=1)
self._slider3 = Slider(ax6, 'seed', 0, 50, valinit=self._seed, valstep=1)
self._slider4 = Slider(ax7, 'Size f', 64, 1024, valinit=self._size_F, valstep=8)
self._slider5 = Slider(ax8, 'Size h', 64, 1024, valinit=self._size_H, valstep=8)
self._slider1.on_changed(self.update_sigma)
self._slider2.on_changed(self.update_shift)
self._slider3.on_changed(self.update_seed)
self._slider4.on_changed(self.update_size_f)
self._slider5.on_changed(self.update_size_h)
fig.subplots_adjust(hspace=0.75)
return fig, [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8]
def update(self):
np.random.seed(self._seed)
g = signal.windows.gaussian(self._size_F, std=self._sigma,sym=True)
f = generate_signal(self._size_F) * signal.windows.hamming(self._size_F)
h = np.zeros((self._size_F))
s = self._size_F//2-self._size_H//2
h[s:s+self._size_H] = (f[s:s+self._size_H] + np.random.normal(0,0.5, self._size_H)) * signal.windows.hamming(self._size_H)
h = np.roll(h,int(self._shift/180*self._size_F//2))
F = fft(f)
G = fft(g)
H = fft(h)
R = H*G/F
r = ifft(R)
s = np.argmax(abs(r))
r2 = np.copy(r)
r2[s-5:s+5] = 0
x = np.linspace(-180,180,self._size_F)
self.plt_g.set_data(x, g)
self.plt_r.set_data(x,abs(r))
self.plt_r2.set_data(x,abs(r2))
self.plt_f.set_data(x,f)
self.plt_h.set_data(x,h)
self._axs[0].set_ylim([min(np.min(h),np.min(f))-1, max(np.max(h),np.max(f))+1])
self._axs[1].set_ylim([0, np.max(r)+0.2])
self._axs[2].set_ylim([0, np.max(g)*1.1])
self.line_r.set_xdata([round((np.argmax(abs(r))/self._size_F-0.5)*360)])
self._fig.canvas.draw_idle()
def update_sigma(self, val):
self._sigma = val
self.update()
def update_size_f(self, val):
if val < self._size_H:
self._slider4.eventson = False
self._slider4.set_val(max(val, self._size_H))
self._slider4.eventson = True
self._size_F = max(val, self._size_H)
self.update()
def update_size_h(self, val):
if val > self._size_F:
self._slider5.eventson = False
self._slider5.set_val(min(val, self._size_F))
self._slider5.eventson = True
self._size_H = min(val, self._size_F)
self.update()
def update_seed(self, val):
self._seed = val
self.update()
def update_shift(self, val):
self._shift = val
self.update()
def show(self) -> None:
plt.show()
if __name__ == '__main__':
ui = UI()
ui.show()