Redo the 'example_fit_multi_datasets.py' with built-in models

[mstekiel]

In the process of learning how to fit multiple datasets with lmfit I found the TODO flag in the multifit example, and decided to do the job.

Comments:

1. Refined amplitude values are slightly different than in the previous example, as the previous gaussian definition did not include the 1/(sigma*sqrt(2 pi) prefactor.
2. The number of datasets is now controlled by single parameter N. It takes few second to fit N=50, really impressive for lmfit!

Cheers!

"""
Fit Multiple Data Sets
======================

Fitting multiple (simulated) Gaussian data sets simultaneously.

All minimizers require the residual array to be one-dimensional. Therefore, in
the ``objective`` function we need to ``flatten`` the array before returning it.

"""
import matplotlib.pyplot as plt
import numpy as np

from lmfit import Parameters, minimize, report_fit
from lmfit.models import GaussianModel

###############################################################################
# Create N simulated Gaussian data sets
N = 5
np.random.seed(2021)
x = np.linspace(-1, 2, 151)
data = []
for _ in np.arange(N):
    params = Parameters()
    params.add('amplitude', value=0.60 + 9.50*np.random.rand())
    params.add('center', value=-0.20 + 1.20*np.random.rand())
    params.add('sigma', value=0.25 + 0.03*np.random.rand())
    dat = GaussianModel().eval(x=x, params=params) + np.random.normal(size=x.size, scale=0.1)
    data.append(dat)
data = np.array(data)

###############################################################################
# The objective function will extract and evaluate a Gaussian from the compound model
def objective(params, x, data, model):
    """Calculate total residual for fits of Gaussians to several data sets."""
    ndata, _ = data.shape
    resid = 0.0*data[:]

    # make residual per data set
    for i in range(ndata):
        components = model.components[i].eval(params=params, x=x)
        resid[i, :] = data[i, :] - components

    # now flatten this to a 1D array, as minimize() needs
    return resid.flatten()

###############################################################################
# Create a composite model by adding Gaussians
model_arr = [GaussianModel(prefix=f'n{i+1}_') for i,_ in enumerate(data)]
model = sum(model_arr[1:], start=model_arr[0])

###############################################################################
# Prepare the fitting parameters and constrain n2_sigma, ..., nN_sigma to be equal to n1_sigma
fit_params = model.make_params()
for iy, y in enumerate(data):
    fit_params.add(f'n{iy+1}_amplitude', value=0.5, min=0.0, max=200)
    fit_params.add(f'n{iy+1}_center', value=0.4, min=-2.0, max=2.0)
    fit_params.add(f'n{iy+1}_sigma', value=0.3, min=0.01, max=3.0)

    if iy>0:
        fit_params[f'n{iy+1}_sigma'].expr = 'n1_sigma'

###############################################################################
# Run the global fit and show the fitting result
out = minimize(objective, fit_params, args=(x, data, model))
report_fit(out.params)


###############################################################################
# Plot the data sets and fits
plt.figure()
for i, y in enumerate(data):
    components = model.eval_components(params=out.params, x=x)
    plt.plot(x, y, 'o', x, components[f'n{i+1}_'], '-')

plt.show()

[newville]

@mstekiel Thanks! If you're willing to make a pull request for that, we'll accept it as is. If not, let us know and we will update the example.

[mstekiel]

Glad it's helpful. Feel free to update it on your own.

[SchmelleDev]

Many thanks to @mstekiel.
I needed a more general version with datasets of different x values.
I got it working (with the help of AI )...

I would like to see the original version as well as the updated one in the examples section.

import matplotlib.pyplot as plt
import numpy as np
from lmfit import Parameters, minimize, report_fit
from lmfit.models import GaussianModel

###############################################################################
# Create N simulated Gaussian data sets with variable x lengths
N = 5
np.random.seed(2021)
data = []
x_values = []

for _ in range(N):
    x = np.linspace(-1, 2, np.random.randint(100, 200))  # Variable x length
    params = Parameters()
    params.add('amplitude', value=0.60 + 9.50 * np.random.rand())
    params.add('center', value=-0.20 + 1.20 * np.random.rand())
    params.add('sigma', value=0.25 + 0.03 * np.random.rand())
    dat = GaussianModel().eval(x=x, params=params) + np.random.normal(size=x.size, scale=0.1)
    data.append(dat)
    x_values.append(x)

data = np.array(data, dtype=object)  # Store data as an object array to allow variable lengths

###############################################################################
# The objective function will extract and evaluate a Gaussian from the compound model
def objective(params, x_values, data, model):
    """Calculate total residual for fits of Gaussians to several data sets."""
    resid = []

    # Make residual per data set
    for i in range(len(data)):
        components = model.components[i].eval(params=params, x=x_values[i])
        resid.append(data[i] - components)

    # Flatten the residuals into a 1D array
    return np.hstack(resid)

###############################################################################
# Create a composite model by adding Gaussians
model_arr = [GaussianModel(prefix=f'n{i+1}_') for i in range(N)]
model = sum(model_arr[1:], start=model_arr[0])

###############################################################################
# Prepare the fitting parameters and constrain n2_sigma, ..., nN_sigma to be equal to n1_sigma
fit_params = model.make_params()
for iy in range(N):
    fit_params.add(f'n{iy+1}_amplitude', value=0.5, min=0.0, max=200)
    fit_params.add(f'n{iy+1}_center', value=0.4, min=-2.0, max=2.0)
    fit_params.add(f'n{iy+1}_sigma', value=0.3, min=0.01, max=3.0)

    if iy > 0:
        fit_params[f'n{iy+1}_sigma'].expr = 'n1_sigma'

###############################################################################
# Run the global fit and show the fitting result
out = minimize(objective, fit_params, args=(x_values, data, model))
report_fit(out.params)

###############################################################################
# Plot the data sets and fits
plt.figure()
for i in range(N):
    components = model.eval_components(params=out.params, x=x_values[i])
    plt.plot(x_values[i], data[i], 'o', label=f'Dataset {i+1}')
    plt.plot(x_values[i], components[f'n{i+1}_'], '-', label=f'Fit {i+1}')

plt.legend()
plt.show()