{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b1eb5bfa",
   "metadata": {},
   "source": [
    "# Multidimensional models\n",
    "\n",
    "\n",
    "Fitting is not limited to 1D models. The following example demonstrates how to fit a 2D Gaussian peak to a 2D DataArray."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4f666d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import lmfit\n",
    "import xarray_lmfit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5787089a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate synthetic 2D data\n",
    "x = np.linspace(-10, 10, 50)\n",
    "y = np.linspace(-10, 10, 50)\n",
    "x_arr = xr.DataArray(x, dims=(\"x\",), coords={\"x\": x})\n",
    "y_arr = xr.DataArray(y, dims=(\"y\",), coords={\"y\": y})\n",
    "z_arr = lmfit.lineshapes.gaussian2d(\n",
    "    x_arr,\n",
    "    y_arr,\n",
    "    amplitude=4.0,\n",
    "    centerx=0.0,\n",
    "    centery=0.0,\n",
    "    sigmax=1.0,\n",
    "    sigmay=2.0,\n",
    ").rename(\"z\")\n",
    "\n",
    "# Add some noise with fixed seed for reproducibility\n",
    "rng = np.random.default_rng(5)\n",
    "z_arr = z_arr.copy(data=rng.normal(z_arr, 0.01))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "008ca0ee",
   "metadata": {},
   "source": [
    "Fitting a 2D model is as simple as providing multiple coordinate names for different independent variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e48c2d2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lmfit.models.Gaussian2dModel()\n",
    "\n",
    "result_ds = z_arr.xlm.modelfit(\n",
    "    (\"x\", \"y\"),\n",
    "    model=model,\n",
    "    params=model.make_params(\n",
    "        amplitude=2.0, centerx=0.0, centery=0.0, sigmax=1.0, sigmay=2.0\n",
    "    ),\n",
    ")\n",
    "result_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40bf7d77",
   "metadata": {},
   "source": [
    "Let's take a look at the best fit and residuals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75bbfc1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(12, 3), layout=\"compressed\")\n",
    "\n",
    "z_arr.plot(ax=axs[0], center=False)\n",
    "axs[0].set_title(\"Data\")\n",
    "\n",
    "result_ds.modelfit_best_fit.plot(ax=axs[1])\n",
    "axs[1].set_title(\"Fit\")\n",
    "\n",
    "(z_arr - result_ds.modelfit_best_fit).plot(ax=axs[2])\n",
    "axs[2].set_title(\"Data $-$ Fit\")\n",
    "\n",
    "for ax in axs:\n",
    "    ax.set_aspect(\"equal\")"
   ]
  }
 ],
 "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": 5
}