{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.formula.api as smf\n",
    "import statsmodels.api as sm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Invent a dataset\n",
    "nx = 5\n",
    "x  = np.arange(nx)\n",
    "y  = np.array([0.5,0.9,1.7,3.3,2])\n",
    "ye = np.array([0.1,0.3,0.1,0.3,2])\n",
    "ye = np.array([0.6,0.6,0.6,0.6,2])\n",
    "\n",
    "\n",
    "df = pd.DataFrame({'x':x,'y':y,'ye':ye})\n",
    "\n",
    "# Construct design matrix\n",
    "X = sm.add_constant(x)\n",
    "\n",
    "def cov2se(cov):\n",
    "    return np.sqrt( np.diag(cov) )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Intercept    0.60\n",
      "x            0.54\n",
      "dtype: float64\n",
      "Intercept    0.601997\n",
      "x            0.245764\n",
      "dtype: float64\n",
      "[0.60199668 0.24576411]\n",
      "[0.60133186 0.42614552]\n",
      "Intercept    0.29533\n",
      "x            0.84467\n",
      "dtype: float64\n",
      "Intercept    0.385145\n",
      "x            0.198227\n",
      "dtype: float64\n",
      "[0.38514458 0.19822654]\n"
     ]
    }
   ],
   "source": [
    "# OLS\n",
    "fitols = smf.ols('y ~ 1 + x',data=df).fit()\n",
    "# WLS\n",
    "fitwls = smf.wls('y ~ 1 + x',data=df,weights=1/ye**2).fit()\n",
    "\n",
    "# Analytical OLS\n",
    "# The Moore-Penrose pseudo-inverse Xpinv = (X.T @ X )**-1 @ X.T\n",
    "Xpinv = np.linalg.pinv(X)\n",
    "b  = Xpinv @ y\n",
    "# Two equivalent ways of computing the parameter covariances\n",
    "bcov = np.var(fitols.resid,ddof=2) * np.linalg.inv(X.T @ X )\n",
    "bcov = Xpinv @ (np.var(fitols.resid,ddof=2) * np.eye(nx) ) @ Xpinv.T\n",
    "# Instead of using the residuals to extimate the noise covariance, use independent information \n",
    "bcov2= Xpinv @ np.diag(ye**2) @ Xpinv.T\n",
    "# parameter standard error\n",
    "bse = cov2se( bcov )\n",
    "bse2= cov2se( bcov2)\n",
    "\n",
    "\n",
    "# Analytical WLS\n",
    "W = np.diag(1/ye**2) / np.sum(1/ye**2)\n",
    "bwls = np.linalg.inv(X.T @ W @ X )@ X.T @ W @ y \n",
    "\n",
    "# Equivalent\n",
    "Xw = np.sqrt(W) @ X\n",
    "yw = np.sqrt(W) @ y\n",
    "Xwpinv = np.linalg.pinv(Xw)\n",
    "bwls = Xwpinv @ yw\n",
    "\n",
    "bcovwls = np.sum( W @ fitwls.resid**2) / (nx-2) * np.linalg.inv(Xw.T @ Xw)\n",
    "bsewls = cov2se( bcovwls )\n",
    "\n",
    "# Confirm that results are same as OLS\n",
    "print(fitols.params)\n",
    "print(fitols.bse)\n",
    "print(bse)\n",
    "print(bse2)\n",
    "\n",
    "print(fitwls.params)\n",
    "print(fitwls.bse)\n",
    "print(bsewls)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17908 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PARAMETERS\n",
      "Intercept    0.60\n",
      "x            0.54\n",
      "dtype: float64\n",
      "[0.60316503 0.53858571]\n",
      "[0.6034374  0.53658175]\n",
      "\n",
      "UNCERTAINTY\n",
      "\n",
      "OLS\n",
      "Intercept    0.601997\n",
      "x            0.245764\n",
      "dtype: float64\n",
      "\n",
      "OLS + Bootstrap\n",
      "[0.42147252 0.40261013]\n",
      "\n",
      "WLS\n",
      "Intercept    0.174821\n",
      "x            0.114890\n",
      "dtype: float64\n"
     ]
    }
   ],
   "source": [
    "# Bootstrap\n",
    "n=30000\n",
    "coef = np.zeros((2,n))\n",
    "for i in range(n):\n",
    "    # Add gaussian noise\n",
    "    y2 = y + ye * np.random.randn(5)\n",
    "    # fit the new data\n",
    "    fit = sm.OLS(y2,X).fit()\n",
    "    coef[:,i] = fit.params\n",
    "    \n",
    "\n",
    "plt.clf()\n",
    "plt.scatter(x,y,label='data')\n",
    "plt.plot(x,fitols.fittedvalues,label='OLS')\n",
    "plt.plot(x,fitwls.fittedvalues,label='WLS')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "## \n",
    "print('PARAMETERS')\n",
    "print(fitols.params)\n",
    "print(np.median(coef,axis=1))\n",
    "print(np.mean(coef,axis=1))\n",
    "\n",
    "print('\\nUNCERTAINTY')\n",
    "print('\\nOLS')\n",
    "print(fitols.bse)\n",
    "print('\\nOLS + Bootstrap')\n",
    "print(np.std(coef,axis=1))\n",
    "print('\\nWLS')\n",
    "print(fitwls.bse)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
