Fitting arbitrary functions with lmfit

sample.py

import numpy as np
import matplotlib.pyplot as plt

from lmfit import Parameters, fit_report, minimize

# First generate our testdata

real_mu = 3.4
real_sigma = 3.1
real_c = 7.1
real_offset = -3
real_noise_lvl = 1

data_x = np.linspace(-10, 10, 100)
data_y = real_c * np.exp(-0.5 * (real_mu - data_x)**2 / (real_sigma**2)) * 1.0 / (real_sigma*np.sqrt(2 * np.pi)) + real_offset
data_y_noisy = data_y + (np.random.rand(data_y.shape[0]) - 0.5) * real_noise_lvl

fig, ax = plt.subplots()
ax.set_title("Simulated input data")
ax.plot(data_x, data_y, 'g--', label = "Real data")
ax.plot(data_x, data_y_noisy, label = "Noisy data")
ax.grid()
ax.legend()
plt.savefig("lmfit_01_indata.png")
plt.savefig("lmfit_01_indata.svg")
plt.show()

# Now the fit function

def modelFunction(pars, x, data = None):
    # Convert Parameters to a dictionary

    vals = pars.valuesdict()

    # For readability give the parameters variables

    amp = pars["amp"]
    mu = pars["mu"]
    sigma = pars["sigma"]
    offset = pars["offset"]

    # Evaluate the current model at position x

    modelValue = 1.0 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-0.5 * (mu - x)**2 / (sigma**2)) * amp + offset

    # Return the residual if we've also got data, else the model value

    if data is None:
        return modelValue
    else:
        return modelValue - data

# Lets initialize parameters with a naive guess

fit_pars = Parameters()
fit_pars.add("amp", value = np.max(data_y_noisy), vary = True)
fit_pars.add("mu", value = data_x[np.argmax(data_y_noisy)], vary = True)
fit_pars.add("sigma", value = 1, vary = True)
fit_pars.add("offset", value = np.min(data_y_noisy), vary = True)

out = minimize(modelFunction, fit_pars, args=(data_x,), kws = { 'data' : data_y_noisy })

# Show only chi squared

print(f"Chi squared: {out.chisqr}")

# Show our fit report

print(fit_report(out))

# Plot the result

fig, ax = plt.subplots()
ax.plot(data_x, data_y_noisy, 'x', label = "Noisy input data")
ax.plot(data_x, modelFunction(out.params, data_x), 'r--', label = "Fitted function")
ax.plot(data_x, data_y, 'g', label = "Original function")
ax.grid()
ax.legend()
plt.savefig("lmfit_01_outdata.png")
plt.savefig("lmfit_01_outdata.svg")
plt.show()