simulate and fit 2d gaussian that decays with exponential lifetime of 3 picoseconds in python

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

# Define the 2D Gaussian function
def gaussian_2d(x, y, amplitude, x0, y0, sigma_x, sigma_y):
    return amplitude * np.exp(-((x-x0)**2 / (2*sigma_x**2) + (y-y0)**2 / (2*sigma_y**2)))

# Define the exponential decay function
def exponential_decay(t, amplitude, lifetime):
    return amplitude * np.exp(-t / lifetime)

# Define the combined function
def combined_function(x, y, t, amplitude, x0, y0, sigma_x, sigma_y, lifetime):
    return gaussian_2d(x, y, amplitude, x0, y0, sigma_x, sigma_y) * exponential_decay(t, 1, lifetime)

# Generate some sample data
x = np.linspace(-10, 10, 100)
y = np.linspace(-10, 10, 100)
x, y = np.meshgrid(x, y)
t = np.linspace(0, 10, 100)

# Set the parameters
amplitude = 1
x0 = 0
y0 = 0
sigma_x = 2
sigma_y = 2
lifetime = 3  # in picoseconds

# Simulate the data
data = np.zeros((len(t), len(x), len(y)))
for i in range(len(t)):
    data[i] = combined_function(x, y, t[i], amplitude, x0, y0, sigma_x, sigma_y, lifetime)

# Add some noise to the data
data += np.random.normal(0, 0.1, data.shape)

# Define the function to fit
def func_to_fit(params, x, y, t):
    amplitude, x0, y0, sigma_x, sigma_y, lifetime = params
    return combined_function(x, y, t, amplitude, x0, y0, sigma_x, sigma_y, lifetime)

# Perform the fit
params = [amplitude, x0, y0, sigma_x, sigma_y, lifetime]
for i in range(len(t)):
    popt, pcov = curve_fit(func_to_fit, (x.ravel(), y.ravel(), np.full_like(x.ravel(), t[i])), data[i].ravel(), p0=params)
    params = popt

    # Print the fit parameters
    print(f"t = {t[i]} ps")
    print(f"Amplitude: {popt[0]}")
    print(f"x0: {popt[1]}")
    print(f"y0: {popt[2]}")
    print(f"Sigma x: {popt[3]}")
    print(f"Sigma y: {popt[4]}")
    print(f"Lifetime: {popt[5]} ps")
    print()
1863 chars
60 lines