{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstrate standard major axis (SMA) fitting #\n",
    "\n",
    "We will compare how two statistical fitting procedures compare to each other in the presence of noise and outliers.\n",
    "The two methods are Standard Major Axis (SMA) and Ordinary Least Squares (OLS) regression\n",
    "\n",
    "Steps\n",
    "1. Generate some synthetic data from a known model\n",
    "2. Add noise\n",
    "3. Fit the noisy data with two methods\n",
    "4. Print and visualize the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "## User-defined variables ##\n",
    "\n",
    "# Specify true slope and intercept to use\n",
    "slope_true      = 2\n",
    "intercept_true  = 1\n",
    "\n",
    "# standard deviation of noise\n",
    "# low noise ~ 0.1\n",
    "# high noise > ~1 \n",
    "noise_std = 1\n",
    "\n",
    "# Set whether to add outliers or not\n",
    "add_outliers = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.formula.api as smf\n",
    "import acgc.figstyle\n",
    "from acgc.stats import sma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create 1000 points with specified slope and intercept\n",
    "# x is normally distributed\n",
    "n = 1000\n",
    "xtrue = 3 + np.random.randn(n) \n",
    "ytrue = xtrue * slope_true + intercept_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Noisy versions of x and y\n",
    "\n",
    "# Add noise to x and y\n",
    "xnoisy = xtrue + np.random.randn(n) * noise_std\n",
    "ynoisy = ytrue + np.random.randn(n) * noise_std\n",
    "\n",
    "# Add outliers if requested\n",
    "if add_outliers:\n",
    "\n",
    "    # Index location of outliers\n",
    "    idx = np.random.randint(0,n-1,4)\n",
    "\n",
    "    # Specify outlier values\n",
    "    xnoisy[idx] = xnoisy[idx] + 100 * noise_std\n",
    "    ynoisy[idx] = ynoisy[idx] + 1000 * noise_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['slope', 'intercept', 'slope_ste', 'intercept_ste', 'slope_interval', 'intercept_interval', 'df_model', 'df_resid', 'params', 'nobs', 'fittedvalues', 'resid'])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# SMA fit, use robust methods to minimize effect of outliers\n",
    "sma_result = sma( xnoisy, ynoisy, \n",
    "                cl=0.95 )\n",
    "\n",
    "# Robust SMA\n",
    "rsma_result = sma( xnoisy, ynoisy, \n",
    "                cl=0.95,\n",
    "                robust=True )\n",
    "sma_result.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OLS for comparison\n",
    "ols_result = smf.ols('y ~ x + 1', {'x':xnoisy, 'y':ynoisy} ).fit()\n",
    "\n",
    "# standard error of fit parameters\n",
    "ols_ste_params = np.sqrt( np.diag( ols_result.cov_params() ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Parameters\n",
      "True\n",
      "slope            2.0000+/-  0.0000\n",
      "intercept        1.0000+/-  0.0000\n",
      "\n",
      "OLS\n",
      "slope            1.0610+/-  0.0388\n",
      "intercept        3.8387+/-  0.1284\n",
      "\n",
      "SMA\n",
      "slope            1.6219+/-  0.0388\n",
      "intercept        2.1693+/-  0.1309\n",
      "\n",
      "Robust SMA\n",
      "slope            1.5967+/-  0.0381\n",
      "intercept        2.2644+/-  0.1280\n"
     ]
    }
   ],
   "source": [
    "# Print slope and intercept\n",
    "fmt = '{:15s}{:8.4f}+/-{:8.4f}'\n",
    "\n",
    "print('\\nParameters')\n",
    "print('True')\n",
    "print(fmt.format('slope',       slope_true, 0))\n",
    "print(fmt.format('intercept',   intercept_true, 0))\n",
    "\n",
    "print('\\nOLS')\n",
    "print(fmt.format('slope',       ols_result.params.iloc[1], ols_ste_params[1] ))\n",
    "print(fmt.format('intercept',   ols_result.params.iloc[0], ols_ste_params[0] ))\n",
    "\n",
    "print('\\nSMA')\n",
    "print(fmt.format('slope',       sma_result['slope'],        sma_result['slope_ste']))\n",
    "print(fmt.format('intercept',   sma_result['intercept'],    sma_result['intercept_ste']))\n",
    "\n",
    "print('\\nRobust SMA')\n",
    "print(fmt.format('slope',       rsma_result['slope'],        rsma_result['slope_ste']))\n",
    "print(fmt.format('intercept',   rsma_result['intercept'],    rsma_result['intercept_ste']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x12ae13f70>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 94700 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 400,
       "width": 400
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot data & fit lines\n",
    "\n",
    "# Noisy data\n",
    "plt.scatter( xnoisy, ynoisy, \n",
    "            s=1, color='black', label='Noisy data') \n",
    "# True relationship\n",
    "plt.plot( xtrue, ytrue, \n",
    "          color='black', linewidth=1, label='True')\n",
    "\n",
    "# Fitted relationships\n",
    "plt.plot( xnoisy, sma_result['slope'] * xnoisy + sma_result['intercept'],\n",
    "         color='C1', linewidth=1, label='SMA')\n",
    "plt.plot( xnoisy, ols_result.fittedvalues.values,\n",
    "         color='C0', linewidth=1, label='OLS')\n",
    "if add_outliers:\n",
    "    plt.plot( xnoisy, rsma_result['slope'] * xnoisy + rsma_result['intercept'],\n",
    "            color='C2', linewidth=1, label='rSMA')\n",
    "\n",
    "# Legend\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "general",
   "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.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
