simulate and fit 2d gaussian in python

main.py
import numpy as np
import matplotlib.pyplot as plt

# Define the parameters of the Gaussian
x0, y0 = 0, 0  # center
sigma_x, sigma_y = 1, 1  # standard deviations
amp = 1  # amplitude

# Create a grid of x and y values
x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)

# Simulate the 2D Gaussian
Z = amp * np.exp(-((X - x0) / sigma_x)**2 - ((Y - y0) / sigma_y)**2)

# Plot the simulated Gaussian
plt.imshow(Z, cmap='viridis', origin='lower')
plt.colorbar()
plt.show()
500 chars
21 lines

main.py
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

# Define the function to fit
def gaussian_2d(xy, amp, x0, y0, sigma_x, sigma_y):
    x, y = xy
    return amp * np.exp(-((x - x0) / sigma_x)**2 - ((y - y0) / sigma_y)**2)

# Create a grid of x and y values
x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)
XY = (X.ravel(), Y.ravel())

# Simulate the data
Z = gaussian_2d(XY, 1, 0, 0, 1, 1) + np.random.normal(0, 0.1, len(Z))

# Fit the data
popt, pcov = curve_fit(gaussian_2d, XY, Z, p0=[1, 0, 0, 1, 1])

# Print the fitted parameters
print(popt)

# Plot the fitted Gaussian
Z_fit = gaussian_2d(XY, *popt).reshape(X.shape)
plt.imshow(Z_fit, cmap='viridis', origin='lower')
plt.colorbar()
plt.show()
769 chars
30 lines