sim/exe-inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "if '..' not in sys.path:\n",
    "    sys.path.append('..')\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import copy\n",
    "import scipy as sp\n",
    "import math\n",
    "import seaborn\n",
    "import pickle\n",
    "import warnings\n",
    "import matplotlib\n",
    "import re\n",
    "import multiprocessing\n",
    "\n",
    "from lib.mobilitysim import MobilitySimulator\n",
    "from lib.dynamics import DiseaseModel\n",
    "from lib.inference import * \n",
    "from bayes_opt import BayesianOptimization\n",
    "from lib.parallel import *\n",
    "from lib.distributions import CovidDistributions\n",
    "from lib.plot import Plotter\n",
    "from lib.data import collect_data_from_df\n",
    "from lib.measures import (\n",
    "    MeasureList, \n",
    "    BetaMultiplierMeasure, \n",
    "    BetaMultiplierMeasureByType,\n",
    "    SocialDistancingForAllMeasure, \n",
    "    SocialDistancingByAgeMeasure,\n",
    "    SocialDistancingForPositiveMeasure, \n",
    "    Interval)\n",
    "\n",
    "from lib.mobilitysim import MobilitySimulator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Settings\n",
    "\n",
    "Determine settings for inference. Nothing below should have to be changed signficiantly/at all."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# settings used to generate mobility traces on the fly. These settings are used for inference\n",
    "# See town-generator.ipynb for an example on how to create\n",
    "mob_settings = 'lib/tu_settings_20_10.pk'\n",
    "case_downsample = 10\n",
    "\n",
    "# optimizer settings (toy settings)\n",
    "multimodal = False\n",
    "random_repeats_per_call = 4\n",
    "init_points = 1\n",
    "n_iterations = 1\n",
    "cpus_used = multiprocessing.cpu_count() // 2\n",
    "\n",
    "# optimizer bounds\n",
    "param_bounds = {\n",
    "    'beta':  (1.00, 1.80),\n",
    "    'alpha': (0.20, 0.80),\n",
    "    'mu': (0.20, 0.80)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# seed for random states and log file\n",
    "c = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Covid19 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_cases_ = collect_data_from_df('LK Tübingen', 'new')\n",
    "resistant_cases_ = collect_data_from_df('LK Tübingen', 'recovered')\n",
    "fatality_cases_ = collect_data_from_df('LK Tübingen', 'fatality')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Empirical fatality rate per age group from the above data. RKI data defines 6 groups: **0-4y, 5-14y, 15-34y, 35-59y, 60-79y, 80+y**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fatality rate per age group\n",
    "num_age_groups = fatality_cases_.shape[1] \n",
    "fatality_rates_by_age = (fatality_cases_[-1, :] / \\\n",
    "    (new_cases_[-1, :] +  fatality_cases_[-1, :] + resistant_cases_[-1, :]))\n",
    "\n",
    "print('Empirical fatality rates per age group:  ', fatality_rates_by_age.tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scale down cases based on number of people in simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_cases, resistant_cases, fatality_cases = (\n",
    "    1/case_downsample * new_cases_, \n",
    "    1/case_downsample * resistant_cases_, \n",
    "    1/case_downsample * fatality_cases_)\n",
    "new_cases, resistant_cases, fatality_cases = np.ceil(new_cases), np.ceil(resistant_cases), np.ceil(fatality_cases)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Maximum time fixed by real data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_time = int(new_cases.shape[0] * 24.0) # maximum time to simulate, in hours\n",
    "print('Max time T (days):', max_time // 24)\n",
    "print('Positive at t=0:  ', int(new_cases[0, :].sum()))\n",
    "print('Positive at t=T:  ', int(new_cases[-1, :].sum()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define distributions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate correct distributions\n",
    "distributions = CovidDistributions(fatality_rates_by_age=fatality_rates_by_age)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define model parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set simulation parameters and containment strategy for positives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set initial seed count (based on infection counts on March 10)\n",
    "initial_seeds = {\n",
    "    'expo' : 1,\n",
    "    'ipre' : 1,\n",
    "    'isym' : 3,\n",
    "    'iasy' : 3,\n",
    "}\n",
    "\n",
    "# standard quarantine of positive tests and test availablility\n",
    "measure_list = MeasureList([\n",
    "    SocialDistancingForPositiveMeasure(\n",
    "        t_window=Interval(0.0, max_time), p_stay_home=1.0)\n",
    "])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set testing parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testing_params = {\n",
    "    'testing_t_window'    : [0.0, max_time], # in hours\n",
    "    'testing_frequency'   : 24.0,     # in hours\n",
    "    'test_reporting_lag'  : 48.0,     # in hours (actual and self-report delay)\n",
    "    'tests_per_batch'     : 10,       # assume 300 tests/day in LK Tübingen\n",
    "    'test_smart_delta'    : 24.0 * 3, # in hours\n",
    "    'test_smart_duration' : 24.0 * 7, # in hours\n",
    "    'test_smart_action'   : 'isolate', \n",
    "    'test_smart_num_contacts'   : 10, \n",
    "    'test_targets'        : 'isym',\n",
    "    'test_queue_policy'   : 'fifo',\n",
    "    'smart_tracing'       : None, \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run optimization "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load settings as set in header of this notebook and generate example traces to extract information for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(mob_settings, 'rb') as fp:\n",
    "    obj = pickle.load(fp)\n",
    "mob = MobilitySimulator(**obj)\n",
    "%time mob.simulate(max_time=max_time, seed=12345)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initiate objective. `loss_daily` considers aggregate positives over days of simulation, `multimodal_loss_daily` considers a weighted combination aggregate positive, recovered, deceased people over days of simulation. Using the latter for inference is potential future work. Both compute the mean squared error of target and predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initiate objective\n",
    "if multimodal:\n",
    "    # targets\n",
    "    multimodal_targets = (\n",
    "        new_cases.sum(axis=1), \n",
    "        resistant_cases.sum(axis=1), \n",
    "        fatality_cases.sum(axis=1)) \n",
    "\n",
    "    # generate objective function\n",
    "    extra_params = {'weights' : [1.0, 1.0, 1.0]}\n",
    "    loss_function = make_loss_function(\n",
    "        mob_settings=mob_settings,\n",
    "        distributions=distributions,\n",
    "        initial_seeds=initial_seeds, \n",
    "        testing_params=testing_params,\n",
    "        random_repeats=random_repeats_per_call,  \n",
    "        cpu_count=cpus_used, \n",
    "        measure_list=measure_list, \n",
    "        loss='multimodal_loss_daily', \n",
    "        targets=multimodal_targets, \n",
    "        param_bounds=param_bounds,\n",
    "        num_site_types=mob.num_site_types,\n",
    "        num_sites=mob.num_sites,\n",
    "        num_people=mob.num_people,\n",
    "        site_loc=mob.site_loc,\n",
    "        home_loc=mob.home_loc,\n",
    "        time_horizon=mob.max_time,\n",
    "        c=c,\n",
    "        extra_params=extra_params)\n",
    "else:\n",
    "    # targets\n",
    "    targets = new_cases.sum(axis=1)\n",
    "\n",
    "    # generate objective function\n",
    "    loss_function = make_loss_function( \n",
    "        mob_settings=mob_settings,\n",
    "        distributions=distributions,\n",
    "        initial_seeds=initial_seeds, \n",
    "        testing_params=testing_params,\n",
    "        random_repeats=random_repeats_per_call,  \n",
    "        cpu_count=cpus_used, \n",
    "        measure_list=measure_list, \n",
    "        loss='loss_daily', \n",
    "        targets=targets, \n",
    "        param_bounds=param_bounds,\n",
    "        num_site_types=mob.num_site_types,\n",
    "        num_sites=mob.num_sites,\n",
    "        num_people=mob.num_people,\n",
    "        site_loc=mob.site_loc,\n",
    "        home_loc=mob.home_loc,\n",
    "        time_horizon=mob.max_time, \n",
    "        c=c,\n",
    "        extra_params=None)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run Bayesian optimization with the upper confidence bound heuristic. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "print('\\nBayesian optimization:')\n",
    "print('  Init points: ', init_points)\n",
    "print('  Iterations : ', n_iterations)\n",
    "print('  Using cpus:  ', cpus_used)\n",
    "print('  Parameter bounds: ', param_bounds, '\\n')\n",
    "\n",
    "# run optimization\n",
    "optimizer = BayesianOptimization(f=loss_function, pbounds=param_bounds, verbose=2, random_state=c)\n",
    "optimizer.maximize(init_points=init_points, n_iter=n_iterations,\n",
    "    acq=\"ucb\", kappa=1.0, # upper confidence bound heuristic\n",
    "#     acq=\"ei\", xi=0.05, # expected improvement heuristic\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract inferred parameters and store to file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Inferred parameters\n",
    "print('Inferred parameters')\n",
    "print(optimizer.max['params'])\n",
    "\n",
    "\n",
    "filename = f'optimized_params_{c}.txt'\n",
    "with open(filename, 'w+') as fp:\n",
    "    fp.write(f'Optimized parameters run={c}\\n' + str(optimizer.max['params']) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting run for the tuned parameters\n",
    "The optimization can also be skipped and settings simply set here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "with open(mob_settings, 'rb') as fp:\n",
    "        kwargs = pickle.load(fp)\n",
    "mob = MobilitySimulator(**kwargs)\n",
    "mob.verbose = True\n",
    "mob.simulate(max_time=max_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "sim = DiseaseModel(mob, distributions)\n",
    "\n",
    "sim.launch_epidemic(\n",
    "    params={\n",
    "        'betas' : [1.1383] * 5, \n",
    "        'alpha' : 0.3224,\n",
    "        'mu' : 0.2072\n",
    "    },\n",
    "    initial_counts=initial_seeds,\n",
    "    testing_params=testing_params,\n",
    "    measure_list=measure_list,\n",
    "    verbose=True)\n",
    "\n",
    "print('\\nInfections at end: ', \n",
    "      sim.state['isym'].sum() + sim.state['ipre'].sum() + sim.state['iasy'].sum())\n",
    "\n",
    "print('Positives at end: ', \n",
    "      sim.state['posi'].sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# plotting run for inferred parameters\n",
    "random_repeats = 16\n",
    "\n",
    "# inferred parameters (see paper)\n",
    "params = {\n",
    "    'betas' : [1.1383] * 5, # site infectivity by type\n",
    "    'alpha' : 0.3224,\n",
    "    'mu' : 0.2072\n",
    "}\n",
    "\n",
    "summary = launch_parallel_simulations(\n",
    "    mob_settings, \n",
    "    distributions, \n",
    "    random_repeats, multiprocessing.cpu_count(), \n",
    "    params, initial_seeds, testing_params, measure_list, \n",
    "    max_time=mob.max_time, \n",
    "    num_people=mob.num_people, \n",
    "    num_sites=mob.num_sites, \n",
    "    site_loc=mob.site_loc, \n",
    "    home_loc=mob.home_loc)\n",
    "\n",
    "\n",
    "print('\\nInfections at end: ', \n",
    "      summary.state['isym'].sum(axis=1) + summary.state['ipre'].sum(axis=1) + summary.state['iasy'].sum(axis=1))\n",
    "\n",
    "print('Positives at end: ', \n",
    "      summary.state['posi'].sum(axis=1))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the positives over time from the simulations compared to the real case data, adjusted for the reporting lag, as are the real cases, to validate the inference procedure "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plotter = Plotter()\n",
    "plotter.plot_positives_vs_target(\n",
    "    summary, new_cases.sum(axis=1), \n",
    "    test_lag=2,\n",
    "    title='Inferred parameters and COVID-19 case data', \n",
    "    filename='inference_example',\n",
    "    figsize=(10, 5),\n",
    "    errorevery=1, acc=1000, \n",
    "    ymax=80)"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}