{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear regression diagnostics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In real-life, relation between response and target variables are seldom linear. Here, we make use of outputs of ``statsmodels`` to visualise and identify potential problems that can occur from fitting ``linear regression`` model to non-linear relation. Primarily, the aim is to reproduce visualisations discussed in Potential Problems section (Chapter 3.3.3) of *An Introduction to Statistical Learning* (ISLR) book by James et al., Springer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import statsmodels\n",
    "import statsmodels.formula.api as smf\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Simple multiple linear regression\n",
    "\n",
    "Firstly, let us load the Advertising data from Chapter 2 of ISLR book and fit a linear model to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data\n",
    "data_url = \"https://raw.githubusercontent.com/nguyen-toan/ISLR/07fd968ea484b5f6febc7b392a28eb64329a4945/dataset/Advertising.csv\"\n",
    "df = pd.read_csv(data_url).drop('Unnamed: 0', axis=1)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fitting linear model\n",
    "res = smf.ols(formula= \"Sales ~ TV + Radio + Newspaper\", data=df).fit()\n",
    "res.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Diagnostic Figures/Table\n",
    "\n",
    "In the following first we present a base code that we will later use to generate following diagnostic plots:\n",
    "\n",
    "    a. residual\n",
    "    b. qq\n",
    "    c. scale location\n",
    "    d. leverage\n",
    "\n",
    "and a table\n",
    "\n",
    "    a. vif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# base code\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "from statsmodels.tools.tools import maybe_unwrap_results\n",
    "from statsmodels.graphics.gofplots import ProbPlot\n",
    "from statsmodels.stats.outliers_influence import variance_inflation_factor\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import Type\n",
    "\n",
    "style_talk = 'seaborn-talk'    #refer to plt.style.available\n",
    "\n",
    "class LinearRegDiagnostic():\n",
    "    \"\"\"\n",
    "    Diagnostic plots to identify potential problems in a linear regression fit.\n",
    "    Mainly,\n",
    "        a. non-linearity of data\n",
    "        b. Correlation of error terms\n",
    "        c. non-constant variance\n",
    "        d. outliers\n",
    "        e. high-leverage points\n",
    "        f. collinearity\n",
    "\n",
    "    Authors:\n",
    "        Prajwal Kafle (p33ajkafle@gmail.com, where 3 = r)\n",
    "        Does not come with any sort of warranty.\n",
    "        Please test the code one your end before using.\n",
    "        \n",
    "        Matt Spinelli (m3spinelli@gmail.com, where 3 = r)\n",
    "        (1) Fixed incorrect annotation of the top most extreme residuals in\n",
    "            the Residuals vs Fitted and, especially, the Normal Q-Q plots.\n",
    "        (2) Changed Residuals vs Leverage plot to match closer the y-axis\n",
    "            range shown in the equivalent plot in the R package ggfortify.\n",
    "        (3) Added horizontal line at y=0 in Residuals vs Leverage plot to\n",
    "            match the plots in R package ggfortify and base R.\n",
    "        (4) Added option for placing a vertical guideline on the Residuals\n",
    "            vs Leverage plot using the rule of thumb of h = 2p/n to denote\n",
    "            high leverage (high_leverage_threshold=True).\n",
    "        (5) Added two more ways to compute the Cook's Distance (D) threshold:\n",
    "            * 'baseR': D > 1 and D > 0.5 (default)        \n",
    "            * 'convention': D > 4/n\n",
    "            * 'dof': D > 4 / (n - k - 1)\n",
    "        (6) Fixed class name to conform to Pascal casing convention\n",
    "        (7) Fixed Residuals vs Leverage legend to work with loc='best'\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, \n",
    "                 results: Type[statsmodels.regression.linear_model.RegressionResultsWrapper]) -> None:\n",
    "        \"\"\"\n",
    "        For a linear regression model, generates following diagnostic plots:\n",
    "\n",
    "        a. residual\n",
    "        b. qq\n",
    "        c. scale location and\n",
    "        d. leverage\n",
    "\n",
    "        and a table\n",
    "\n",
    "        e. vif\n",
    "\n",
    "        Args:\n",
    "            results (Type[statsmodels.regression.linear_model.RegressionResultsWrapper]):\n",
    "                must be instance of statsmodels.regression.linear_model object\n",
    "\n",
    "        Raises:\n",
    "            TypeError: if instance does not belong to above object\n",
    "\n",
    "        Example:\n",
    "        >>> import numpy as np\n",
    "        >>> import pandas as pd\n",
    "        >>> import statsmodels.formula.api as smf\n",
    "        >>> x = np.linspace(-np.pi, np.pi, 100)\n",
    "        >>> y = 3*x + 8 + np.random.normal(0,1, 100)\n",
    "        >>> df = pd.DataFrame({'x':x, 'y':y})\n",
    "        >>> res = smf.ols(formula= \"y ~ x\", data=df).fit()\n",
    "        >>> cls = Linear_Reg_Diagnostic(res)\n",
    "        >>> cls(plot_context=\"seaborn-v0_8-paper\")\n",
    "\n",
    "        In case you do not need all plots you can also independently make an individual plot/table\n",
    "        in following ways\n",
    "\n",
    "        >>> cls = Linear_Reg_Diagnostic(res)\n",
    "        >>> cls.residual_plot()\n",
    "        >>> cls.qq_plot()\n",
    "        >>> cls.scale_location_plot()\n",
    "        >>> cls.leverage_plot()\n",
    "        >>> cls.vif_table()\n",
    "        \"\"\"\n",
    "\n",
    "        if isinstance(results, statsmodels.regression.linear_model.RegressionResultsWrapper) is False:\n",
    "            raise TypeError(\"result must be instance of statsmodels.regression.linear_model.RegressionResultsWrapper object\")\n",
    "\n",
    "        self.results = maybe_unwrap_results(results)\n",
    "\n",
    "        self.y_true = self.results.model.endog\n",
    "        self.y_predict = self.results.fittedvalues\n",
    "        self.xvar = self.results.model.exog\n",
    "        self.xvar_names = self.results.model.exog_names\n",
    "\n",
    "        self.residual = np.array(self.results.resid)\n",
    "        influence = self.results.get_influence()\n",
    "        self.residual_norm = influence.resid_studentized_internal\n",
    "        self.leverage = influence.hat_matrix_diag\n",
    "        self.cooks_distance = influence.cooks_distance[0]\n",
    "        self.nparams = len(self.results.params)\n",
    "        self.nresids = len(self.residual_norm)\n",
    "\n",
    "    def __call__(self, plot_context='seaborn-v0_8-paper', **kwargs):\n",
    "        # print(plt.style.available)\n",
    "        with plt.style.context(plot_context):\n",
    "            fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(10,10))\n",
    "            self.residual_plot(ax=ax[0,0])\n",
    "            self.qq_plot(ax=ax[0,1])\n",
    "            self.scale_location_plot(ax=ax[1,0])\n",
    "            self.leverage_plot(\n",
    "                ax=ax[1,1],\n",
    "                high_leverage_threshold = kwargs.get('high_leverage_threshold'),\n",
    "                cooks_threshold = kwargs.get('cooks_threshold'))\n",
    "            plt.show()\n",
    "\n",
    "        return self.vif_table(), fig, ax,\n",
    "\n",
    "    def residual_plot(self, ax=None):\n",
    "        \"\"\"\n",
    "        Residual vs Fitted Plot\n",
    "\n",
    "        Graphical tool to identify non-linearity.\n",
    "        (Roughly) Horizontal red line is an indicator that the residual has a linear pattern\n",
    "        \"\"\"\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots()\n",
    "\n",
    "        sns.residplot(\n",
    "            x=self.y_predict,\n",
    "            y=self.residual,\n",
    "            lowess=True,\n",
    "            scatter_kws={'alpha': 0.5},\n",
    "            line_kws={'color': 'red', 'lw': 1, 'alpha': 0.8},\n",
    "            ax=ax)\n",
    "\n",
    "        # annotations\n",
    "        residual_abs = np.abs(self.residual)\n",
    "        abs_resid = np.flip(np.argsort(residual_abs), 0)\n",
    "        abs_resid_top_3 = abs_resid[:3]\n",
    "        for i in abs_resid_top_3:\n",
    "            ax.annotate(\n",
    "                i,\n",
    "                xy=(self.y_predict[i], self.residual[i]),\n",
    "                color='C3')\n",
    "\n",
    "        ax.set_title('Residuals vs Fitted', fontweight=\"bold\")\n",
    "        ax.set_xlabel('Fitted values')\n",
    "        ax.set_ylabel('Residuals')\n",
    "        return ax\n",
    "\n",
    "    def qq_plot(self, ax=None):\n",
    "        \"\"\"\n",
    "        Standarized Residual vs Theoretical Quantile plot\n",
    "\n",
    "        Used to visually check if residuals are normally distributed.\n",
    "        Points spread along the diagonal line will suggest so.\n",
    "        \"\"\"\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots()\n",
    "            \n",
    "        QQ = ProbPlot(self.residual_norm)\n",
    "        fig = QQ.qqplot(line='45', alpha=0.5, lw=1, ax=ax)\n",
    "\n",
    "        # annotations       \n",
    "        abs_norm_resid = np.flip(np.argsort(np.abs(self.residual_norm)), 0)\n",
    "        abs_norm_resid_top_3 = abs_norm_resid[:3]\n",
    "        for i, x, y in self.__qq_top_resid(QQ.theoretical_quantiles, abs_norm_resid_top_3):\n",
    "            ax.annotate(\n",
    "                i, \n",
    "                xy=(x, y), \n",
    "                ha='right', \n",
    "                color='C3')\n",
    "\n",
    "        ax.set_title('Normal Q-Q', fontweight=\"bold\")\n",
    "        ax.set_xlabel('Theoretical Quantiles')\n",
    "        ax.set_ylabel('Standardized Residuals')\n",
    "        return ax\n",
    "\n",
    "    def scale_location_plot(self, ax=None):\n",
    "        \"\"\"\n",
    "        Sqrt(Standarized Residual) vs Fitted values plot\n",
    "\n",
    "        Used to check homoscedasticity of the residuals.\n",
    "        Horizontal line will suggest so.\n",
    "        \"\"\"\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots()\n",
    "\n",
    "        residual_norm_abs_sqrt = np.sqrt(np.abs(self.residual_norm))\n",
    "\n",
    "        ax.scatter(self.y_predict, residual_norm_abs_sqrt, alpha=0.5);\n",
    "        sns.regplot(\n",
    "            x=self.y_predict,\n",
    "            y=residual_norm_abs_sqrt,\n",
    "            scatter=False, ci=False,\n",
    "            lowess=True,\n",
    "            line_kws={'color': 'red', 'lw': 1, 'alpha': 0.8},\n",
    "            ax=ax)\n",
    "\n",
    "        # annotations\n",
    "        abs_sq_norm_resid = np.flip(np.argsort(residual_norm_abs_sqrt), 0)\n",
    "        abs_sq_norm_resid_top_3 = abs_sq_norm_resid[:3]\n",
    "        for i in abs_sq_norm_resid_top_3:\n",
    "            ax.annotate(\n",
    "                i,\n",
    "                xy=(self.y_predict[i], residual_norm_abs_sqrt[i]),\n",
    "                color='C3')\n",
    "            \n",
    "        ax.set_title('Scale-Location', fontweight=\"bold\")\n",
    "        ax.set_xlabel('Fitted values')\n",
    "        ax.set_ylabel(r'$\\sqrt{|\\mathrm{Standardized\\ Residuals}|}$');\n",
    "        return ax\n",
    "\n",
    "    def leverage_plot(self, ax=None, high_leverage_threshold=False, cooks_threshold='baseR'):\n",
    "        \"\"\"\n",
    "        Residual vs Leverage plot\n",
    "\n",
    "        Points falling outside Cook's distance curves are considered observation that can sway the fit\n",
    "        aka are influential.\n",
    "        Good to have none outside the curves.\n",
    "        \"\"\"\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots()\n",
    "\n",
    "        ax.scatter(\n",
    "            self.leverage,\n",
    "            self.residual_norm,\n",
    "            alpha=0.5);\n",
    "\n",
    "        sns.regplot(\n",
    "            x=self.leverage,\n",
    "            y=self.residual_norm,\n",
    "            scatter=False,\n",
    "            ci=False,\n",
    "            lowess=True,\n",
    "            line_kws={'color': 'red', 'lw': 1, 'alpha': 0.8},\n",
    "            ax=ax)\n",
    "\n",
    "        # annotations\n",
    "        leverage_top_3 = np.flip(np.argsort(self.cooks_distance), 0)[:3]\n",
    "        for i in leverage_top_3:\n",
    "            ax.annotate(\n",
    "                i,\n",
    "                xy=(self.leverage[i], self.residual_norm[i]),\n",
    "                color = 'C3')\n",
    "            \n",
    "        factors = []\n",
    "        if cooks_threshold == 'baseR' or cooks_threshold is None:\n",
    "            factors = [1, 0.5]\n",
    "        elif cooks_threshold == 'convention':\n",
    "            factors = [4/self.nresids]        \n",
    "        elif cooks_threshold == 'dof':\n",
    "            factors = [4/ (self.nresids - self.nparams)]\n",
    "        else:\n",
    "            raise ValueError(\"threshold_method must be one of the following: 'convention', 'dof', or 'baseR' (default)\")\n",
    "        for i, factor in enumerate(factors):\n",
    "            label = \"Cook's distance\" if i == 0 else None\n",
    "            xtemp, ytemp = self.__cooks_dist_line(factor)\n",
    "            ax.plot(xtemp, ytemp, label=label, lw=1.25, ls='--', color='red')\n",
    "            ax.plot(xtemp, np.negative(ytemp), lw=1.25, ls='--', color='red')\n",
    "\n",
    "        if high_leverage_threshold:\n",
    "            high_leverage = 2 * self.nparams / self.nresids\n",
    "            if max(self.leverage) > high_leverage:\n",
    "                ax.axvline(high_leverage, label='High leverage', ls='-.', color='purple', lw=1)\n",
    "                \n",
    "        ax.axhline(0, ls='dotted', color='black', lw=1.25)\n",
    "        ax.set_xlim(0, max(self.leverage)+0.01)\n",
    "        ax.set_ylim(min(self.residual_norm)-0.1, max(self.residual_norm)+0.1)\n",
    "        ax.set_title('Residuals vs Leverage', fontweight=\"bold\")\n",
    "        ax.set_xlabel('Leverage')\n",
    "        ax.set_ylabel('Standardized Residuals')\n",
    "        plt.legend(loc='best')\n",
    "        return ax\n",
    "\n",
    "    def vif_table(self):\n",
    "        \"\"\"\n",
    "        VIF table\n",
    "\n",
    "        VIF, the variance inflation factor, is a measure of multicollinearity.\n",
    "        VIF > 5 for a variable indicates that it is highly collinear with the\n",
    "        other input variables.\n",
    "        \"\"\"\n",
    "        vif_df = pd.DataFrame()\n",
    "        vif_df[\"Features\"] = self.xvar_names\n",
    "        vif_df[\"VIF Factor\"] = [variance_inflation_factor(self.xvar, i) for i in range(self.xvar.shape[1])]\n",
    "\n",
    "        return (vif_df\n",
    "                .sort_values(\"VIF Factor\")\n",
    "                .round(2))\n",
    "\n",
    "\n",
    "    def __cooks_dist_line(self, factor):\n",
    "        \"\"\"\n",
    "        Helper function for plotting Cook's distance curves\n",
    "        \"\"\"\n",
    "        p = self.nparams\n",
    "        formula = lambda x: np.sqrt((factor * p * (1 - x)) / x)\n",
    "        x = np.linspace(0.001, max(self.leverage), 50)\n",
    "        y = formula(x)\n",
    "        return x, y\n",
    "\n",
    "\n",
    "    def __qq_top_resid(self, quantiles, top_residual_indices):\n",
    "        \"\"\"\n",
    "        Helper generator function yielding the index and coordinates\n",
    "        \"\"\"\n",
    "        offset = 0\n",
    "        quant_index = 0\n",
    "        previous_is_negative = None\n",
    "        for resid_index in top_residual_indices:\n",
    "            y = self.residual_norm[resid_index]\n",
    "            is_negative = y < 0\n",
    "            if previous_is_negative == None or previous_is_negative == is_negative:\n",
    "                offset += 1\n",
    "            else:\n",
    "                quant_index -= offset\n",
    "            x = quantiles[quant_index] if is_negative else np.flip(quantiles, 0)[quant_index]\n",
    "            quant_index += 1\n",
    "            previous_is_negative = is_negative\n",
    "            yield resid_index, x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Making use of the \n",
    "\n",
    "    * fitted model on the Advertising data above and \n",
    "    * the base code provided\n",
    "now we generate diagnostic plots one by one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls = LinearRegDiagnostic(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**A. Residual vs Fitted values**\n",
    "\n",
    "Graphical tool to identify non-linearity.\n",
    "\n",
    "In the graph red (roughly) horizontal line is an indicator that the residual has a linear pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls.residual_plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**B. Standarized Residual vs Theoretical Quantile**\n",
    "\n",
    "This plot is used to visually check if residuals are normally distributed.\n",
    "\n",
    "Points spread along the diagonal line will suggest so."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls.qq_plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**C. Sqrt(Standarized Residual) vs Fitted values**\n",
    "\n",
    "This plot is used to check homoscedasticity of the residuals.\n",
    "\n",
    "A near horizontal red line in the graph would suggest so."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls.scale_location_plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**D. Residual vs Leverage**\n",
    "\n",
    "Points falling outside the Cook's distance curves are considered observation that can sway the fit aka are influential.\n",
    "\n",
    "Good to have no points outside these curves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls.leverage_plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cook's distance curves can be drawn using other rules of thumbs:\n",
    "\n",
    "| Rule of Thumb    | Threshold(s) |\n",
    "| :----------- | ----------- |\n",
    "| `'baseR'` (**default**) | $$D_i > 1 \\mid D_i > 0.5$$ |\n",
    "| `'convention'` | $$D_i > { 4 \\over n}$$ |\n",
    "| `'dof'` | $$D_i > {4 \\over n - k - 1}$$ |\n",
    "\n",
    "A high leverage guideline can also be displayed using the convention: $h_{ii} > {2p \\over n}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls.leverage_plot(high_leverage_threshold=True, cooks_threshold='dof');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**E. VIF**\n",
    "\n",
    "The variance inflation factor (VIF), is a measure of multicollinearity.\n",
    "\n",
    "VIF > 5 for a variable indicates that it is highly collinear with the \n",
    "other input variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls.vif_table()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Alternatively, all diagnostics can be generated in one go as follows.\n",
    "# Fig and ax can be used to modify axes or plot properties after the fact.\n",
    "cls = LinearRegDiagnostic(res)\n",
    "vif, fig, ax = cls()\n",
    "print(vif)\n",
    "\n",
    "#fig.savefig('../../docs/source/_static/images/linear_regression_diagnostics_plots.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For detail discussions on the interpretation and caveats of the above plots please refer to the ISLR book."
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
