from pyestimate import sin2d_param_estimate
import numpy as np
import matplotlib.pyplot as plt
#
# Define a 2D sinusoidal signal, corrupted by white gaussian noise
#
N = 20
M = 80
fn0 = 0.089
fm0 = 0.05
A = 1.23454
phi = 0.5789
n, m = np.meshgrid(np.arange(N), np.arange(M), indexing='ij')
sigma2 = 0.2
x = A * np.cos(2*np.pi*(fn0*n+fm0*m) + phi) + np.random.default_rng(seed=0).normal(scale=np.sqrt(sigma2), size=(N,M)) # Noisy 2D sinusoid
#
# Estimate sinusoidal parameters and plot estimation
#
A_hat, f_hat, phi_hat = sin2d_param_estimate(x, brute_Ns=100) # parameters estimation
print(f'Estimated amplitude: {A_hat: .3f} (true value: {A: .3f})')
print(f'Estimated frequencies: {f_hat[0]: .3f}, {f_hat[1]: .3f} (true values: {fn0: .3f}, {fm0: .3f})')
print(f'Estimated phase: {phi_hat: .3f} (true values: {phi: .3f})')   
plt.subplot(121)
plt.imshow(x, vmin=np.min(x), vmax=np.max(x))
plt.title('Noisy input 2D sinusoid')
plt.subplot(122)
plt.imshow(A_hat * np.cos(2*np.pi*(f_hat[0]*n+f_hat[1]*m) + phi_hat), vmin=np.min(x), vmax=np.max(x))
plt.title('Estimated 2D sinusoid')
plt.show()