covid19_sir.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%matplotlib inline\n",
    "import datetime\n",
    "\n",
    "import numpy as np\n",
    "from IPython.display import display, Markdown\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import arviz as az\n",
    "import pymc3 as pm\n",
    "import theano.tensor as tt\n",
    "from sunode.wrappers.as_theano import solve_ivp\n",
    "import sunode\n",
    "import sympy as sym\n",
    "\n",
    "import requests\n",
    "import io\n",
    "\n",
    "sns.set_context('talk')\n",
    "plt.style.use('seaborn-whitegrid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import load_covid_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sign = load_covid_data.load_data(filter_n_days_100=10)\n",
    "df_sign = df_sign[df_sign.state.isna()]\n",
    "\n",
    "countries = pd.Series(df_sign.country.unique())\n",
    "countries = countries.loc[~countries.isin(['Cruise Ship (total)', 'Cruise Ship'])]\n",
    "n_countries = len(countries)\n",
    "countries = countries.rename('country').to_frame().set_index('country')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sign = df_sign.reset_index()\n",
    "df_sign.index.name = 'observation'\n",
    "\n",
    "previous = (\n",
    "    df_sign\n",
    "    .reset_index()\n",
    "    .groupby(['country'])\n",
    "    .apply(\n",
    "        lambda x:\n",
    "        x.set_index('date')\n",
    "        .assign(date=x.date.values)\n",
    "        .sort_index()\n",
    "        .shift()\n",
    "    )\n",
    ")\n",
    "\n",
    "df_sign = (\n",
    "    df_sign\n",
    "    .reset_index()\n",
    "    .set_index(['country', 'date'])\n",
    "    .assign(\n",
    "        previous_date=previous.date,\n",
    "        previous_confirmed=previous.confirmed,\n",
    "    )\n",
    "    .reset_index()\n",
    "    .set_index('observation')\n",
    ")\n",
    "\n",
    "df_sign['new_confirmed'] = df_sign.eval('confirmed - previous_confirmed')\n",
    "df_sign['days_since_last_data'] = (df_sign['date'] - df_sign['previous_date']).dt.days\n",
    "df_sign = df_sign.dropna(subset=['new_confirmed'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Growth Rate Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_file = 'https://github.com/datasets/population/raw/master/data/population.csv'\n",
    "populations = pd.read_csv(pop_file)\n",
    "populations = populations.sort_values(['Country Name', 'Year']).groupby(['Country Name']).last()\n",
    "\n",
    "idxs = populations.index.get_indexer([c.replace('(total)', '').strip() for c in countries.index])\n",
    "countries['population'] = [populations.iloc[i].Value if idx != -1 else np.nan for i, idx in enumerate(idxs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "countries.loc['Egypt', 'population'] = populations.loc['Egypt, Arab Rep.'].Value\n",
    "countries.loc['Iran', 'population'] = populations.loc['Iran, Islamic Rep.'].Value\n",
    "countries.loc['Korea, South', 'population'] = populations.loc['Korea, Rep.'].Value\n",
    "countries.loc['Hong Kong', 'population'] = populations.loc['Hong Kong SAR, China'].Value\n",
    "countries.loc['Czechia', 'population'] = populations.loc['Czech Republic'].Value\n",
    "countries.loc['US (total)', 'population'] = populations.loc['United States'].Value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "countries['population'] = countries.population.values.astype('d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_country = True\n",
    "\n",
    "if multi_country:\n",
    "    max_countries = 50\n",
    "    countries = countries.iloc[:n_countries]\n",
    "    n_countries = len(countries)\n",
    "    df_sign = df_sign[df_sign.country.isin(countries.index)]\n",
    "else:\n",
    "    n_countries = 1\n",
    "    country = 'Italy'\n",
    "\n",
    "    countries = countries.loc[country: country]\n",
    "    df_sign = df_sign[df_sign.country == country]    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_day = df_sign.date.min() - datetime.timedelta(days=1)\n",
    "end_data_day = df_sign.date.max()\n",
    "end_day = end_data_day + datetime.timedelta(days=30)\n",
    "\n",
    "data_time = pd.date_range(start_day, end_data_day)\n",
    "time = pd.date_range(start_day, end_day)\n",
    "country = countries.index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/twiecki/miniconda3/envs/pymc3/lib/python3.7/site-packages/theano/tensor/subtensor.py:2197: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "  rval = inputs[0].__getitem__(inputs[1:])\n",
      "/Users/twiecki/miniconda3/envs/pymc3/lib/python3.7/site-packages/theano/tensor/subtensor.py:2197: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "  rval = inputs[0].__getitem__(inputs[1:])\n"
     ]
    },
    {
     "ename": "KeyError",
     "evalue": "'recovered'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m~/miniconda3/envs/pymc3/lib/python3.7/site-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36mget_loc\u001b[0;34m(self, key, method, tolerance)\u001b[0m\n\u001b[1;32m   2645\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2646\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_engine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_loc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2647\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/index.pyx\u001b[0m in \u001b[0;36mpandas._libs.index.IndexEngine.get_loc\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/index.pyx\u001b[0m in \u001b[0;36mpandas._libs.index.IndexEngine.get_loc\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/hashtable_class_helper.pxi\u001b[0m in \u001b[0;36mpandas._libs.hashtable.PyObjectHashTable.get_item\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/hashtable_class_helper.pxi\u001b[0m in \u001b[0;36mpandas._libs.hashtable.PyObjectHashTable.get_item\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'recovered'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-23-d40edb3e475b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    191\u001b[0m         \u001b[0mmu_recovered\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    192\u001b[0m         \u001b[0malpha\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx_country\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 193\u001b[0;31m         \u001b[0mobserved\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaximum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdf_sign\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'recovered'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mdf_sign\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'deaths'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    194\u001b[0m     )\n",
      "\u001b[0;32m~/miniconda3/envs/pymc3/lib/python3.7/site-packages/pandas/core/frame.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m   2798\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolumns\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnlevels\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2799\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_getitem_multilevel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2800\u001b[0;31m             \u001b[0mindexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolumns\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_loc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2801\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mis_integer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2802\u001b[0m                 \u001b[0mindexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mindexer\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/pymc3/lib/python3.7/site-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36mget_loc\u001b[0;34m(self, key, method, tolerance)\u001b[0m\n\u001b[1;32m   2646\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_engine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_loc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2647\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2648\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_engine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_loc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_cast_indexer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2649\u001b[0m         \u001b[0mindexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_indexer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmethod\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtolerance\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtolerance\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2650\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mindexer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndim\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mindexer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/index.pyx\u001b[0m in \u001b[0;36mpandas._libs.index.IndexEngine.get_loc\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/index.pyx\u001b[0m in \u001b[0;36mpandas._libs.index.IndexEngine.get_loc\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/hashtable_class_helper.pxi\u001b[0m in \u001b[0;36mpandas._libs.hashtable.PyObjectHashTable.get_item\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/hashtable_class_helper.pxi\u001b[0m in \u001b[0;36mpandas._libs.hashtable.PyObjectHashTable.get_item\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'recovered'"
     ]
    }
   ],
   "source": [
    "log_scale = False\n",
    "\n",
    "if not log_scale:\n",
    "    def sir_rhs(t, y, p):\n",
    "        exp = np.vectorize(sym.exp)\n",
    "        infections = exp(y.log_infection_rate) * y.S * y.I\n",
    "        remissions = p.remission_rate * y.I\n",
    "\n",
    "        change = [\n",
    "            sunode.symode.lambdify.interpolate_spline(t, p.infection_change_polys[i], 0, len(data_time), degree=4)\n",
    "            for i in range(len(infections))\n",
    "        ]\n",
    "\n",
    "        return {\n",
    "            'S': -infections,\n",
    "            'I': infections - remissions,\n",
    "            'R': remissions,\n",
    "            'log_infection_rate': change,\n",
    "        }\n",
    "else:\n",
    "    def sir_rhs(t, y, p):\n",
    "        exp = np.vectorize(sym.exp)\n",
    "        infections = exp(y.log_infection_rate + y.I + y.S)\n",
    "        remissions = p.remission_rate * exp(y.I)\n",
    "\n",
    "        change = [\n",
    "            sunode.symode.lambdify.interpolate_spline(t, p.infection_change_polys[i], 0, len(data_time), degree=4)\n",
    "            for i in range(len(infections))\n",
    "        ]\n",
    "\n",
    "        return {\n",
    "            'S': -infections / exp(y.S),\n",
    "            'I': (infections - remissions) / exp(y.I),\n",
    "            'R': remissions / exp(y.R),\n",
    "            'log_infection_rate': change,\n",
    "        }\n",
    "\n",
    "\n",
    "with pm.Model() as model:\n",
    "    log_pop = np.log(countries['population'].values)\n",
    "    if multi_country:\n",
    "        sd = pm.HalfNormal('init_logI_sd')\n",
    "        mu = pm.Normal('init_logI_mu')\n",
    "        raw = pm.Normal('init_logI_raw', shape=n_countries)\n",
    "        init_logI = mu + sd * raw - log_pop\n",
    "        pm.Deterministic('init_logI', init_logI)\n",
    "    else:\n",
    "        init_logI = pm.Normal('init_logI_raw', shape=n_countries)\n",
    "        init_logI = init_logI - log_pop\n",
    "        pm.Deterministic('init_logI', init_logI)\n",
    "\n",
    "    if multi_country:\n",
    "        sd = pm.HalfNormal('log_infect_sd')\n",
    "        mu = pm.Normal('log_infect_mu')\n",
    "        raw = pm.Normal('log_infect_raw', shape=n_countries)\n",
    "        log_infect = mu + sd * raw\n",
    "        pm.Deterministic('log_infect', log_infect)\n",
    "    else:\n",
    "        log_infect = pm.Normal('log_infect', shape=n_countries)\n",
    "\n",
    "    if multi_country:\n",
    "        sd = pm.HalfNormal('log_remiss_sd')  # In 1/day/infected\n",
    "        mu = pm.Normal('log_remiss_mu')\n",
    "        raw = pm.Normal('log_remiss_raw', shape=n_countries)\n",
    "        log_remiss = mu + sd * raw\n",
    "        pm.Deterministic('log_remiss', log_remiss)\n",
    "    else:\n",
    "        log_remiss = pm.Normal('log_remiss', shape=n_countries)\n",
    "        \n",
    "    sd = pm.HalfNormal('infection_change_polys_sd')\n",
    "    raw = pm.StudentT('infection_change_polys_raw', shape=(n_countries, 8), nu=5)\n",
    "    infection_change_polys = sd * raw / 100\n",
    "\n",
    "    if not log_scale:\n",
    "        init = {\n",
    "            'S': (1 - tt.exp(init_logI), (n_countries,)),\n",
    "            'I': (tt.exp(init_logI), (n_countries,)),\n",
    "            'R': (np.zeros(n_countries), (n_countries,)),\n",
    "            'log_infection_rate': (log_infect, (n_countries,)),\n",
    "        }\n",
    "    else:\n",
    "        init = {\n",
    "            'S': (tt.log1p(-tt.exp(init_logI)), (n_countries,)),\n",
    "            'I': (init_logI, (n_countries,)),\n",
    "            'R': np.full(n_countries, -20, dtype='d'),\n",
    "            'log_infection_rate': (log_infect, (n_countries,)),\n",
    "        }\n",
    "    \n",
    "    params = {\n",
    "        #'infection_rate': (tt.exp(log_infect), (n_countries,)),\n",
    "        'remission_rate': (tt.exp(log_remiss), (n_countries,)),\n",
    "        'infection_change_polys': (infection_change_polys, (n_countries, 8)),\n",
    "        '_placeholder': (np.zeros(n_countries), (n_countries,)),\n",
    "    }\n",
    "    \n",
    "    solutions = []\n",
    "    problems = []\n",
    "    solvers = []\n",
    "    for i in range(n_countries):\n",
    "        init_i = {\n",
    "            name: (var[i:i+1], (1, *shape_others))\n",
    "            for i, (name, (var, (_, *shape_others))) in enumerate(init.items())\n",
    "        }\n",
    "        params_i = {\n",
    "            name: (var[i:i+1], (1, *shape_others))\n",
    "            for i, (name, (var, (_, *shape_others))) in enumerate(params.items())\n",
    "        }\n",
    "        \n",
    "        solution_i, flat_solution, problem, solver, *_ = solve_ivp(\n",
    "            t0=0,\n",
    "            y0=init_i,\n",
    "            params=params_i,\n",
    "            tvals=np.arange(len(data_time)),\n",
    "            rhs=sir_rhs,\n",
    "            solver_kwargs={\n",
    "                'solver': 'ADAMS',\n",
    "                #'constraints': np.r_[np.ones(n_countries * 3), np.zeros(n_countries)],\n",
    "                'constraints': np.array([1., 1, 1, 0]),\n",
    "            }\n",
    "        )\n",
    "        problems.append(problem)\n",
    "        solvers.append(solver)\n",
    "        solutions.append(solution_i)\n",
    "    \n",
    "    solution = {}\n",
    "    for name in init:\n",
    "        solution[name] = tt.join(1, *[vals[name] for vals in solutions])\n",
    "\n",
    "        \n",
    "    for name in init:\n",
    "        if log_scale:\n",
    "            pm.Deterministic(f'solution_{name}', tt.exp(solution[name]))\n",
    "        else:\n",
    "            pm.Deterministic(f'solution_{name}', solution[name])\n",
    "\n",
    "    # The number of people who have at some time been infected\n",
    "    if log_scale:\n",
    "        log_cases = tt.log1p(-tt.exp(solution['S'])) + np.log(countries.population.values[None, :])\n",
    "    else:\n",
    "        cases = (1 - solution['S']) * countries.population.values[None, :]\n",
    "    #cases = solution['I'] * countries.population.values[None, :]\n",
    "\n",
    "    idx_country = countries.index.get_indexer(df_sign.country)\n",
    "    idx_data_time = data_time.get_indexer(df_sign.date)\n",
    "    idx_previous_time = data_time.get_indexer(df_sign.previous_date)\n",
    "    \n",
    "    assert -1 not in idx_country\n",
    "    assert -1 not in idx_data_time\n",
    "    assert -1 not in idx_previous_time\n",
    "    \n",
    "    if log_scale:\n",
    "        log_current_infected = log_cases[idx_data_time, idx_country]\n",
    "        log_previous_infected = log_cases[idx_previous_time, idx_country]\n",
    "        #mu = tt.exp(log_current_infected) - tt.exp(log_previous_infected) + 1\n",
    "        mu_confirmed = tt.exp(log_current_infected)\n",
    "    else:\n",
    "        current_infected = cases[idx_data_time, idx_country]\n",
    "        previous_infected = cases[idx_previous_time, idx_country]\n",
    "        #mu = current_infected - previous_infected + 1\n",
    "        mu_confirmed = current_infected\n",
    "    \n",
    "    alpha = pm.HalfNormal('confirmed_alpha', shape=n_countries, sd=10)\n",
    "    pm.Deterministic('mu_confirmed', mu_confirmed)\n",
    "    pm.NegativeBinomial(\n",
    "        'confirmed_count',\n",
    "        mu_confirmed,\n",
    "        alpha[idx_country],\n",
    "        observed=np.maximum(0, df_sign['confirmed'].values)\n",
    "    )\n",
    "    \n",
    "    if log_scale:\n",
    "        log_cases = solution['R'] + np.log(countries.population.values[None, :])\n",
    "    else:\n",
    "        cases = solution['R'] * countries.population.values[None, :]\n",
    "    \n",
    "    if log_scale:\n",
    "        log_current_recovered = log_cases[idx_data_time, idx_country]\n",
    "        log_previous_recovered = log_cases[idx_previous_time, idx_country]\n",
    "        #mu = tt.exp(log_current_infected) - tt.exp(log_previous_infected) + 1\n",
    "        mu_recovered = tt.exp(log_current_recovered)\n",
    "    else:\n",
    "        current_recovered = cases[idx_data_time, idx_country]\n",
    "        previous_recovered = cases[idx_previous_time, idx_country]\n",
    "        #mu = current_infected - previous_infected + 1\n",
    "        mu_recovered = current_recovered\n",
    "    \n",
    "    alpha = pm.HalfNormal('recovered_alpha', shape=n_countries, sd=10)\n",
    "    pm.Deterministic('mu_recovered', mu_recovered)\n",
    "    pm.NegativeBinomial(\n",
    "        'recovered_count',\n",
    "        mu_recovered,\n",
    "        alpha[idx_country],\n",
    "        observed=np.maximum(0, df_sign['recovered'].values + df_sign['deaths'].values)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solutions[1]['R'].tag.test_value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases.tag.test_value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lib = sunode._cvodes.lib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    jac_prod = problem.make_sundials_jac_prod()\n",
    "    lin_solver = lib.SUNLinSol_SPGMR(solver._state_buffer.c_ptr, lib.PREC_NONE, 5)\n",
    "    lib.SUNLinSolInitialize_SPGMR(lin_solver)\n",
    "\n",
    "    lib.CVodeSetLinearSolver(solver._ode, lin_solver, sunode.basic.ffi.NULL)\n",
    "\n",
    "    lib.CVodeSetJacTimes(solver._ode, sunode.basic.ffi.NULL, jac_prod.cffi)\n",
    "\n",
    "    jac_prod = problem.make_sundials_adjoint_jac_prod()\n",
    "    lin_solverB = lib.SUNLinSol_SPGMR(solver._state_buffer.c_ptr, lib.PREC_NONE, 5)\n",
    "    lib.SUNLinSolInitialize_SPGMR(lin_solverB)\n",
    "\n",
    "    lib.CVodeSetLinearSolverB(solver._ode, solver._odeB, lin_solverB, sunode.basic.ffi.NULL)\n",
    "\n",
    "    lib.CVodeSetJacTimesB(solver._ode, solver._odeB, sunode.basic.ffi.NULL, jac_prod.cffi)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for problem, solver in zip(problems, solvers):\n",
    "    lib.CVodeSStolerances(solver._ode, 1e-7, 1e-7)\n",
    "    lib.CVodeSStolerancesB(solver._ode, solver._odeB, 1e-6, 1e-6)\n",
    "    lib.CVodeQuadSStolerancesB(solver._ode, solver._odeB, 1e-6, 1e-6)\n",
    "    lib.CVodeSetMaxNumSteps(solver._ode, 500)\n",
    "    lib.CVodeSetMaxNumStepsB(solver._ode, solver._odeB, 500)\n",
    "    #lib.CVodeSetMinStep(solver._ode, 1e-12)\n",
    "    #lib.CVodeSetMinStepB(solver._ode, solver._odeB, 1e-12)\n",
    "    lib.CVodeSetMaxOrd(solver._ode, 4)\n",
    "    lib.CVodeSetMaxOrdB(solver._ode, solver._odeB, 4)\n",
    "    #lib.CVodeSetInitStep(solver._ode, 1e-7)\n",
    "    #lib.CVodeSetInitStepB(solver._ode, solver._odeB, -1e-6)\n",
    "    #lib.CVodeSetMaxStep(solver._ode, 50)\n",
    "    #lib.CVodeSetMaxStepB(solver._ode, solver._odeB, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = model.logp_dlogp_function()\n",
    "func.set_extra_values({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/adr/anaconda3/envs/pymc/lib/python3.7/site-packages/theano/tensor/subtensor.py:2197: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.\n",
      "  rval = inputs[0].__getitem__(inputs[1:])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array(-47547.31119204),\n",
       " array([-4.97753393e+03, -1.98826271e+03, -1.52161765e+00, -3.26821370e+03,\n",
       "         2.34153375e-01,  2.34136957e-01, -1.57921282e+00, -7.67434729e-01,\n",
       "         4.69474386e-01, -5.42560044e-01,  4.63417693e-01,  4.65729754e-01,\n",
       "        -2.41962272e-01,  1.91328024e+00,  1.72491783e+00,  5.62287529e-01,\n",
       "         1.01283112e+00, -3.14247333e-01,  9.08024076e-01,  1.41230370e+00,\n",
       "        -1.92681031e+04, -4.00767998e+04, -6.75282047e-02,  1.42474819e+00,\n",
       "         5.44382725e-01, -1.73547728e+05,  1.15099358e+00, -3.75698018e-01,\n",
       "         6.00638690e-01,  2.91693750e-01,  6.01706612e-01, -1.85227818e+00,\n",
       "         1.34972247e-02,  1.05771093e+00, -8.22544912e-01,  1.22084365e+00,\n",
       "        -2.08863595e-01,  1.95967012e+00,  1.32818605e+00, -1.96861236e-01,\n",
       "        -6.66215336e+03,  2.75131562e+04,  5.75779173e+04,  3.01103696e-01,\n",
       "         1.47852199e+00,  7.19844208e-01,  4.60638771e-01, -1.05712223e+00,\n",
       "        -3.43618290e-01,  1.76304016e+00, -3.24083969e-01,  3.85082280e-01,\n",
       "         6.76922000e-01, -6.11676289e-01, -1.03099952e+00, -9.31280119e-01,\n",
       "         8.39217523e-01,  3.09212376e-01, -3.31263431e-01, -9.75545127e-01,\n",
       "        -2.32288255e+03,  2.21265401e-01,  1.06652222e+00,  1.11605347e+00,\n",
       "        -8.61304631e-01, -1.18979086e+00,  8.63226214e-02, -1.00235041e+00,\n",
       "        -4.22901760e-01, -2.45155171e+01, -5.64908086e+02, -1.31813745e+03,\n",
       "        -6.42093462e+02, -9.14923053e+01, -1.66338689e+00, -8.80596998e-01,\n",
       "        -1.04303339e-01,  3.52505639e-01, -1.09927813e-01,  1.33238439e+00,\n",
       "         2.61086481e-01, -4.17876772e-01, -1.23429196e+00,  5.90217313e-01,\n",
       "         8.58021172e-01,  5.73244434e-01, -9.40810100e-01, -3.86154427e-01,\n",
       "         6.01926601e-01, -5.85093128e-01, -1.16273905e-01, -9.78713816e-01,\n",
       "         7.66869052e-01,  3.84929185e-01,  4.56492741e-01,  1.22952144e+00,\n",
       "        -3.49219908e-01, -3.09053930e-01, -6.13611588e-03,  2.78439987e-01,\n",
       "         1.21260771e+00,  4.87521726e-01,  4.01818444e-01,  8.52934648e-01,\n",
       "         1.92541136e-01, -4.69530201e-01, -1.32244868e+00, -2.08224149e-01,\n",
       "        -3.05014018e-01,  8.92361861e-02,  1.32608243e+00,  3.18121778e-02,\n",
       "        -7.22238510e-02, -1.33538471e+00,  2.29137415e-01, -3.55393567e-01,\n",
       "         4.16440882e-02,  1.10152047e+00, -1.08735951e+00, -8.10650723e-01,\n",
       "        -8.43657691e-01,  9.36388925e-01, -1.20794617e+00,  1.20759243e+00,\n",
       "        -6.58846969e-01, -1.34135591e+00,  9.93656039e-01,  6.38598450e-01,\n",
       "        -1.19344610e-01,  5.75018745e-01,  1.25652088e+00, -8.21982889e-02,\n",
       "         1.04003158e+00, -5.43912129e-01,  9.43750840e-01, -1.25631366e+00,\n",
       "         8.37183880e-01,  3.78619463e-01, -8.62110010e-01,  1.13356183e+00,\n",
       "        -2.70156449e-01, -1.16907113e+00,  1.27174267e+00, -2.20060274e-01,\n",
       "        -3.07702959e-01, -8.35988369e-01,  1.13654683e+00,  1.17485216e+00,\n",
       "        -5.93967782e-01, -3.50204016e-01, -2.96865951e-01, -4.05991905e-01,\n",
       "         7.46946986e-01, -2.75729766e-01, -3.45747613e-01,  7.77836071e-01,\n",
       "        -1.31994999e+00, -5.44164595e-01,  1.11350581e+00, -7.25331655e-01,\n",
       "         9.82871553e-01, -8.40378289e-01, -1.09605751e+00,  8.67908172e-01,\n",
       "        -9.75061664e-01, -4.79013474e-01, -8.69018608e-01, -1.32367892e+00,\n",
       "         2.90961669e-01,  8.12198452e-01,  9.21579893e-01,  8.63969924e-01,\n",
       "         9.24121792e-02, -4.00069968e-01, -3.27021746e-01, -8.73134556e-01,\n",
       "        -1.56017428e-02, -1.22613484e+00,  3.13200682e-01, -1.31628304e+00,\n",
       "        -6.96287155e-01,  8.96808592e-01,  1.04531525e+00, -5.53211619e-01,\n",
       "         2.65503724e-01, -7.77524688e-01, -5.43539593e-01,  8.73020843e-02,\n",
       "         8.88702001e-01,  1.24597392e+00,  5.15271432e-01, -8.96217922e-01,\n",
       "        -2.54578712e-01,  1.14080991e+00, -2.06577989e-01, -4.49046903e-01,\n",
       "         9.17307994e-01, -1.83602371e-01, -6.98031600e-02,  1.08744170e+00,\n",
       "        -4.18626994e-01, -6.33120759e-01, -1.05269844e+00, -1.03474570e+00,\n",
       "         1.19832541e+00, -1.59908004e+02, -6.11705130e+02, -8.62561141e+02,\n",
       "        -2.36422868e+02, -8.87968820e+03, -9.70604735e+02, -8.78299176e+02,\n",
       "        -1.89798482e+03, -1.21455159e+03, -3.61493042e+02, -1.20464743e+03,\n",
       "        -1.81803587e+02, -3.18191975e+02, -5.28797736e+01, -2.76740937e+02,\n",
       "        -3.86149413e+03, -5.84999233e+01, -5.87752386e+02, -8.17712775e+01,\n",
       "        -2.59847776e+02, -1.82653088e+03, -2.32731723e+02, -1.41356571e+02,\n",
       "        -3.00064440e+02, -7.00352254e+02, -4.08492074e+02, -8.90550476e+02,\n",
       "        -6.07873945e+02, -3.12459010e+02, -4.73296992e+03, -9.58080827e+02,\n",
       "         2.27362579e+01, -4.40776317e+02, -2.80599280e+02, -1.72701319e+03,\n",
       "        -1.57198789e+02]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x0 = np.random.randn(func.size)\n",
    "func(x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "150"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = sunode._cvodes.ffi.new('long[1]')\n",
    "lib.CVodeGetNumNonlinSolvIters(solver._ode, out)\n",
    "out[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "61.5 ms ± 187 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit func(x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import theano\n",
    "func_dets = theano.function(model.free_RVs, model.deterministics)\n",
    "dets = {var.name: vals for var, vals in zip(model.deterministics, func_dets(**func.array_to_dict(x0)))}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import theano\n",
    "func_logps = theano.function(model.free_RVs, [var.logp_elemwiset for var in model.observed_RVs])\n",
    "logps = {var.name: vals for var, vals in zip(model.observed_RVs, func_logps(**func.array_to_dict(x0)))}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(~np.isfinite(logps['recovered_count'])).nonzero()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dets['solution_R'].min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    trace = pm.sample(tune=1000, chains=6, cores=6, discard_tuned_samples=True, compute_convergence_checks=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz\n",
    "tr = arviz.from_pymc3(\n",
    "    trace=trace,\n",
    "    coords={\n",
    "        'data_time': data_time,\n",
    "        'country': country,\n",
    "        'time': time,\n",
    "        'observation': df_sign.index,\n",
    "        'infection_poly': np.arange(8),\n",
    "    },\n",
    "    dims={\n",
    "        'init_logI_raw': ['country'],\n",
    "        'init_logI': ['country'],\n",
    "        'log_infect_raw': ['country'],\n",
    "        'log_infect': ['country'],\n",
    "        'log_remiss_raw': ['country'],\n",
    "        'log_remiss': ['country'],\n",
    "        'solution_I': ['data_time', 'country'],\n",
    "        'solution_S': ['data_time', 'country'],\n",
    "        'solution_R': ['data_time', 'country'],\n",
    "        'solution_log_infection_rate': ['data_time', 'country'],\n",
    "        'infection_change_polys_raw': ['country', 'infection_poly'],\n",
    "        'infection_change_polys': ['country', 'infection_poly'],\n",
    "        'confirmed_alpha': ['country'],\n",
    "        'recovered_alpha': ['country'],\n",
    "        'mu_confirmed': ['observation'],\n",
    "        'mu_recovered': ['observation'],\n",
    "    },\n",
    "    log_likelihood=False,\n",
    ").posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xr.DataArray(countries['population'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arviz.plot_trace(tr, var_names=['init_logI'], compact=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(tr.mu_confirmed - xr.DataArray(df_sign.confirmed, dims=('observation',))).isel(draw=-1).plot.line(x='observation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(tr.mu_recovered - xr.DataArray(df_sign.recovered + df_sign.deaths, dims=('observation',))).isel(draw=-1).plot.line(x='observation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "((1 - tr.solution_S) * xr.DataArray(countries['population'])).isel(country=1, draw=-1).plot.line(x='data_time', add_legend=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "((tr.solution_I) * xr.DataArray(countries['population'])).isel(country=1, draw=-1).plot.line(x='data_time', add_legend=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 430,
   "metadata": {},
   "outputs": [],
   "source": [
    "#tr.solution_log_infection_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    (tr.infection_change_polys_sd * tr.infection_change_polys_raw)\n",
    "    .cumsum('infection_poly')\n",
    "    .median(['draw', 'chain'])\n",
    "    .plot\n",
    "    .line(x='infection_poly')\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vals = (tr.infection_change_polys_raw * tr.infection_change_polys_sd / 100).transpose('infection_poly', ...).values\n",
    "degree = 4\n",
    "lower = 0\n",
    "upper = len(tr.data_time)\n",
    "\n",
    "n_knots = degree + len(vals) + 1\n",
    "knots = np.arange(n_knots)\n",
    "lower_knot = degree #+ lower * (n_knots - 2 * degree - 1)\n",
    "upper_knot = degree + (n_knots - 2 * degree - 1)\n",
    "\n",
    "knots = (knots - lower_knot) / (upper_knot - lower_knot)\n",
    "knots = knots * (upper - lower) + lower\n",
    "\n",
    "x = np.arange(upper)\n",
    "spline = interpolate.BSpline(knots, vals.reshape((vals.shape[0], -1)), degree)\n",
    "spline_int = spline.antiderivative()\n",
    "infection_rate_change = (spline_int(x) - spline_int(lower)).reshape((upper, *vals.shape[1:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr['infection_rate_change'] = (('data_time', *tr.infection_change_polys_raw.transpose('infection_poly', ...).dims[1:]), infection_rate_change)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr.infection_rate_change.sel(country='Germany', chain=2).isel(draw=slice(-100, None)).plot.line(x='data_time', add_legend=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.lineplot(\n",
    "    x='data_time',\n",
    "    y='infection_rate',\n",
    "    hue='country',\n",
    "    units='draw',\n",
    "    estimator=None,\n",
    "    data = (\n",
    "        np.exp(tr.solution_log_infection_rate)\n",
    "        .rename('infection_rate')\n",
    "        #.sel(country=['Germany', 'Japan'], chain=2)\n",
    "        .sel(chain=2)\n",
    "        .isel(draw=slice(-100, None))\n",
    "        .to_dataframe()\n",
    "        .reset_index()\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(df_sign[df_sign.country == 'Germany'].set_index('date').confirmed).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(df_sign[df_sign.country == 'Italy'].set_index('date').confirmed).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr.infection_change_polys_sd.plot.line(x='draw', add_legend=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sign.query(\"country == 'Japan'\").set_index('date').confirmed.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "((1 - tr.solution_S) * xr.DataArray(countries['population'])).isel(country=0, draw=-1).plot.line(x='data_time', add_legend=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sign.query(\"country == 'Japan'\").set_index('date').eval('recovered + deaths').plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(tr.solution_R * xr.DataArray(countries['population'])).isel(country=0, draw=-1).plot.line(x='data_time', add_legend=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr.solution_S.diff('data_time').max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sign.plot.line(x='date', y='cases')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pymc3",
   "language": "python",
   "name": "pymc3"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}