{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import statsmodels.api as sm\n",
    "import matplotlib.pyplot as plt\n",
    "import dautil as dl\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def set_labels(ax):\n",
    "    ax.set_xlabel('Year')\n",
    "    ax.set_ylabel('Sunactivity')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_fit(df, ax, results):\n",
    "    x = df['YEAR']\n",
    "    cp = dl.plotting.CyclePlotter(ax)\n",
    "    cp.plot(x[1:], df['SUNACTIVITY'][1:], label='Data')\n",
    "    cp.plot(x[2:], results.predict()[1:], label='Fit')\n",
    "    ax.legend(loc='best')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "df = sm.datasets.sunspots.load_pandas().data\n",
    "vals = df['SUNACTIVITY'].values\n",
    "\n",
    "# Outlier added by malicious person, because noone\n",
    "# laughs at his jokes.\n",
    "vals[0] = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rlm_model = sm.RLM(vals[1:], sm.add_constant(vals[:-1]),\n",
    "                   M=sm.robust.norms.TrimmedMean())\n",
    "\n",
    "rlm_results = rlm_model.fit()\n",
    "hb = dl.report.HTMLBuilder()\n",
    "hb.h1('Fitting a robust linear model')\n",
    "hb.h2('Robust Linear Model')\n",
    "hb.add(rlm_results.summary().tables[1].as_html())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hb.h2('Ordinary Linear Model')\n",
    "ols_model = sm.OLS(vals[1:], sm.add_constant(vals[:-1]))\n",
    "ols_results = ols_model.fit()\n",
    "hb.add(ols_results.summary().tables[1].as_html())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "dl.options.mimic_seaborn()\n",
    "context = dl.nb.Context('rlm_demo')\n",
    "dl.nb.RcWidget(context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, [ax, ax2] = plt.subplots(2, 1)\n",
    "\n",
    "plot_fit(df, ax, rlm_results)\n",
    "ax.set_title('Robust Linear Model')\n",
    "set_labels(ax)\n",
    "\n",
    "ax2.set_title('Ordinary Least Squares')\n",
    "plot_fit(df, ax2, ols_results)\n",
    "set_labels(ax2)\n",
    "plt.tight_layout()\n",
    "HTML(hb.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
