{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit 3rd order Polynomial to graph data using numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "plt.rcParams[\"font.size\"] = 20\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype='d')\n",
    "dx =  np.array([0.1,0.1,0.5,0.1,0.5,0.1,0.5,0.1,0.5,0.1], dtype='d')\n",
    "y = np.array([1.1 ,2.3 ,2.7 ,3.2 ,3.1 ,2.4 ,1.7 ,1.5 ,1.5  ,1.7 ], dtype='d')\n",
    "dy = np.array([0.15,0.22,0.29,0.39,0.31,0.21,0.13,0.15,0.19,0.13], dtype='d')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create numpy array with weights as 1/error "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ones = np.ones(10, dtype='d')\n",
    "weight = ones/dy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "various fit options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 3 \n",
    "#model = np.polyfit(x, y, n, 0 , 0, weight, cov='unscaled' )\n",
    "#model, fitCov = np.polyfit(x, y, n , 0 , 0, weight, cov='unscaled' )\n",
    "model, fitCov = np.polyfit(x, y, n , 0 , 0, weight, cov='unscaled' )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "monitor printout of coefficient polynomial vector (model)  and covariance matrix (fitCov)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.02590585 -0.48405665  2.52045767 -0.98073638]\n",
      "[[ 1.19801683e-05 -2.07833174e-04  9.95675425e-04 -9.97263377e-04]\n",
      " [-2.07833174e-04  3.65835615e-03 -1.78482005e-02  1.83096634e-02]\n",
      " [ 9.95675425e-04 -1.78482005e-02  8.93452262e-02 -9.59887600e-02]\n",
      " [-9.97263377e-04  1.83096634e-02 -9.59887600e-02  1.19001618e-01]]\n"
     ]
    }
   ],
   "source": [
    "print (model)\n",
    "print (fitCov)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "prepare errors for plotting, see\n",
    "https://stackoverflow.com/questions/28505008/numpy-polyfit-how-to-get-1-sigma-uncertainty-around-the-estimated-curve/28528966\n",
    "write polynomial as np.dot(yy, model)  with yy=[x**n, x*n-1, ..., 1] and  x can be a single value or \n",
    "or a vector and  model as coefficient vector. Since this a linear equation, with the covariance matrix\n",
    "fitCov of model, the covariance matrix of the values is np.dot(yy, np.dot(model, yy.T))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the interpolation for plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_plot = np.linspace(0.1, 10.5, 500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matrix with rows 1, x, x**2, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "YY = np.vstack([x_plot**(n-i) for i in range(n+1)]).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "matrix multiplication calculates the polynomial values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_plot = np.dot(YY, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Covariance_y = YY*Covariance*YY.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitCov_y_plot = np.dot(YY, np.dot(fitCov, YY.T))\n",
    "Sigma_y_plot = np.sqrt(np.diag(fitCov_y_plot))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "model contains parameters in order highest power first! Define fit function "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitpol3(model):\n",
    "    return model[3] + model[2]*x + model[1]*x**2 + model[0]*x**3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "construct polynomial from coefficients in model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict = np.poly1d(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "r_2 coefficient of the determination of the regression score function\n",
    "input are the y data values and the prediction from the fit with x data values \n",
    "this is borrowed from scikit-learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9486643210376879\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import r2_score\n",
    "p = r2_score(y, predict(x))\n",
    "print (p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "plot data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36232 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.errorbar(x, y, dy , dx, fmt=\"o\")\n",
    "plt.plot(x, fitpol3(model) )\n",
    "plt.title(\"numpy Fit Test\")\n",
    "plt.xlim(-0.1, 10.1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "plot data with covariance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45080 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fg, ax = plt.subplots(1, 1)\n",
    "ax.set_title(\"numpy Fit Test with $\\pm1\\sigma$ interval\")\n",
    "ax.fill_between(x_plot, y_plot+Sigma_y_plot, y_plot-Sigma_y_plot, alpha=.25)\n",
    "ax.plot(x_plot, y_plot,'-')\n",
    "ax.errorbar(x, y,dy , dx, fmt='o')\n",
    "ax.axis('tight')\n",
    "\n",
    "fg.canvas.draw()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
