{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "name": "Change_point_model.ipynb", "provenance": [], "collapsed_sections": [], "toc_visible": true }, "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.10" } }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "kqXl71t9suiG" }, "source": [ "# Estimating the Date of COVID-19 Changes\n", "\n", "References: \n", "https://nbviewer.jupyter.org/github/jramkiss/jramkiss.github.io/blob/master/_posts/notebooks/covid19-changes.ipynb " ] }, { "cell_type": "code", "metadata": { "id": "fcSBMN5p_CY8" }, "source": [ "# Parameters\n", "country_ = \"China\"\n", "end_date = \"2020-04-01\"\n", "cad_start = \"2020-01-22\"\n", "cad_end = \"2020-03-16\"" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "gFnvD8OysuiI" }, "source": [ "import pandas as pd\n", "import numpy as np\n", "\n", "import seaborn as sns; sns.set()\n", "import matplotlib.pyplot as plt\n", "import matplotlib.dates as mdates\n", "\n", "\n", "from sklearn.linear_model import LinearRegression\n", "\n", "from scipy import stats\n", "import statsmodels.api as sm\n", "import pylab\n", "\n", "# from google.colab import files\n", "# from io import StringIO\n", "# uploaded = files.upload()\n", "\n", "url = 'https://raw.githubusercontent.com/assemzh/ProbProg-COVID-19/master/full_grouped.csv'\n", "data = pd.read_csv(url)\n", "\n", "data.Date = pd.to_datetime(data.Date)\n", "\n", "# for fancy python printing\n", "from IPython.display import Markdown, display\n", "def printmd(string):\n", " display(Markdown(string))\n", " \n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "import matplotlib as mpl\n", "mpl.rcParams['figure.dpi'] = 250" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "hzvPpvVvphTD" }, "source": [ "## Create country\n" ] }, { "cell_type": "code", "metadata": { "id": "koX5yGHrsuib" }, "source": [ "# function to make the time series of confirmed and daily confirmed cases for a specific country\n", "def create_country (country, end_date, state = False) : \n", " if state :\n", " df = data.loc[data[\"Province/State\"] == country, [\"Province/State\", \"Date\", \"Confirmed\", \"Deaths\", \"Recovered\"]]\n", " else : \n", " df = data.loc[data[\"Country/Region\"] == country, [\"Country/Region\", \"Date\", \"Confirmed\", \"Deaths\", \"Recovered\"]]\n", " df.columns = [\"country\", \"date\", \"confirmed\", \"deaths\", \"recovered\"]\n", "\n", " # group by country and date, sum(confirmed, deaths, recovered). do this because countries have multiple cities \n", " df = df.groupby(['country','date'])['confirmed', 'deaths', 'recovered'].sum().reset_index()\n", "\n", " # convert date string to datetime\n", " std_dateparser = lambda x: str(x)[5:10]\n", " df.date = pd.to_datetime(df.date)\n", " df['date_only'] = df.date.apply(std_dateparser)\n", " df = df.sort_values(by = \"date\")\n", " df = df[df.date <= end_date]\n", "\n", "\n", " # make new confirmed cases every day:\n", " cases_shifted = np.array([0] + list(df.confirmed[:-1]))\n", " daily_confirmed = np.array(df.confirmed) - cases_shifted\n", " df[\"daily_confirmed\"] = daily_confirmed \n", "\n", " fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 6))\n", " ax = [ax]\n", " sns.lineplot(x = df.date, \n", " y = df.daily_confirmed, \n", " ax = ax[0])\n", "\n", " ax[0].set(ylabel='Daily Confirmed Cases')\n", "\n", " ax[0].axvline(pd.to_datetime('2020-01-22'), \n", " linestyle = '--', linewidth = 1.5,\n", " label = \"Policy start: Jan 22, 2020\" ,\n", " color = \"red\") \n", "\n", " \n", " ax[0].xaxis.get_label().set_fontsize(22)\n", " ax[0].yaxis.get_label().set_fontsize(22)\n", " x = df.date\n", " # ax[0].set_xticks(x[::30])\n", " ax[0].xaxis.set_major_locator(mdates.MonthLocator(interval=1)) #to get a tick every month\n", "\n", " ax[0].title.set_fontsize(20)\n", " ax[0].tick_params(labelsize=22)\n", " myFmt = mdates.DateFormatter('%b %-d, %Y')\n", " ax[0].xaxis.set_major_formatter(myFmt)\n", "\n", " ax[0].set(ylabel='', xlabel='');\n", " ax[0].legend(loc = \"bottom right\", fontsize=22)\n", "\n", "\n", " sns.set_style(\"ticks\")\n", " plt.tight_layout()\n", " sns.despine()\n", " plt.savefig('/content/sample_data/country_daily.pdf')\n", " print(df.tail())\n", " return df\n", "\n", "\n", "def summary(samples):\n", " site_stats = {}\n", " for k, v in samples.items():\n", " site_stats[k] = {\n", " \"mean\": torch.mean(v, 0),\n", " \"std\": torch.std(v, 0),\n", " \"5%\": v.kthvalue(int(len(v) * 0.05), dim=0)[0],\n", " \"95%\": v.kthvalue(int(len(v) * 0.95), dim=0)[0],\n", " }\n", " return site_stats" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "w_A0fd4Zsuiw" }, "source": [ "cad = create_country(country_, end_date)" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "UR0BM7TysujG" }, "source": [ "cad = cad[cad.date >= cad_start].reset_index(drop = True)\n", "cad[\"days_since_start\"] = np.arange(cad.shape[0]) + 1" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "OaTHo6I2sujp" }, "source": [ "cad.shape\n", "cad_tmp = cad[cad.date < cad_end]\n", "cad_tmp.shape" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "loi3CtjSsuoz" }, "source": [ "## Data for Regression" ] }, { "cell_type": "code", "metadata": { "id": "Os_M7r4Tsuo4" }, "source": [ "reg_data = cad_tmp.copy()\n", "reg_data.head()" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "JkO0Z8M0supC" }, "source": [ "## Change Point Estimation in Pyro" ] }, { "cell_type": "code", "metadata": { "id": "aIUed4Ny3-oq" }, "source": [ "!pip install pyro-ppl\n", "!pip install numpyro" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "3ZS9fTPxsupD" }, "source": [ "import torch\n", "\n", "import pyro\n", "import pyro.distributions as dist\n", "from torch import nn\n", "from pyro.nn import PyroModule, PyroSample\n", "\n", "from pyro.infer import MCMC, NUTS, HMC\n", "from pyro.infer.autoguide import AutoGuide, AutoDiagonalNormal\n", "\n", "from pyro.infer import SVI, Trace_ELBO\n", "from pyro.infer import Predictive" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "9gPrEaEJsupP" }, "source": [ "# we should be able to have an empirical estimate for the mean of the prior for the 2nd regression bias term\n", "# this will be something like b = log(max(daily_confirmed))\n", "\n", "# might be able to have 1 regression model but change the data so that we have new terms for (tau < t) \n", "# like an interaction term\n", "\n", "class COVID_change(PyroModule):\n", " def __init__(self, in_features, out_features, b1_mu, b2_mu):\n", " super().__init__()\n", " self.linear1 = PyroModule[nn.Linear](in_features, out_features, bias = False)\n", " self.linear1.weight = PyroSample(dist.Normal(0.5, 0.25).expand([1, 1]).to_event(1))\n", " self.linear1.bias = PyroSample(dist.Normal(b1_mu, 1.))\n", " \n", " # could possibly have stronger priors for the 2nd regression line, because we wont have as much data\n", " self.linear2 = PyroModule[nn.Linear](in_features, out_features, bias = False)\n", " self.linear2.weight = PyroSample(dist.Normal(0., 0.25).expand([1, 1])) #.to_event(1))\n", " self.linear2.bias = PyroSample(dist.Normal(b2_mu, b2_mu/4))\n", "\n", " def forward(self, x, y=None):\n", " tau = pyro.sample(\"tau\", dist.Beta(4, 3))\n", " sigma = pyro.sample(\"sigma\", dist.Uniform(0., 3.))\n", " # fit lm's to data based on tau\n", " sep = int(np.ceil(tau.detach().numpy() * len(x)))\n", " mean1 = self.linear1(x[:sep]).squeeze(-1)\n", " mean2 = self.linear2(x[sep:]).squeeze(-1)\n", " mean = torch.cat((mean1, mean2))\n", " obs = pyro.sample(\"obs\", dist.StudentT(2, mean, sigma), obs=y)\n", " return mean" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "if0toOMysupU" }, "source": [ "tensor_data = torch.tensor(reg_data[[\"confirmed\", \"days_since_start\", \"daily_confirmed\"]].values, dtype=torch.float)\n", "x_data = tensor_data[:, 1].unsqueeze_(1)\n", "y_data = np.log(tensor_data[:, 0])\n", "y_data_daily = np.log(tensor_data[:, 2])\n", "# prior hyper params\n", "# take log of the average of the 1st quartile to get the prior mean for the bias of the 2nd regression line\n", "q1 = np.quantile(y_data, q = 0.25)\n", "bias_1_mean = np.mean(y_data.numpy()[y_data <= q1])\n", "print(\"Prior mean for Bias 1: \", bias_1_mean)\n", "\n", "# take log of the average of the 4th quartile to get the prior mean for the bias of the 2nd regression line\n", "q4 = np.quantile(y_data, q = 0.75)\n", "bias_2_mean = np.mean(y_data.numpy()[y_data >= q4])\n", "print(\"Prior mean for Bias 2: \", bias_2_mean)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "nrm8RrFasupc" }, "source": [ "## Approximate Inference with Stochastic Variational Inference" ] }, { "cell_type": "markdown", "metadata": { "id": "x0nDLSPisupm" }, "source": [ "# HMC with NUTS" ] }, { "cell_type": "code", "metadata": { "id": "X1rSXXtKsupm" }, "source": [ "model = COVID_change(1, 1, \n", " b1_mu = bias_1_mean,\n", " b2_mu = bias_2_mean)\n", "# need more than 400 samples/chain if we want to use a flat prior on b_2 and w_2\n", "num_samples = 400 \n", "# mcmc \n", "nuts_kernel = NUTS(model)\n", "mcmc = MCMC(nuts_kernel, \n", " num_samples=num_samples,\n", " warmup_steps = 200, \n", " num_chains = 1)\n", "mcmc.run(x_data, y_data)" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "RJPSAgdbLTTJ" }, "source": [ "samples = mcmc.get_samples()" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "7Z968a5xsupv" }, "source": [ "# extract individual posteriors\n", "weight_1_post = samples[\"linear1.weight\"].detach().numpy()\n", "weight_2_post = samples[\"linear2.weight\"].detach().numpy()\n", "bias_1_post = samples[\"linear1.bias\"].detach().numpy()\n", "bias_2_post = samples[\"linear2.bias\"].detach().numpy()\n", "tau_post = samples[\"tau\"].detach().numpy()\n", "sigma_post = samples[\"sigma\"].detach().numpy()\n", "\n", "# build likelihood distribution:\n", "tau_days = list(map(int, np.ceil(tau_post * len(x_data))))\n", "mean_ = torch.zeros(len(tau_days), len(x_data))\n", "obs_ = torch.zeros(len(tau_days), len(x_data))\n", "for i in range(len(tau_days)) : \n", " mean_[i, :] = torch.cat((x_data[:tau_days[i]] * weight_1_post[i] + bias_1_post[i],\n", " x_data[tau_days[i]:] * weight_2_post[i] + bias_2_post[i])).reshape(len(x_data))\n", " obs_[i, :] = dist.Normal(mean_[i, :], sigma_post[i]).sample()\n", "samples[\"_RETURN\"] = mean_\n", "samples[\"obs\"] = obs_\n", "pred_summary = summary(samples)\n", "mu = pred_summary[\"_RETURN\"] # mean\n", "y = pred_summary[\"obs\"] # samples from likelihood: mu + sigma\n", "y_shift = np.exp(y[\"mean\"]) - np.exp(torch.cat((y[\"mean\"][0:1], y[\"mean\"][:-1])))\n", "print(y_shift)\n", "predictions = pd.DataFrame({\n", " \"days_since_start\": x_data[:, 0],\n", " \"mu_mean\": mu[\"mean\"], # mean of likelihood\n", " \"mu_perc_5\": mu[\"5%\"],\n", " \"mu_perc_95\": mu[\"95%\"],\n", " \"y_mean\": y[\"mean\"], # mean of likelihood + noise\n", " \"y_perc_5\": y[\"5%\"],\n", " \"y_perc_95\": y[\"95%\"],\n", " \"true_confirmed\": y_data,\n", " \"true_daily_confirmed\": y_data_daily,\n", " \"y_daily_mean\": y_shift\n", "})\n", "\n", "w1_ = pred_summary[\"linear1.weight\"]\n", "w2_ = pred_summary[\"linear2.weight\"]\n", "\n", "b1_ = pred_summary[\"linear1.bias\"]\n", "b2_ = pred_summary[\"linear2.bias\"]\n", "\n", "tau_ = pred_summary[\"tau\"]\n", "sigma_ = pred_summary[\"sigma\"]\n", "\n", "ind = int(np.ceil(tau_[\"mean\"] * len(x_data)))" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "KegzMbLOsuqC" }, "source": [ "## Model Diagnostics\n", "\n", "- Residual plots: Should these be samples from the likelihood compared with the actual data? Or just the mean of the likelihood?\n", "- $\\hat{R}$: The factor that the scale of the current distribution will be reduced by if we were to run the simulations forever. As n tends to $\\inf$, $\\hat{R}$ tends to 1. So we want values close to 1.\n", "- Mixing and Stationarity: I sampled 4 chains. Do I then take these chains, split them in half and plot them. If they converge to the same stationary distribution, does that mean the MCMC converged? What do I do with more sampled chains?" ] }, { "cell_type": "code", "metadata": { "id": "QUh6fBjtsuqV" }, "source": [ "mcmc.summary()\n", "diag = mcmc.diagnostics()" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "LJ2a6Epnsuqf" }, "source": [ "## Posterior Plots" ] }, { "cell_type": "code", "metadata": { "id": "bPpET7b6suqg" }, "source": [ "print(ind)\n", "print(reg_data.date[ind])\n", "\n", "sns.distplot(weight_1_post, \n", " kde_kws = {\"label\": \"Weight posterior before CP\"}, \n", " # color = \"red\",\n", " norm_hist = True,\n", " kde = True)\n", "plt.axvline(x = w1_[\"mean\"], linestyle = '--',label = \"Mean weight before CP\" ,\n", " # color = \"red\"\n", " )\n", "\n", "sns.distplot(weight_2_post, \n", " kde_kws = {\"label\": \"Weight posterior after CP\"}, \n", " color = \"red\",\n", " norm_hist = True,\n", " kde = True)\n", "plt.axvline(x = w2_[\"mean\"], linestyle = '--',label = \"Mean weight after CP\" ,\n", " color = \"red\")\n", "\n", "legend = plt.legend(loc='upper right')\n", "legend.get_frame().set_alpha(1)\n", "sns.set_style(\"ticks\")\n", "plt.tight_layout()\n", "sns.despine()\n", "plt.savefig('/content/sample_data/country_weights.pdf')\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "Q3-M1WQtrTbx" }, "source": [ "print(w1_[\"mean\"])\n", "print(w2_[\"mean\"])" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "A4glDCQRO9Ss" }, "source": [ "predictions['date'] = pd.to_datetime(reg_data.date)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "DAvMkdZagKu2" }, "source": [ "# Final plot" ] }, { "cell_type": "code", "metadata": { "id": "gejnd0Qxsuq1" }, "source": [ "start_date_ = str(reg_data.date[0]).split(' ')[0]\n", "change_date_ = str(reg_data.date[ind]).split(' ')[0]\n", "print(\"Date of change for {}: {}\".format(country_, change_date_))\n", "import seaborn as sns\n", "\n", "# plot data:\n", "fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(7, 5))\n", "ax = [ax]\n", "# log regression model\n", "ax[0].scatter(y = np.exp(y_data[:ind]), x = x_data[:ind], s = 15);\n", "ax[0].scatter(y = np.exp(y_data[ind:]), x = x_data[ind:], s = 15, color = \"red\");\n", "\n", "ax[0].plot(predictions[\"days_since_start\"],\n", " np.exp(predictions[\"y_mean\"]), \n", " color = \"green\",\n", " label = \"Fitted line by MCMC-NUTS model\") \n", "# ax[0].axvline(2, \n", "# linestyle = '--', linewidth = 1.5,\n", "# label = \"Date of Lockdown: Jan 23, 2020\" ,\n", "# color = \"red\")\n", "\n", "# ax[0].axvline(ind, \n", "# linestyle = '--', linewidth = 1.5,\n", "# label = \"Date of Change: Feb 8, 2020\",\n", "# color = \"black\")\n", "\n", "ax[0].fill_between(predictions[\"days_since_start\"], \n", " np.exp(predictions[\"y_perc_5\"]), \n", " np.exp(predictions[\"y_perc_95\"]), \n", " alpha = 0.25,\n", " label = \"90% CI of predictions\",\n", " color = \"teal\");\n", "ax[0].fill_betweenx([0, 1], \n", " tau_[\"5%\"] * len(x_data), \n", " tau_[\"95%\"] * len(x_data), \n", " alpha = 0.25,\n", " label = \"90% CI of changing point\",\n", " color = \"lightcoral\",\n", " transform=ax[0].get_xaxis_transform());\n", "ax[0].set(ylabel = \"Total Cases\",)\n", "ax[0].legend(loc = \"lower right\", fontsize=12.8)\n", "ax[0].set_ylim([100,150000])\n", "ax[0].xaxis.get_label().set_fontsize(16)\n", "ax[0].yaxis.get_label().set_fontsize(16)\n", "ax[0].title.set_fontsize(20)\n", "ax[0].tick_params(labelsize=16)\n", "\n", "# plt.xticks(ticks=[1,17,34,51], labels=[\"Jan 22\",\n", " # \"Feb 8\",\n", " # \"Feb 25\",\n", " # \"Mar 13\"], fontsize=15)\n", "ax[0].set_yscale('log')\n", "plt.setp(ax[0].get_xticklabels(), rotation=0, horizontalalignment='center')\n", "print(reg_data.columns)\n", "myFmt = mdates.DateFormatter('%m-%d')\n", "sns.set_style(\"ticks\")\n", "sns.despine()\n", "ax[0].figure.savefig('/content/sample_data/country_cp.pdf')\n" ], "execution_count": null, "outputs": [] } ] }