{
 "cells": [
  {
   "cell_type": "raw",
   "id": "33c91c6520a84dfb",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Polynomial Fit Function Array Tutorial\n",
    "======================================\n",
    "\n",
    "A demonstration of creating a function array, fitting a polynomial to it, and then comparing\n",
    "the fit if a different degree polynomial is used.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "887eedc318fabcad",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import named_arrays as na\n",
    "import astropy.units as u\n",
    "import astropy.visualization\n",
    "\n",
    "astropy.visualization.quantity_support();"
   ]
  },
  {
   "cell_type": "raw",
   "id": "971456c65a8ec009",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Instances of :class:`named_arrays.PolynomialFitFunctionArray` are defined similarly to :class:`named_arrays.FunctionArray` by a set of `inputs` (indepent variables) and `outputs` (dependent variables), but with additional parameters `degree`, `axis_polynomial`, and `components_polynomial` used to fit a polynomial function to the function.\n",
    "\n",
    "Start by defining a three-dimensional input grid dependent on wavelength, position.x, and position.y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1393884e560bc1b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = na.SpectralPositionalVectorLinearSpace(\n",
    "    start=na.SpectralPositionalVectorArray(\n",
    "        wavelength=100 * u.AA,\n",
    "        position=na.Cartesian2dVectorArray(\n",
    "            x=-50 * u.arcsec,\n",
    "            y=-50 * u.arcsec,\n",
    "        )\n",
    "    ),\n",
    "    stop=na.SpectralPositionalVectorArray(\n",
    "        wavelength=500 * u.AA,\n",
    "        position=na.Cartesian2dVectorArray(\n",
    "            x=50 * u.arcsec,\n",
    "            y=50 * u.arcsec,\n",
    "        )\n",
    "    ),\n",
    "    num=na.SpectralPositionalVectorArray(\n",
    "        wavelength=2,\n",
    "        position=na.Cartesian2dVectorArray(\n",
    "            x=11,\n",
    "            y=11,\n",
    "        ),\n",
    "    ),\n",
    "    axis=na.SpectralPositionalVectorArray(\n",
    "        wavelength='wavelength',\n",
    "        position=na.Cartesian2dVectorArray(\n",
    "            x='x',\n",
    "            y='y',\n",
    "        ),\n",
    "    ),\n",
    ")\n",
    "\n",
    "inputs"
   ]
  },
  {
   "cell_type": "raw",
   "id": "a8183a186530ddb0",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Using the defined inputs, we calculate a second degree polynomial in :code:`outputs.x` and :code:`outputs.y` with 6 total non-zero coefficients.  Each :code:`x` and :code:`y` has a constant offset, one linear term and one quadratic term.  Coefficients `c` and `d` depend on both time and wavlength."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3734706d26357898",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "t = na.ScalarLinearSpace(-10 * u.s, 10 * u.s, num=3, axis='time')\n",
    "\n",
    "a = 1 * (u.mm / u.arcsec)\n",
    "b = .2 * u.mm / (u.arcsec ** 2)\n",
    "c = t * (u.mm / (u.arcsec * u.s))\n",
    "d = .001 *  inputs.wavelength * u.mm / (u.AA * u.arcsec ** 2)\n",
    "\n",
    "outputs = na.Cartesian2dVectorArray(\n",
    "    x=1 * u.mm + a * inputs.position.x + b * inputs.position.y ** 2 ,\n",
    "    y=5 * u.mm + c * inputs.position.y + d * inputs.position.x ** 2,\n",
    ")"
   ]
  },
  {
   "cell_type": "raw",
   "id": "ccc6d512e02a30a",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Defining a :class:`named_arrays.PolynomialFitFunctionArray` of `degree` 2 with components, :code:`position.x` and :code:`position.y`, and two logical axes, :code:`\"x\"` and :code:`\"y\"`, will fit a second degree polynomial to the spatial part of `inputs` and `outputs`. The attribute :code:`fit.coefficients` gives the coefficients of the linear least squares fit to the inputs and outputs. The fit coefficients match those used to create the outputs to machine precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b736fe5de8bf2d7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = na.PolynomialFitFunctionArray(\n",
    "    inputs=inputs,\n",
    "    outputs=outputs,\n",
    "    degree=2,\n",
    "    components_polynomial=('position.x', 'position.y'),\n",
    "    axis_polynomial=('x', 'y'),\n",
    ")\n",
    "\n",
    "coefficients = fit.coefficients\n",
    "print(coefficients.components['position.x'].x-a)\n",
    "print(coefficients.components['position.y*position.y'].x-b)\n",
    "print(coefficients.components['position.y'].y-c)\n",
    "print(coefficients.components['position.x*position.x'].y-d)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "9f74444af3cd0c4f",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Calling a polynomal function array with a set of inputs uses the fit polynomial to calculate a new set of outputs and returns a :class:`named_arrays.FunctionArray` with those specfied inputs and calculated outputs.  The rms error between the fit and original outputs is low."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2b44484163c4ad",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "best_fit_quad = fit(fit.inputs)\n",
    "rms_error = np.sqrt(np.square(best_fit_quad.outputs - outputs).sum())\n",
    "rms_error"
   ]
  },
  {
   "cell_type": "raw",
   "id": "6fc5a84f1e009567",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "We can also fit a linear function to the original data, and resulting in a much greater error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d8fe6270948c56b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_linear = na.PolynomialFitFunctionArray(\n",
    "    inputs=inputs,\n",
    "    outputs=outputs,\n",
    "    degree=1,\n",
    "    components_polynomial=('position.x', 'position.y'),\n",
    "    axis_polynomial=('x', 'y'),\n",
    ")\n",
    "best_fit_linear = fit_linear(fit.inputs)\n",
    "rms_error = np.sqrt(np.square(best_fit_linear.outputs - outputs).sum())\n",
    "rms_error"
   ]
  },
  {
   "cell_type": "raw",
   "id": "b2796f50bbf1d7a4",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ":mod:`named_arrays` plotting routines make it easy to visualize the orginal outputs, and fit outputs, on their inputs grids as a function of wavlength and time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd63c9428020f63",
   "metadata": {},
   "outputs": [],
   "source": [
    "original_output_y = fit.outputs.y.value\n",
    "quadratic_fit_output_y = best_fit_quad.outputs.y.value\n",
    "linear_fit_output_y = best_fit_linear.outputs.y.value\n",
    "\n",
    "fig, ax = na.plt.subplots(\n",
    "    axis_cols='wavelength',\n",
    "    ncols=fit.shape['wavelength'],\n",
    "    axis_rows='time',\n",
    "    nrows=fit.shape['time'],\n",
    "    sharex=True,\n",
    "    sharey=True\n",
    ")\n",
    "na.plt.pcolormesh(\n",
    "    fit.broadcasted.inputs.position,\n",
    "    C=original_output_y,\n",
    "    ax=ax,\n",
    ")\n",
    "fig.suptitle('Orginal Function Array');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fb0817f55f4679c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = na.plt.subplots(\n",
    "    axis_cols='wavelength',\n",
    "    ncols=fit.shape['wavelength'],\n",
    "    axis_rows='time',\n",
    "    nrows=fit.shape['time'],\n",
    "      sharex=True,\n",
    "    sharey=True\n",
    ")\n",
    "na.plt.pcolormesh(\n",
    "    fit.broadcasted.inputs.position,\n",
    "    C=quadratic_fit_output_y,\n",
    "    ax=ax,\n",
    ")\n",
    "fig.suptitle('Quadratic Fit');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c2f891801be4bd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = na.plt.subplots(\n",
    "    axis_cols='wavelength',\n",
    "    ncols=fit.shape['wavelength'],\n",
    "    axis_rows='time',\n",
    "    nrows=fit.shape['time'],\n",
    "    sharex=True,\n",
    "    sharey=True\n",
    ")\n",
    "na.plt.pcolormesh(\n",
    "    fit.broadcasted.inputs.position,\n",
    "    C=linear_fit_output_y ,\n",
    "    ax=ax,\n",
    ")\n",
    "fig.suptitle('Linear Fit');"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}