{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Curve fitting with `lmfit`\n",
    "\n",
    "In this section, we will cover basic curve fitting using [lmfit](https://lmfit.github.io/lmfit-py/) for reference purposes. For detailed information, please refer to the [lmfit documentation](https://lmfit.github.io/lmfit-py/).\n",
    "\n",
    "If you are already familiar with [lmfit](https://lmfit.github.io/lmfit-py/), you can skip to the [next section](./modelfit)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden",
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = [\"svg\", \"pdf\"]\n",
    "plt.rcParams[\"figure.dpi\"] = 96\n",
    "plt.rcParams[\"image.cmap\"] = \"viridis\"\n",
    "plt.rcParams[\"figure.figsize\"] = (4, 2.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by defining a model function and the data to fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poly1(x, a, b):\n",
    "    return a * x + b\n",
    "\n",
    "\n",
    "# Generate some toy data\n",
    "x = np.linspace(0, 10, 20)\n",
    "y = poly1(x, 1, 2)\n",
    "\n",
    "# Add some noise with fixed seed for reproducibility\n",
    "rng = np.random.default_rng(1)\n",
    "yerr = np.full_like(x, 0.5)\n",
    "y = rng.normal(y, yerr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Models\n",
    "\n",
    "A lmfit model can be created by calling {class}`lmfit.Model <lmfit.model.Model>` with the model function and the independent variable(s) as arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lmfit\n",
    "\n",
    "model = lmfit.Model(poly1)\n",
    "params = model.make_params(a=1.0, b=2.0)\n",
    "result = model.fit(y, x=x, params=params, weights=1 / yerr)\n",
    "\n",
    "result.plot()\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By passing dictionaries to `make_params`, we can set the initial values of the parameters and also set the bounds for the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lmfit.Model(poly1)\n",
    "params = model.make_params(\n",
    "    a={\"value\": 1.0, \"min\": 0.0},\n",
    "    b={\"value\": 2.0, \"vary\": False},\n",
    ")\n",
    "result = model.fit(y, x=x, params=params, weights=1 / yerr)\n",
    "_ = result.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`result` is a {class}`lmfit.model.ModelResult` object that contains the results of the\n",
    "fit. The best-fit parameters can be accessed through the `result.params` attribute.\n",
    "\n",
    "\n",
    ":::{note}\n",
    "\n",
    "Since all weights are the same in this case, it has little effect on the fit results. However, if we are confident that we have a good estimate of `yerr`, we can pass `scale_covar=True` to the `fit` method to obtain accurate uncertainties.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.params[\"a\"].value, result.params[\"a\"].stderr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameters can also be retrieved in a form that allows easy error propagation calculation, enabled by the [uncertainties](https://github.com/lmfit/uncertainties) package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_uvar = result.uvars[\"a\"]\n",
    "print(a_uvar)\n",
    "print(a_uvar**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composite models\n",
    "\n",
    "Before fitting, let us generate a Gaussian peak on a linear background."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate toy data\n",
    "x = np.linspace(0, 10, 50)\n",
    "y = -0.1 * x + 2 + 3 * np.exp(-((x - 5) ** 2) / (2 * 1**2))\n",
    "\n",
    "# Add some noise with fixed seed for reproducibility\n",
    "rng = np.random.default_rng(5)\n",
    "yerr = np.full_like(x, 0.3)\n",
    "y = rng.normal(y, yerr)\n",
    "\n",
    "# Plot the data\n",
    "plt.errorbar(x, y, yerr, fmt=\"o\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A composite model can be created by adding multiple models together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lmfit.models import GaussianModel, LinearModel\n",
    "\n",
    "model = GaussianModel() + LinearModel()\n",
    "params = model.make_params(slope=-0.1, center=5.0, sigma={\"value\": 0.1, \"min\": 0})\n",
    "params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.fit(y, x=x, params=params, weights=1 / yerr)\n",
    "result.plot()\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How about multiple gaussian peaks? Since the parameter names overlap between the models, we must use the `prefix` argument to distinguish between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GaussianModel(prefix=\"p0_\") + GaussianModel(prefix=\"p1_\") + LinearModel()\n",
    "model.make_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more information, see the [lmfit documentation](https://lmfit.github.io/lmfit-py/model.html)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}