{
 "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": "markdown",
   "metadata": {},
   "source": [
    "## Generate synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
    "# Number of data points to create\n",
    "n = 1000\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 = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Create points with specified slope and intercept\n",
    "# x is normally distributed\n",
    "xtrue = 3 + np.random.randn(n) \n",
    "ytrue = xtrue * slope_true + intercept_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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\n",
    "\n",
    "    # A different kind of noise\n",
    "    # ynoisy = ynoisy + np.random.standard_cauchy(n)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit the data with multiple methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from acgc.stats import sma, bivariate_line_equation\n",
    "\n",
    "# SMA fit, use robust methods to minimize effect of outliers\n",
    "sma_result = sma( xnoisy, ynoisy, \n",
    "                alpha=0.95 )\n",
    "\n",
    "# Robust SMA\n",
    "rsma_result = sma( xnoisy, ynoisy, \n",
    "                alpha=0.95,\n",
    "                robust=True )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import statsmodels.formula.api as smf\n",
    "\n",
    "# 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": "markdown",
   "metadata": {},
   "source": [
    "## Print parameters of the fitted lines\n",
    "\n",
    "`sma_results` is a dict containing the following keys: ['slope', 'intercept', 'slope_ste', 'intercept_ste', 'slope_interval', 'intercept_interval', 'alpha', 'df_model', 'df_resid', 'params', 'nobs', 'fittedvalues', 'resid']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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            9.5144 ±   0.0606\n",
      "intercept      -21.0556 ±   0.4436\n",
      "\n",
      "SMA\n",
      "slope            9.7049 ±   0.0606\n",
      "intercept      -21.6979 ±   0.4454\n",
      "\n",
      "Robust SMA\n",
      "slope            1.5131 ±   0.0385\n",
      "intercept        2.5436 ±   0.1269\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": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_SMA = (9.7049 ± 0.0606) x + (-21.6979 ± 0.4454)\n",
      "y_rSMA = (1.5131 ± 0.0385) x + (2.5436 ± 0.1269)\n"
     ]
    }
   ],
   "source": [
    "# For SMA, the equation for the fitted line can be written as a string\n",
    "print( bivariate_line_equation(sma_result,  floatformat='{:.4f}', include_error=True) )\n",
    "print( bivariate_line_equation(rsma_result, floatformat='{:.4f}', include_error=True) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph the fitted lines and noisy data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x12101ad80>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 65908 characters]",
      "text/plain": [
       "<Figure size 400x330 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 341,
       "width": 411
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import acgc.figstyle\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.12.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
