Realtime R0.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimating COVID-19's $R_t$ in Real-Time\n",
    "Kevin Systrom - April 12\n",
    "\n",
    "In any epidemic, $R_t$ is the measure known as the effective reproduction number. It's the number of people who become infected per infectious person at time $t$. The most well-known version of this number is the basic reproduction number: $R_0$ when $t=0$. However, $R_0$ is a single measure that does not adapt with changes in behavior and restrictions.\n",
    "\n",
    "As a pandemic evolves, increasing restrictions (or potential releasing of restrictions) change $R_t$. Knowing the current $R_t$ is essential. When $R>1$, the pandemic will spread through the entire population. If $R_t<1$, the pandemic will grow to some fixed number less than the population. The lower $R_t$, the more manageable the situation. The value of $R_t$ helps us (1) understand how effective our measures have been controlling an outbreak and (2) gives us vital information about whether we should increase or reduce restrictions based on our competing goals of economic prosperity and human safety. [Well-respected epidemiologists argue](https://www.nytimes.com/2020/04/06/opinion/coronavirus-end-social-distancing.html) that tracking $R_t$ is the only way to manage through this crisis.\n",
    "\n",
    "Yet, today, to my knowledge there is no real-time tracking of $R_t$ in United States. In fact, the only real-time measure I've seen has been for [Hong Kong](https://covid19.sph.hku.hk/dashboard). More importantly, it is not useful to understand $R_t$ at a national level. Instead, to manage this crisis effectively, we need a local (state, county and/or city) level granularity of $R_t$.\n",
    "\n",
    "What follows is a solution to this problem at the US State level. It's a modified version of a solution created by [Bettencourt & Ribeiro 2008](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0002185) to estimate real-time $R_t$ using a Bayesian approach. While I have stayed true to most of their process, my solution differs in an important way that I will call out clearly.\n",
    "\n",
    "If you have questions, comments, or improvments feel free to get in touch: [hello@systrom.com](mailto:hello@systrom.com). And if it's not entirely clear, I'm not an epidemiologist. At the same time, data is data, and statistics are statistics and this is based on work by well-known epidemiologists so calibrate accordingly. In the meantime, I hope you can learn something new as I did by reading through this example. Feel free to take this work and apply it elsewhere – internationally or to counties in the United States."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.dates import date2num, num2date\n",
    "from matplotlib import dates as mdates\n",
    "from matplotlib import ticker\n",
    "from matplotlib.colors import ListedColormap\n",
    "from matplotlib.patches import Patch\n",
    "\n",
    "from scipy import stats as sps\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "FILTERED_REGIONS = [\n",
    "    'Virgin Islands',\n",
    "    'American Samoa',\n",
    "    'Northern Mariana Islands',\n",
    "    'Guam',\n",
    "    'Puerto Rico']\n",
    "\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bettencourt & Ribeiro's Approach\n",
    "\n",
    "Every day, we learn how many more people have COVID-19. This new case count gives us a clue about the current value of $R_t$. We also, figure that the value of $R_t$ today is related to the value of $R_{t-1}$ (yesterday's value) and every previous value of $R_{t-m}$ for that matter.\n",
    "\n",
    "With these insights, the authors use [Bayes' rule](https://en.wikipedia.org/wiki/Bayes%27_theorem) to update their beliefs about the true value of $R_t$ based on how many new cases have been reported each day.\n",
    "\n",
    "This is Bayes' Theorem as we'll use it:\n",
    "\n",
    "$$ P(R_t|k)=\\frac{P(R_t)\\cdot\\mathcal{L}(R_t|k)}{P(k)} $$\n",
    "\n",
    "This says that, having seen $k$ new cases, we believe the distribution of $R_t$ is equal to:\n",
    "\n",
    "- The __prior__ beliefs of the value of $P(R_t)$ without the data ...\n",
    "- times the __likelihood__ of $R_t$ given that we've seen $k$ new cases ...\n",
    "- divided by the probability of seeing this many cases in general.\n",
    "\n",
    "Importantly, $P(k)$ is a constant, so the numerator is proportional to the posterior. Since all probability distributions sum to 1.0, we can ignore $P(k)$ and normalize our posterior to sum to 1.0:\n",
    "\n",
    "$$ P(R_t|k) \\propto P(R_t) \\cdot \\mathcal{L}(R_t|k) $$\n",
    "\n",
    "This is for a single day. To make it iterative: every day that passes, we use yesterday's conclusion (ie. posterior) $P(R_{t-1}|k_{t-1})$ to be today's prior $P(R_t)$ so on day two:\n",
    "\n",
    "$$ P(R_2|k) \\propto P(R_0)\\cdot\\mathcal{L}(R_2|k_2)\\cdot\\mathcal{L}(R_1|k_1) $$\n",
    "\n",
    "And more generally:\n",
    "\n",
    "$$ P(R_t|k_t) \\propto P(R_0) \\cdot {\\displaystyle \\prod^{T}_{t=0}}\\mathcal{L}(R_t|k_t) $$\n",
    "\n",
    "With a uniform prior $P(R_0)$, this reduces to:\n",
    "\n",
    "$$ P(R_t|k_t) \\propto {\\displaystyle \\prod^{T}_{t=0}}\\mathcal{L}\\left(R_t|k_t\\right) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### My Proposed Modification\n",
    "\n",
    "This works fine, but it suffers from an issue: the posterior on any given day is equally influenced by the distant past as much as the recent day. For epidemics that have $R_t>1$ for a long time and then become under control ($R_t<1$), the posterior gets stuck. It cannot forget about the many days where $R_t>1$ so eventually $P(R_t|k)$ asymptotically approaches 1 when we know it's well under 1. The authors note this in the paper as a footnote. Unfortunately this won't work for us. __The most critical thing to know is when we've dipped below the 1.0 threshold!__\n",
    "\n",
    "So, I propose to only incorporate the last $m$ days of the likelihood function. By doing this, the algorithm's prior is built based on the recent past which is a much more useful prior than the entire history of the epidemic. So this simple, but important change leads to the following:\n",
    "\n",
    "$$ P(R_t|k_t) \\propto {\\displaystyle \\prod^{T}_{t=T-m}}\\mathcal{L}\\left(R_t|k_t\\right) $$\n",
    "\n",
    "While this takes the last $m$ priors into account equally, you can decide to apply a windowing function (such as an exponential) to favor recent priors over more distant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choosing a Likelihood Function $\\mathcal{L}\\left(R_t|k_t\\right)$\n",
    "\n",
    "A likelihood function function says how likely a value of $R_t$ is given an observed number of new cases $k$.\n",
    "\n",
    "Any time you need to model 'arrivals' over some time period of time, statisticians like to use the [Poisson Distribution](https://en.wikipedia.org/wiki/Poisson_distribution). Given an average arrival rate of $\\lambda$ new cases per day, the probability of seeing $k$ new cases is distributed according to the Poisson distribution:\n",
    "\n",
    "$$P(k|\\lambda) = \\frac{\\lambda^k e^{-\\lambda}}{k!}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Column vector of k\n",
    "k = np.arange(0, 70)[:, None]\n",
    "\n",
    "# Different values of Lambda\n",
    "lambdas = [10, 20, 30, 40]\n",
    "\n",
    "# Evaluated the Probability Mass Function (remember: poisson is discrete)\n",
    "y = sps.poisson.pmf(k, lambdas)\n",
    "\n",
    "# Show the resulting shape\n",
    "print(y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> __Note__: this was a terse expression which makes it tricky. All I did was to make $k$ a column. By giving it a column for $k$ and a 'row' for lambda it will evaluate the pmf over both and produce an array that has $k$ rows and lambda columns. This is an efficient way of producing many distributions all at once, and __you will see it used again below__!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set(title='Poisson Distribution of Cases\\n $p(k|\\lambda)$')\n",
    "\n",
    "plt.plot(k, y,\n",
    "         marker='o',\n",
    "         markersize=3,\n",
    "         lw=0)\n",
    "\n",
    "plt.legend(title=\"$\\lambda$\", labels=lambdas);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Poisson distribution says that if you think you're going to have $\\lambda$ cases per day, you'll probably get that many, plus or minus some variation based on chance.\n",
    "\n",
    "But in our case, we know there have been $k$ cases and we need to know what value of $\\lambda$ is most likely. In order to do this, we fix $k$ in place while varying $\\lambda$. __This is called the likelihood function.__\n",
    "\n",
    "For example, imagine we observe $k=20$ new cases, and we want to know how likely each $\\lambda$ is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 20\n",
    "\n",
    "lam = np.linspace(1, 45, 90)\n",
    "\n",
    "likelihood = pd.Series(data=sps.poisson.pmf(k, lam),\n",
    "                       index=pd.Index(lam, name='$\\lambda$'),\n",
    "                       name='lambda')\n",
    "\n",
    "likelihood.plot(title=r'Likelihood $L\\left(\\lambda|k_t\\right)$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This says that if we see 20 cases, the most likely value of $\\lambda$ is (not surprisingly) 20. But we're not certain: it's possible lambda was 21 or 17 and saw 20 new cases by chance alone. It also says that it's unlikely $\\lambda$ was 40 and we saw 20.\n",
    "\n",
    "Great. We have $\\mathcal{L}\\left(\\lambda_t|k_t\\right)$ which is parameterized by $\\lambda$ but we were looking for $\\mathcal{L}\\left(R_t|k_t\\right)$ which is parameterized by $R_t$. We need to know the relationship between $\\lambda$ and $R_t$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connecting $\\lambda$ and $R_t$\n",
    "\n",
    "__The key insight to making this work is to realize there's a connection between $R_t$ and $\\lambda$__. [The derivation](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0002185) is beyond the scope of this notebook, but here it is:\n",
    "\n",
    "$$ \\lambda = k_{t-1}e^{\\gamma(R_t-1)}$$\n",
    "\n",
    "where $\\gamma$ is the reciprocal of the serial interval ([about 4 days for COVID19](https://wwwnc.cdc.gov/eid/article/26/6/20-0357_article)). Since we know every new case count on the previous day, we can now reformulate the likelihood function as a Poisson parameterized by fixing $k$ and varying $R_t$.\n",
    "\n",
    "$$ \\lambda = k_{t-1}e^{\\gamma(R_t-1)}$$\n",
    "\n",
    "$$\\mathcal{L}\\left(R_t|k\\right) = \\frac{\\lambda^k e^{-\\lambda}}{k!}$$\n",
    "\n",
    "### Evaluating the Likelihood Function\n",
    "\n",
    "To contiue our example, let's imagine a sample of new case counts $k$. What is the likelihood of different values of $R_t$ on each of those days?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = np.array([20, 40, 55, 90])\n",
    "\n",
    "# We create an array for every possible value of Rt\n",
    "R_T_MAX = 12\n",
    "r_t_range = np.linspace(0, R_T_MAX, R_T_MAX*100+1)\n",
    "\n",
    "# Gamma is 1/serial interval\n",
    "# https://wwwnc.cdc.gov/eid/article/26/6/20-0357_article\n",
    "GAMMA = 1/4\n",
    "\n",
    "# Map Rt into lambda so we can substitute it into the equation below\n",
    "# Note that we have N-1 lambdas because on the first day of an outbreak\n",
    "# you do not know what to expect.\n",
    "lam = k[:-1] * np.exp(GAMMA * (r_t_range[:, None] - 1))\n",
    "\n",
    "# Evaluate the likelihood on each day and normalize sum of each day to 1.0\n",
    "likelihood_r_t = sps.poisson.pmf(k[1:], lam)\n",
    "likelihood_r_t / np.sum(likelihood_r_t, axis=0)\n",
    "\n",
    "# Plot it\n",
    "ax = pd.DataFrame(\n",
    "    data = likelihood_r_t,\n",
    "    index = r_t_range\n",
    ").plot(\n",
    "    title='Likelihood of $R_t$ given $k$',\n",
    "    xlim=(0,7)\n",
    ")\n",
    "\n",
    "ax.legend(labels=k[1:], title='New Cases')\n",
    "ax.set_xlabel('$R_t$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that each day we have a independent guesses for $R_t$. The goal is to combine the information we have about previous days with the current day. To do this, we use Bayes' theorem.\n",
    "\n",
    "### Performing the Bayesian Update\n",
    "\n",
    "To perform the Bayesian update, we need to multiply the likelihood by the prior (which is just the previous day's likelihood) to get the posteriors. Let's do that using the cumulative product of each successive day:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posteriors = likelihood_r_t.cumprod(axis=1)\n",
    "posteriors = posteriors / np.sum(posteriors, axis=0)\n",
    "\n",
    "columns = pd.Index(range(1, posteriors.shape[1]+1), name='Day')\n",
    "posteriors = pd.DataFrame(\n",
    "    data = posteriors,\n",
    "    index = r_t_range,\n",
    "    columns = columns)\n",
    "\n",
    "ax = posteriors.plot(\n",
    "    title='Posterior $P(R_t|k)$',\n",
    "    xlim=(0,7)\n",
    ")\n",
    "ax.legend(title='Day')\n",
    "ax.set_xlabel('$R_t$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how on Day 1, our posterior matches Day 1's likelihood from above? That's because we have no information other than that day. However, when we update the prior using Day 2's information, you can see the curve has moved left, but not nearly as left as the likelihood for Day 2 from above. This is because Bayesian updating uses information from both days and effectively averages the two. Since Day 3's likelihood is in between the other two, you see a small shift to the right, but more importantly: a narrower distribution. We're becoming __more__ confident in our believes of the true value of $R_t$.\n",
    "\n",
    "From these posteriors, we can answer important questions such as \"What is the most likely value of $R_t$ each day?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "most_likely_values = posteriors.idxmax(axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also obtain the [highest density intervals](https://www.sciencedirect.com/topics/mathematics/highest-density-interval) for $R_t$:\n",
    "\n",
    "> Note: I apologize in advance for the clunky brute force HDI algorithm. Please let me know if there are better ones out there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def highest_density_interval(pmf, p=.95):\n",
    "    \n",
    "    # If we pass a DataFrame, just call this recursively on the columns\n",
    "    if(isinstance(pmf, pd.DataFrame)):\n",
    "        return pd.DataFrame([highest_density_interval(pmf[col]) for col in pmf],\n",
    "                            index=pmf.columns)\n",
    "    \n",
    "    cumsum = np.cumsum(pmf.values)\n",
    "    best = None\n",
    "    for i, value in enumerate(cumsum):\n",
    "        for j, high_value in enumerate(cumsum[i+1:]):\n",
    "            if (high_value-value > p) and (not best or j<best[1]-best[0]):\n",
    "                best = (i, i+j+1)\n",
    "                break\n",
    "            \n",
    "    low = pmf.index[best[0]]\n",
    "    high = pmf.index[best[1]]\n",
    "    return pd.Series([low, high], index=['Low', 'High'])\n",
    "\n",
    "hdi = highest_density_interval(posteriors, p=.95)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can plot both the most likely values for $R_t$ and the HDIs over time. This is the most useful representation as it shows how our beliefs change with every day."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = most_likely_values.plot(marker='o',\n",
    "                             label='Most Likely',\n",
    "                             title=f'$R_t$ by day',\n",
    "                             c='k',\n",
    "                             markersize=4)\n",
    "\n",
    "ax.fill_between(hdi.index,\n",
    "                hdi['Low'],\n",
    "                hdi['High'],\n",
    "                color='k',\n",
    "                alpha=.1,\n",
    "                lw=0,\n",
    "                label='HDI')\n",
    "\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the most likely value of $R_t$ changes with time and the highest-density interval narrows as we become more sure of the true value of $R_t$ over time. Note that since we only had four days of history, I did not apply my windowing modification to this sample. Next, however, we'll turn to a real-world application where this modification is necessary."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Real-World Application to US Data\n",
    "\n",
    "### Setup\n",
    "\n",
    "Load US state case data from the NYT archive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'https://raw.githubusercontent.com/nytimes/covid-19-data/master/us-states.csv'\n",
    "states = pd.read_csv(url,\n",
    "                     usecols=[0,1,3],\n",
    "                     index_col=['state', 'date'],\n",
    "                     parse_dates=['date'],\n",
    "                     squeeze=True).sort_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taking a look at the state, we need to start the analysis when there are a consistent number of cases each day. Find the last zero new case day and start on the day after that.\n",
    "\n",
    "Also, case reporting is very erratic based on testing backlogs, etc. To get the best view of the 'true' data we can, I've applied a gaussian filter to the time series. This is obviously an arbitrary choice, but you'd imagine the real world process is not nearly as stochastic as the actual reporting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_name = 'New York'\n",
    "\n",
    "def prepare_cases(cases):\n",
    "    new_cases = cases.diff()\n",
    "\n",
    "    smoothed = new_cases.rolling(7,\n",
    "        win_type='gaussian',\n",
    "        min_periods=1,\n",
    "        center=True).mean(std=2).round()\n",
    "    \n",
    "    zeros = smoothed.index[smoothed.eq(0)]\n",
    "    if len(zeros) == 0:\n",
    "        idx_start = 0\n",
    "    else:\n",
    "        last_zero = zeros.max()\n",
    "        idx_start = smoothed.index.get_loc(last_zero) + 1\n",
    "    smoothed = smoothed.iloc[idx_start:]\n",
    "    original = new_cases.loc[smoothed.index]\n",
    "    \n",
    "    return original, smoothed\n",
    "\n",
    "cases = states.xs(state_name).rename(f\"{state_name} cases\")\n",
    "\n",
    "original, smoothed = prepare_cases(cases)\n",
    "\n",
    "original.plot(title=f\"{state_name} New Cases per Day\",\n",
    "               c='k',\n",
    "               linestyle=':',\n",
    "               alpha=.5,\n",
    "               label='Actual',\n",
    "               legend=True,\n",
    "             figsize=(600/72, 400/72))\n",
    "\n",
    "ax = smoothed.plot(label='Smoothed',\n",
    "                   legend=True)\n",
    "ax.get_figure().set_facecolor('w')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like the example before, we create lambda based on the previous day's counts from all values of $R_t$. Unlike the previous example, I now evaluate the __log__ of the Poisson. Why? It makes windowing easier.\n",
    "\n",
    "Since $\\log{ab}=\\log{a}+\\log{b}$, we can do a rolling sum over the last $m$ periods and then exponentiate to get the rolling product of the original values. This does not change any of the numbers – it's just a convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_posteriors(sr, window=7, min_periods=1):\n",
    "    lam = sr[:-1].values * np.exp(GAMMA * (r_t_range[:, None] - 1))\n",
    "\n",
    "    # Note: if you want to have a Uniform prior you can use the following line instead.\n",
    "    # I chose the gamma distribution because of our prior knowledge of the likely value\n",
    "    # of R_t.\n",
    "    \n",
    "    # prior0 = np.full(len(r_t_range), np.log(1/len(r_t_range)))\n",
    "    prior0 = np.log(sps.gamma(a=3).pdf(r_t_range) + 1e-14)\n",
    "\n",
    "    likelihoods = pd.DataFrame(\n",
    "        # Short-hand way of concatenating the prior and likelihoods\n",
    "        data = np.c_[prior0, sps.poisson.logpmf(sr[1:].values, lam)],\n",
    "        index = r_t_range,\n",
    "        columns = sr.index)\n",
    "\n",
    "    # Perform a rolling sum of log likelihoods. This is the equivalent\n",
    "    # of multiplying the original distributions. Exponentiate to move\n",
    "    # out of log.\n",
    "    posteriors = likelihoods.rolling(window,\n",
    "                                     axis=1,\n",
    "                                     min_periods=min_periods).sum()\n",
    "    posteriors = np.exp(posteriors)\n",
    "\n",
    "    # Normalize to 1.0\n",
    "    posteriors = posteriors.div(posteriors.sum(axis=0), axis=1)\n",
    "    \n",
    "    return posteriors\n",
    "\n",
    "posteriors = get_posteriors(smoothed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Result\n",
    "\n",
    "Below you can see every day (row) of the posterior distribution plotted simultaneously. The posteriors start without much confidence (wide) and become progressively more confident (narrower) about the true value of $R_t$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = posteriors.plot(title=f'{state_name} - Daily Posterior for $R_t$',\n",
    "           legend=False, \n",
    "           lw=1,\n",
    "           c='k',\n",
    "           alpha=.3,\n",
    "           xlim=(0.4,4))\n",
    "\n",
    "ax.set_xlabel('$R_t$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting in the Time Domain with Credible Intervals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since our results include uncertainty, we'd like to be able to view the most likely value of $R_t$ along with its highest-density interval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that this takes a while to execute - it's not the most efficient algorithm\n",
    "hdis = highest_density_interval(posteriors)\n",
    "\n",
    "most_likely = posteriors.idxmax().rename('ML')\n",
    "\n",
    "# Look into why you shift -1\n",
    "result = pd.concat([most_likely, hdis], axis=1)\n",
    "\n",
    "result.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_rt(result, ax, state_name):\n",
    "    \n",
    "    ax.set_title(f\"{state_name}\")\n",
    "    \n",
    "    # Colors\n",
    "    ABOVE = [1,0,0]\n",
    "    MIDDLE = [1,1,1]\n",
    "    BELOW = [0,0,0]\n",
    "    cmap = ListedColormap(np.r_[\n",
    "        np.linspace(BELOW,MIDDLE,25),\n",
    "        np.linspace(MIDDLE,ABOVE,25)\n",
    "    ])\n",
    "    color_mapped = lambda y: np.clip(y, .5, 1.5)-.5\n",
    "    \n",
    "    index = result['ML'].index.get_level_values('date')\n",
    "    values = result['ML'].values\n",
    "    \n",
    "    # Plot dots and line\n",
    "    ax.plot(index, values, c='k', zorder=1, alpha=.25)\n",
    "    ax.scatter(index,\n",
    "               values,\n",
    "               s=40,\n",
    "               lw=.5,\n",
    "               c=cmap(color_mapped(values)),\n",
    "               edgecolors='k', zorder=2)\n",
    "    \n",
    "    # Aesthetically, extrapolate credible interval by 1 day either side\n",
    "    lowfn = interp1d(date2num(index),\n",
    "                     result['Low'].values,\n",
    "                     bounds_error=False,\n",
    "                     fill_value='extrapolate')\n",
    "    \n",
    "    highfn = interp1d(date2num(index),\n",
    "                      result['High'].values,\n",
    "                      bounds_error=False,\n",
    "                      fill_value='extrapolate')\n",
    "    \n",
    "    extended = pd.date_range(start=pd.Timestamp('2020-03-01'),\n",
    "                             end=index[-1]+pd.Timedelta(days=1))\n",
    "    \n",
    "    ax.fill_between(extended,\n",
    "                    lowfn(date2num(extended)),\n",
    "                    highfn(date2num(extended)),\n",
    "                    color='k',\n",
    "                    alpha=.1,\n",
    "                    lw=0,\n",
    "                    zorder=3)\n",
    "\n",
    "    ax.axhline(1.0, c='k', lw=1, label='$R_t=1.0$', alpha=.25);\n",
    "    \n",
    "    # Formatting\n",
    "    ax.xaxis.set_major_locator(mdates.MonthLocator())\n",
    "    ax.xaxis.set_major_formatter(mdates.DateFormatter('%b'))\n",
    "    ax.xaxis.set_minor_locator(mdates.DayLocator())\n",
    "    \n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    ax.yaxis.set_major_formatter(ticker.StrMethodFormatter(\"{x:.1f}\"))\n",
    "    ax.yaxis.tick_right()\n",
    "    ax.spines['left'].set_visible(False)\n",
    "    ax.spines['bottom'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.margins(0)\n",
    "    ax.grid(which='major', axis='y', c='k', alpha=.1, zorder=-2)\n",
    "    ax.margins(0)\n",
    "    ax.set_ylim(0.0,3.5)\n",
    "    ax.set_xlim(pd.Timestamp('2020-03-01'), result.index.get_level_values('date')[-1]+pd.Timedelta(days=1))\n",
    "    fig.set_facecolor('w')\n",
    "\n",
    "    \n",
    "fig, ax = plt.subplots(figsize=(600/72,400/72))\n",
    "\n",
    "plot_rt(result, ax, state_name)\n",
    "ax.set_title(f'Real-time $R_t$ for {state_name}')\n",
    "ax.set_ylim(.5,3.5)\n",
    "ax.xaxis.set_major_locator(mdates.WeekdayLocator())\n",
    "ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %d'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Repeat the Process for Every State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {}\n",
    "\n",
    "states_to_process = states.loc[~states.index.get_level_values('state').isin(FILTERED_REGIONS)]\n",
    "\n",
    "for state_name, cases in states_to_process.groupby(level='state'):\n",
    "    clear_output(wait=True)\n",
    "    print(f'Processing {state_name}')\n",
    "    new, smoothed = prepare_cases(cases)\n",
    "    print('\\tGetting Posteriors')\n",
    "    try:\n",
    "        posteriors = get_posteriors(smoothed)\n",
    "    except:\n",
    "        display(cases)\n",
    "    print('\\tGetting HDIs')\n",
    "    hdis = highest_density_interval(posteriors)\n",
    "    print('\\tGetting most likely values')\n",
    "    most_likely = posteriors.idxmax().rename('ML')\n",
    "    result = pd.concat([most_likely, hdis], axis=1)\n",
    "    results[state_name] = result.droplevel(0)\n",
    "    \n",
    "clear_output(wait=True)\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot All US States"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ncols = 4\n",
    "nrows = int(np.ceil(len(results) / ncols))\n",
    "\n",
    "# fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, nrows*3))\n",
    "fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, nrows*3))\n",
    "\n",
    "for i, (state_name, result) in enumerate(results.items()):\n",
    "    plot_rt(result, axes.flat[i], state_name)\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.set_facecolor('w')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export Data to CSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "overall = None\n",
    "\n",
    "for state_name, result in results.items():\n",
    "    r = result.copy()\n",
    "    r.index = pd.MultiIndex.from_product([[state_name], result.index])\n",
    "    if overall is None:\n",
    "        overall = r\n",
    "    else:\n",
    "        overall = pd.concat([overall, r])\n",
    "\n",
    "overall.sort_index(inplace=True)\n",
    "\n",
    "# Uncomment this line if you'd like to export\n",
    "# overall.to_csv('data/rt.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "overall.to_csv('data/rt.csv')# , index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# overall.index.set_names(['state', 'date'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('data/rt.csv')# , index=['date', 'state'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.set_index(['Unnamed: 0', 'date'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f'{overall.index.get_level_values(1).max():%Y-%m-%d}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states.index.get_level_values(1).max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "overall.index.get_level_values(1).max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As of 4/12\n",
    "no_lockdown = [\n",
    "    'North Dakota',\n",
    "    'South Dakota',\n",
    "    'Nebraska',\n",
    "    'Iowa',\n",
    "    'Arkansas'\n",
    "]\n",
    "partial_lockdown = [\n",
    "    'Utah',\n",
    "    'Wyoming',\n",
    "    'Oklahoma'\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered = overall.index.get_level_values(0).isin(FILTERED_REGIONS)\n",
    "mr = overall.loc[~filtered].groupby(level=0)[['ML', 'High', 'Low']].last()\n",
    "\n",
    "def plot_standings(mr, figsize=None, title='Most Recent $R_t$ by State'):\n",
    "    if not figsize:\n",
    "        figsize = ((15.9/50)*len(mr)+.1,2.5)\n",
    "        \n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "\n",
    "    ax.set_title(title)\n",
    "    err = mr[['Low', 'High']].sub(mr['ML'], axis=0).abs()\n",
    "    bars = ax.bar(mr.index,\n",
    "                  mr['ML'],\n",
    "                  width=.825,\n",
    "                  color=FULL_COLOR,\n",
    "                  ecolor=ERROR_BAR_COLOR,\n",
    "                  capsize=2,\n",
    "                  error_kw={'alpha':.5, 'lw':1},\n",
    "                  yerr=err.values.T)\n",
    "\n",
    "    for bar, state_name in zip(bars, mr.index):\n",
    "        if state_name in no_lockdown:\n",
    "            bar.set_color(NONE_COLOR)\n",
    "        if state_name in partial_lockdown:\n",
    "            bar.set_color(PARTIAL_COLOR)\n",
    "\n",
    "    labels = mr.index.to_series().replace({'District of Columbia':'DC'})\n",
    "    ax.set_xticklabels(labels, rotation=90, fontsize=11)\n",
    "    ax.margins(0)\n",
    "    ax.set_ylim(0,2.)\n",
    "    ax.axhline(1.0, linestyle=':', color='k', lw=1)\n",
    "t\n",
    "    leg = ax.legend(handles=[\n",
    "                        Patch(label='Full', color=FULL_COLOR),\n",
    "                        Patch(label='Partial', color=PARTIAL_COLOR),\n",
    "                        Patch(label='None', color=NONE_COLOR)\n",
    "                    ],\n",
    "                    title='Lockdown',\n",
    "                    ncol=3,\n",
    "                    loc='upper left',\n",
    "                    columnspacing=.75,\n",
    "                    handletextpad=.5,\n",
    "                    handlelength=1)\n",
    "\n",
    "    leg._legend_box.align = \"left\"\n",
    "    fig.set_facecolor('w')\n",
    "    return fig, ax\n",
    "\n",
    "mr.sort_values('ML', inplace=True)\n",
    "plot_standings(mr);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.sort_values('High', inplace=True)\n",
    "plot_standings(mr);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show = mr[mr.High.le(1.1)].sort_values('ML')\n",
    "fig, ax = plot_standings(show, title='Likely Under Control');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show = mr[mr.Low.ge(1.05)].sort_values('Low')\n",
    "fig, ax = plot_standings(show, title='Likely Not Under Control');\n",
    "ax.get_legend().remove()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}