{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting censored data\n",
    "\n",
    "Experimental measurements are sometimes censored such that we only know partial information about a particular data point. For example, in measuring the lifespan of mice, a portion of them might live through the duration of the study, in which case we only know the lower bound.\n",
    "\n",
    "One of the ways we can deal with this is to use Maximum Likelihood Estimation ([MLE](http://en.wikipedia.org/wiki/Maximum_likelihood)). However, censoring often make analytical solutions difficult even for well known distributions.\n",
    "\n",
    "We can overcome this challenge by converting the MLE into a convex optimization problem and solving it using [CVXPY](http://www.cvxpy.org/en/latest/).\n",
    "\n",
    "This example is adapted from a homework problem from Boyd's [CVX 101: Convex Optimization Course](https://class.stanford.edu/courses/Engineering/CVX101/Winter2014/info).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We will use similar notation here. Suppose we have a linear model:\n",
    "\n",
    "$$ y^{(i)} = c^Tx^{(i)} +\\epsilon^{(i)} $$\n",
    "\n",
    "where $y^{(i)} \\in \\mathbf{R}$, $c \\in \\mathbf{R}^n$, $k^{(i)} \\in \\mathbf{R}^n$, and $\\epsilon^{(i)}$ is the error and has a normal distribution $N(0, \\sigma^2)$ for $ i = 1,\\ldots,K$.\n",
    "\n",
    "Then the MLE estimator $c$ is the vector that minimizes the sum of squares of the errors $\\epsilon^{(i)}$, namely:\n",
    "\n",
    "$$\n",
    "\\begin{array}{ll}\n",
    "  \\underset{c}{\\mbox{minimize}} & \\sum_{i=1}^K (y^{(i)} - c^T x^{(i)})^2\n",
    "\\end{array}\n",
    "$$\n",
    "\n",
    "In the case of right censored data, only $M$ observations are fully observed and all that is known for the remaining observations is that $y^{(i)} \\geq D$ for $i=\\mbox{M+1},\\ldots,K$ and some constant $D$.\n",
    "\n",
    "Now let's see how this would work in practice.\n",
    "\n",
    "\n",
    "## Data Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "n = 30 # number of variables\n",
    "M = 50 # number of censored observations\n",
    "K = 200 # total number of observations\n",
    "\n",
    "np.random.seed(n*M*K)\n",
    "X = np.random.randn(K*n).reshape(K, n)\n",
    "c_true = np.random.rand(n)\n",
    "\n",
    "# generating the y variable\n",
    "y = X.dot(c_true) + .3*np.sqrt(n)*np.random.randn(K)\n",
    "\n",
    "# ordering them based on y\n",
    "order = np.argsort(y)\n",
    "y_ordered = y[order]\n",
    "X_ordered = X[order,:]\n",
    "\n",
    "#finding boundary\n",
    "D = (y_ordered[M-1] + y_ordered[M])/2. \n",
    "\n",
    "# applying censoring\n",
    "y_censored = np.concatenate((y_ordered[:M], np.ones(K-M)*D))\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Show plot inline in ipython.\n",
    "%matplotlib inline\n",
    "\n",
    "def plot_fit(fit, fit_label):\n",
    "    plt.figure(figsize=(10,6))\n",
    "    plt.grid()\n",
    "    plt.plot(y_censored, 'bo', label = 'censored data')\n",
    "    plt.plot(y_ordered, 'co', label = 'uncensored data')\n",
    "    plt.plot(fit, 'ro', label=fit_label)\n",
    "    plt.ylabel('y')\n",
    "    plt.legend(loc=0)\n",
    "    plt.xlabel('observations');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Regular OLS\n",
    "\n",
    "Let's see what the OLS result looks like. We'll use the `np.linalg.lstsq` function to solve for our coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29240 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "c_ols = np.linalg.lstsq(X_ordered, y_censored, rcond=None)[0]\n",
    "fit_ols = X_ordered.dot(c_ols)\n",
    "plot_fit(fit_ols, 'OLS fit')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that we are systematically overestimating low values of $y$ and vice versa (red vs. cyan). This is caused by our use of censored (blue) observations, which are exerting a lot of leverage and pulling down the trendline to reduce the error between the red and blue points.\n",
    "\n",
    "## OLS using uncensored data\n",
    "\n",
    "A simple way to deal with this while maintaining analytical tractability is to simply ignore all censored observations. \n",
    "\n",
    "$$\n",
    "\\begin{array}{ll}\n",
    "  \\underset{c}{\\mbox{minimize}} & \\sum_{i=1}^M (y^{(i)} - c^T x^{(i)})^2\n",
    "\\end{array}\n",
    "$$\n",
    "\n",
    "Give that our $M$ is much smaller than  $K$, we are  throwing away the majority of the dataset in order to accomplish this, let's see how this new regression does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32752 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "c_ols_uncensored = np.linalg.lstsq(X_ordered[:M], y_censored[:M], rcond=None)[0]\n",
    "fit_ols_uncensored = X_ordered.dot(c_ols_uncensored)\n",
    "plot_fit(fit_ols_uncensored, 'OLS fit with uncensored data only')\n",
    "\n",
    "bad_predictions = (fit_ols_uncensored<=D) & (np.arange(K)>=M)\n",
    "plt.plot(np.arange(K)[bad_predictions], fit_ols_uncensored[bad_predictions], color='orange', marker='o', lw=0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the fit for the uncensored portion is now vastly improved. Even the fit for the censored data is now relatively unbiased i.e. the fitted values (red points) are now centered around the uncensored obsevations (cyan points).\n",
    "\n",
    "The one glaring issue with this arrangement is that we are now predicting many observations to be _below_ $D$ (orange) even though we are well aware that this is not the case. Let's try to fix this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##Using constraints to take into account of censored data\n",
    "Instead of throwing away all censored observations, lets leverage these observations to enforce the additional information that we know, namely that $y$ is bounded from below. We can do this by setting additional constraints:\n",
    "\n",
    "$$\n",
    "\\begin{array}{ll}\n",
    "  \\underset{c}{\\mbox{minimize}} & \\sum_{i=1}^M (y^{(i)} - c^T x^{(i)})^2 \\\\\n",
    "  \\mbox{subject to} & c^T x^{(i)} \\geq D\\\\\n",
    "  & \\mbox{for } i=\\mbox{M+1},\\ldots,K\n",
    "\\end{array}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31876 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cvxpy as cp\n",
    "X_uncensored = X_ordered[:M, :]\n",
    "c = cp.Variable(shape=n)\n",
    "objective = cp.Minimize(cp.sum_squares(X_uncensored*c - y_ordered[:M]))\n",
    "constraints = [ X_ordered[M:,:]*c >= D]\n",
    "prob = cp.Problem(objective, constraints)\n",
    "result = prob.solve()\n",
    "\n",
    "c_cvx = np.array(c.value).flatten()\n",
    "fit_cvx = X_ordered.dot(c_cvx)\n",
    "plot_fit(fit_cvx, 'CVX fit')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Qualitatively, this already looks better than before as it no longer predicts inconsistent values with respect to the censored portion of the data. But does it do a good job of actually finding coefficients $c$ that are close to our original data?\n",
    "\n",
    "We'll use a simple Euclidean distance $\\|c_\\mbox{true} - c\\|_2$ to compare:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "norm(c_true - c_cvx): 1.49\n",
      "norm(c_true - c_ols_uncensored): 2.23\n"
     ]
    }
   ],
   "source": [
    "print(\"norm(c_true - c_cvx): {:.2f}\".format(np.linalg.norm((c_true - c_cvx))))\n",
    "print(\"norm(c_true - c_ols_uncensored): {:.2f}\".format(np.linalg.norm((c_true - c_ols_uncensored))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "Fitting censored data to a parametric distribution can be challenging as the MLE solution is often not analytically tractable. However, many MLEs can be converted into a convex optimization problems as show above. With the advent of simple-to-use and robust numerical packages, we can now solve these problems easily while taking into account the entirety of our information set by enforcing consistency conditions on various portions of the data."
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
