diff --git a/2020-06-compartmental/.gitignore b/2020-06-compartmental/.gitignore new file mode 100644 index 0000000..2385b03 --- /dev/null +++ b/2020-06-compartmental/.gitignore @@ -0,0 +1,2 @@ +data/ +results/ diff --git a/2020-06-compartmental/Makefile b/2020-06-compartmental/Makefile new file mode 100644 index 0000000..912e74e --- /dev/null +++ b/2020-06-compartmental/Makefile @@ -0,0 +1,20 @@ +.PHONY: all lint clean mrclean short_uni_synth + +all: lint + +lint: FORCE + flake8 + +short_uni_synth: FORCE + python runner.py --experiment=short_uni_synth + +long_uni_synth: FORCE + python runner.py --experiment=long_uni_synth + +clean: FORCE + rm -rf temp logs errors + +mrclean: FORCE + rm -rf data results + +FORCE: diff --git a/2020-06-compartmental/analyze.ipynb b/2020-06-compartmental/analyze.ipynb new file mode 100644 index 0000000..ffd90b5 --- /dev/null +++ b/2020-06-compartmental/analyze.ipynb @@ -0,0 +1,190 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Analysis" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "import matplotlib.pyplot as plt\n", + "from runner import short_uni_synth, long_uni_synth\n", + "\n", + "%matplotlib inline\n", + "%config InlineBackend.figure_formats = ['svg']" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "results = list(short_uni_synth.results)\n", + "print(len(results))\n", + "for r in results:\n", + " if r[\"args\"].infer == \"mcmc\":\n", + " break\n", + "print(r.keys())\n", + "print(r[\"times\"].keys())\n", + "print(r[\"evaluate\"].keys())\n", + "print(r[\"evaluate\"][\"R0\"].keys())\n", + "print(r[\"infer\"].keys())\n", + "print(r[\"infer\"][\"R0\"].keys())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def plot_accuracy(variable, metric, experiment):\n", + " view = defaultdict(list)\n", + " for result in experiment.results:\n", + " args = result['args']\n", + " view[args.infer, args.num_bins, args.svi_steps].append(result)\n", + " markers = [\"o\", \"d\", \"s\", \"<\", \"v\", \"^\", \">\"]\n", + " assert len(view) <= len(markers)\n", + "\n", + " plt.figure(figsize=(6, 5)).patch.set_color(\"white\")\n", + " for (key, value), marker in zip(sorted(view.items()), markers):\n", + " algo, num_bins, svi_steps = key\n", + " if algo == \"svi\":\n", + " label = f\"SVI steps={svi_steps}\"\n", + " elif algo == \"mcmc\":\n", + " if num_bins == 1:\n", + " label = \"MCMC relaxed\"\n", + " else:\n", + " label = f\"MCMC num_bins={num_bins}\"\n", + " X = [v[\"times\"][\"infer\"] for v in value]\n", + " Y = [v[\"evaluate\"][variable][metric] for v in value]\n", + " plt.scatter(X, Y, marker=marker, label=label, alpha=0.8)\n", + " plt.ylim(0, None)\n", + " plt.xscale(\"log\")\n", + " plt.xlabel(\"inference time (sec)\")\n", + " plt.ylabel(metric.upper())\n", + " plt.title(f\"{variable} accuracy ({experiment.__name__})\")\n", + " plt.legend(loc=\"best\", prop={'size': 8})\n", + " plt.tight_layout()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_accuracy(\"R0\", \"crps\", short_uni_synth)\n", + "plot_accuracy(\"rho\", \"crps\", short_uni_synth)\n", + "plot_accuracy(\"obs\", \"crps\", short_uni_synth)\n", + "plot_accuracy(\"I\", \"crps\", short_uni_synth)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_accuracy(\"R0\", \"crps\", long_uni_synth)\n", + "plot_accuracy(\"rho\", \"crps\", long_uni_synth)\n", + "plot_accuracy(\"obs\", \"crps\", long_uni_synth)\n", + "plot_accuracy(\"I\", \"crps\", long_uni_synth)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def plot_convergence(variable, experiment, metrics=[\"n_eff\", \"r_hat\"]):\n", + " view = defaultdict(list)\n", + " for result in results:\n", + " args = result['args']\n", + " if args.infer == \"mcmc\":\n", + " view[args.num_bins].append(result)\n", + " markers = [\"o\", \"d\", \"s\"]\n", + " assert len(view) <= len(markers)\n", + "\n", + " fig, axes = plt.subplots(len(metrics), 1, figsize=(6, 5), sharex=True)\n", + " fig.patch.set_color(\"white\")\n", + " for (num_bins, value), marker in zip(sorted(view.items()), markers):\n", + " if num_bins == 1:\n", + " label = \"MCMC relaxed\"\n", + " else:\n", + " label = f\"MCMC num_bins={num_bins}\"\n", + " X = [v[\"times\"][\"infer\"] for v in value]\n", + " for metric, ax in zip(metrics, axes):\n", + " Y = [v[\"infer\"][variable][metric] for v in value]\n", + " ax.scatter(X, Y, marker=marker, label=label, alpha=0.8)\n", + " ax.set_xscale(\"log\")\n", + " ax.set_yscale(\"log\")\n", + " ax.set_ylabel(metric)\n", + " axes[0].set_title(f\"{variable} convergence ({experiment.__name__})\")\n", + " axes[1].set_ylim(1, None)\n", + " axes[-1].legend(loc=\"best\", prop={'size': 8})\n", + " axes[-1].set_xlabel(\"inference time (sec)\")\n", + " plt.subplots_adjust(hspace=0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_convergence(\"R0\", short_uni_synth)\n", + "plot_convergence(\"rho\", short_uni_synth)\n", + "plot_convergence(\"auxiliary_haar_split_0\", short_uni_synth)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_convergence(\"R0\", long_uni_synth)\n", + "plot_convergence(\"rho\", long_uni_synth)\n", + "plot_convergence(\"auxiliary_haar_split_0\", long_uni_synth)" + ] + }, + { + "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": 4 +} diff --git a/2020-06-compartmental/runner.py b/2020-06-compartmental/runner.py new file mode 100644 index 0000000..e1380e7 --- /dev/null +++ b/2020-06-compartmental/runner.py @@ -0,0 +1,123 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +import pickle +import subprocess +import sys +from importlib import import_module + +from util import get_filename + + +class Experiment: + """ + An experiment consists of a collection of tasks. + Each task generates a datapoint by running a python script. + Result datapoints are cached in pickle files named by fingerprint. + """ + def __init__(self, generate_tasks): + self.__name__ = generate_tasks.__name__ + self.tasks = [[sys.executable] + task for task in generate_tasks()] + self.files = [] + for task in self.tasks: + script = task[1] + parser = import_module(script.replace(".py", "")).Parser() + outfile = get_filename(script, parser.parse_args(task[2:])) + self.files.append(outfile) + + @property + def results(self): + """ + Iterates over the subset of experiment results that have been generated. + """ + for outfile in self.files: + if os.path.exists(outfile): + with open(outfile, "rb") as f: + result = pickle.load(f) + yield result + + +@Experiment +def short_uni_synth(): + base = [ + "uni_synth.py", + "--population=1000", + "--duration=20", "--forecast=10", + "--R0=3", "--incubation-time=2", "--recovery-time=4", + ] + for svi_steps in [1000, 2000, 5000, 10000]: + for rng_seed in range(10): + yield base + ["--svi", + "--num-samples=1000", + f"--svi-steps={svi_steps}", + f"--rng-seed={rng_seed}"] + for num_bins in [1, 2, 4]: + for num_samples in [200, 500, 1000]: + num_warmup = int(round(0.4 * num_samples)) + if num_bins == 1: + num_seeds = 10 + else: + num_seeds = 2 + for rng_seed in range(num_seeds): + yield base + ["--mcmc", + f"--warmup-steps={num_warmup}", + f"--num-samples={num_samples}", + f"--num-bins={num_bins}", + f"--rng-seed={rng_seed}"] + + +@Experiment +def long_uni_synth(): + base = [ + "uni_synth.py", + "--population=100000", + "--duration=100", "--forecast=30", + "--R0=2.5", "--incubation-time=4", "--recovery-time=10", + ] + for svi_steps in [1000, 2000, 5000, 10000]: + for rng_seed in range(10): + yield base + ["--svi", + "--num-samples=1000", + f"--svi-steps={svi_steps}", + f"--rng-seed={rng_seed}"] + for num_samples in [200, 500, 1000, 2000, 5000]: + num_warmup = int(round(0.4 * num_samples)) + for rng_seed in range(10): + yield base + ["--mcmc", + "--num-bins=1", + f"--warmup-steps={num_warmup}", + f"--num-samples={num_samples}", + f"--rng-seed={rng_seed}"] + for num_bins in [2, 4]: + for num_samples in [200, 500, 1000]: + num_warmup = int(round(0.4 * num_samples)) + for rng_seed in range(2): + yield base + ["--mcmc", + f"--warmup-steps={num_warmup}", + f"--num-samples={num_samples}", + f"--num-bins={num_bins}", + f"--rng-seed={rng_seed}"] + + +def main(args): + experiment = globals()[args.experiment] + for task, outfile in zip(experiment.tasks, experiment.files): + print(" \\\n ".join(task)) + if args.dry_run or os.path.exists(outfile): + continue + subprocess.check_call(task) + + print("-------------------------") + print("COMPLETED {} TASKS".format(len(experiment.tasks))) + print("-------------------------") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="experiment runner") + parser.add_argument("--experiment") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + main(args) diff --git a/2020-06-compartmental/setup.cfg b/2020-06-compartmental/setup.cfg new file mode 100644 index 0000000..370a859 --- /dev/null +++ b/2020-06-compartmental/setup.cfg @@ -0,0 +1,6 @@ +[flake8] +max-line-length = 120 + +[isort] +line_length = 120 +multi_line_output=3 diff --git a/2020-06-compartmental/uni_real.py b/2020-06-compartmental/uni_real.py new file mode 100644 index 0000000..622ddca --- /dev/null +++ b/2020-06-compartmental/uni_real.py @@ -0,0 +1,493 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import datetime +import logging +import math +import os +import pickle +import resource +import sys +import urllib.request +from collections import OrderedDict +from timeit import default_timer + +import numpy as np +import pandas as pd +import torch +from torch.nn.functional import pad + +import pyro +import pyro.distributions as dist +from pyro.contrib.epidemiology import CompartmentalModel, binomial_dist, infection_dist +from pyro.contrib.forecast.evaluate import eval_crps, eval_mae, eval_rmse +from pyro.infer.mcmc.util import summary +from pyro.ops.tensor_utils import convolve +from util import DATA, get_filename + +fmt = '%(process)d %(message)s' +logging.getLogger("pyro").handlers[0].setFormatter(logging.Formatter(fmt)) +logging.basicConfig(format=fmt, level=logging.INFO) + +# Misc California county populations as of early 2020. +counties = OrderedDict([ + # Bay Area counties. + ("Santa Clara", 1.763e6), + ("Alameda", 1.495e6), + ("Contra Costa", 1.038e6), + ("San Francisco", 871e3), + ("San Mateo", 712e3), + ("Sonoma", 479e3), + ("Solano", 412e3), + ("Marin", 251e3), + ("Napa", 135e3), + # Misc non Bay Area counties. + ("Los Angeles", 10.04e6), + ("Riverside", 2.471e6), + ("San Diego", 3.338e6), + ("Orange", 3.176e6), + ("San Bernardino", 2.18e6), + ("Imperial", 181e3), + ("Kern", 900e3), + ("Fresno", 999e3), + ("Tulare", 466e3), + ("Santa Barbara", 446e3), +]) +counties = OrderedDict((k, int(v)) for k, v in counties.items()) + + +def load_df(basename): + url = ("https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/" + "csse_covid_19_data/csse_covid_19_time_series/") + local_path = os.path.join(DATA, basename) + if not os.path.exists(local_path): + urllib.request.urlretrieve(url + basename, local_path) + return pd.read_csv(local_path) + + +def load_data(args): + cum_cases_df = load_df("time_series_covid19_confirmed_US.csv") + cum_deaths_df = load_df("time_series_covid19_deaths_US.csv") + + # Convert to torch.Tensor. + county = list(counties)[args.county] + population = counties[county] + i = list(cum_cases_df["Admin2"]).index(county) + cum_cases = cum_cases_df.iloc[i, 11:] + i = list(cum_deaths_df["Admin2"]).index(county) + cum_deaths = cum_deaths_df.iloc[i, 12:] + cum_cases = torch.tensor(cum_cases, dtype=torch.get_default_dtype()).contiguous() + cum_deaths = torch.tensor(cum_deaths, dtype=torch.get_default_dtype()).contiguous() + assert cum_cases.shape == cum_deaths.shape + start_date = datetime.datetime.strptime(cum_cases_df.columns[11], "%m/%d/%y") + + # Distribute reported cases and deaths among previous few days. + if args.report_lag: + kernel = torch.ones(args.report_lag) / args.report_lag + cum_cases = convolve(cum_cases, kernel, mode="valid").round() + cum_deaths = convolve(cum_deaths, kernel, mode="valid").round() + + # Convert from cumulative to difference data, and clamp to ensure positivity. + new_cases = (cum_cases - pad(cum_cases[:-1], (1, 0), value=0)).clamp(min=0) + new_deaths = (cum_deaths - pad(cum_deaths[:-1], (1, 0), value=0)).clamp(min=0) + + # Truncate. + truncate = (args.start_date - start_date).days + assert truncate > 0, "start date is too early" + new_cases = new_cases[truncate:].contiguous() + new_deaths = new_deaths[truncate:].contiguous() + start_date += datetime.timedelta(days=truncate) + logging.info(f"{county} data shape = {tuple(new_cases.shape)}") + + return {"population": population, + "new_cases": new_cases, + "new_deaths": new_deaths, + "start_date": start_date} + + +sigmoid = torch.distributions.transforms.SigmoidTransform() + + +class Model(CompartmentalModel): + def __init__(self, args, population, new_cases, new_deaths): + assert len(new_cases) == len(new_deaths) + duration = len(new_cases) + compartments = ("S", "E", "I") # R is implicit. + super().__init__(compartments, duration, population) + + self.incubation_time = args.incubation_time + self.recovery_time = args.recovery_time + self.new_cases = new_cases + self.new_deaths = new_deaths + + # Intervene via a step function. + t1 = (args.intervene_date - args.start_date).days + t2 = self.duration + args.forecast + self.intervene = torch.cat([torch.zeros(t1), torch.ones(t2 - t1)]) + + def global_model(self): + tau_e = self.incubation_time + tau_i = self.recovery_time + R0 = pyro.sample("R0", dist.LogNormal(1., 0.5)) # Weak prior. + external_rate = pyro.sample("external_rate", dist.LogNormal(-2, 2)) + rho0 = pyro.sample("rho0", dist.Beta(10, 10)) # About 50% response rate. + mu = pyro.sample("mu", dist.Beta(2, 100)) # About 2% mortality rate. + R_drift = pyro.sample("R_drift", dist.LogNormal(math.log(0.1), 1.)) + rho_drift = pyro.sample("rho_drift", dist.LogNormal(math.log(0.01), 1.)) + od = pyro.sample("od", dist.Beta(2, 6)) + return R0, external_rate, tau_e, tau_i, rho0, mu, R_drift, rho_drift, od + + def initialize(self, params): + R0, external_rate, tau_e, tau_i, rho0, mu, R_drift, rho_drift, od = params + + # Start with no local infections and initial Brownian motion. + return {"S": self.population, "E": 0, "I": 0, + "R_motion": sigmoid.inv(torch.tensor(0.95)), + "rho_motion": torch.tensor(0.)} + + def transition(self, params, state, t): + R0, external_rate, tau_e, tau_i, rho0, mu, R_drift, rho_drift, od = params + + # Assume drift is 4x larger after various interventions begin. + R_drift = R_drift * (0.25 + 0.75 * self.intervene[t]) + rho_drift = rho_drift * (0.25 + 0.75 * self.intervene[t]) + + # Assume reproductive number Rt and response rate rho vary in time. + R_motion = pyro.sample("R_motion_{}".format(t), + dist.Normal(state["R_motion"], R_drift)) + rho_motion = pyro.sample("rho_motion_{}".format(t), + dist.Normal(state["rho_motion"], rho_drift)) + Rt = pyro.deterministic("Rt_{}".format(t), + R0 * sigmoid(R_motion), event_dim=0) + rho = pyro.deterministic("rho_{}".format(t), + sigmoid(sigmoid.inv(rho0) + rho_motion), event_dim=0) + I_external = external_rate * tau_i / Rt + + # Sample flows between compartments. + S2E = pyro.sample("S2E_{}".format(t), + infection_dist(individual_rate=Rt / tau_i, + num_susceptible=state["S"], + num_infectious=state["I"] + I_external, + population=self.population, + overdispersion=od)) + E2I = pyro.sample("E2I_{}".format(t), + binomial_dist(state["E"], 1 / tau_e, overdispersion=od)) + I2R = pyro.sample("I2R_{}".format(t), + binomial_dist(state["I"], 1 / tau_i, overdispersion=od)) + + # Update compartments and heterogeneous variables. + state["S"] = state["S"] - S2E + state["E"] = state["E"] + S2E - E2I + state["I"] = state["I"] + E2I - I2R + state["R_motion"] = R_motion + state["rho_motion"] = rho_motion + + # Condition on observations. + t_is_observed = isinstance(t, slice) or t < self.duration + pyro.sample("new_cases_{}".format(t), + binomial_dist(S2E, rho, overdispersion=od), + obs=self.new_cases[t] if t_is_observed else None) + pyro.sample("new_deaths_{}".format(t), + binomial_dist(I2R, mu, overdispersion=od), + obs=self.new_deaths[t] if t_is_observed else None) + + +def _item(x): + if isinstance(x, torch.Tensor): + x = x.reshape(-1).median().item() + elif isinstance(x, dict): + for key, value in x.items(): + x[key] = _item(value) + return x + + +def infer_mcmc(args, model): + parallel = args.num_chains > 1 + + mcmc = model.fit_mcmc(heuristic_num_particles=args.smc_particles, + warmup_steps=args.warmup_steps, + num_samples=args.num_samples, + num_chains=args.num_chains, + mp_context="spawn" if parallel else None, + max_tree_depth=args.max_tree_depth, + num_quant_bins=args.num_bins, + haar=True, + haar_full_mass=args.haar_full_mass, + jit_compile=args.jit) + + result = summary(mcmc._samples) + result = _item(result) + return result + + +def infer_svi(args, model): + losses = model.fit_svi(heuristic_num_particles=args.smc_particles, + num_samples=args.num_samples, + num_steps=args.svi_steps, + num_particles=args.svi_particles, + guide_rank=args.guide_rank, + init_scale=args.init_scale, + learning_rate=args.learning_rate, + learning_rate_decay=args.learning_rate_decay, + betas=args.betas, + jit=args.jit) + + return {"loss_initial": losses[0], "loss_final": losses[-1]} + + +def predict(args, model, truth): + samples = model.predict(forecast=args.forecast) + if not args.plot: + return samples + + import matplotlib.pyplot as plt + import matplotlib.dates as mdates + fig, axes = plt.subplots(5, 1, figsize=(8, 8), sharex=True) + shelter_in_place = datetime.datetime.strptime("3/16/20", "%m/%d/%y") + axes[-1].text(shelter_in_place + datetime.timedelta(days=1), -0.25, + "shelter in place", horizontalalignment='center') + for ax in axes: + ax.axvline(shelter_in_place, color="black", linestyle=":", lw=1, alpha=0.3) + ax.axvline(truth["start_date"] + datetime.timedelta(days=model.duration), + color="black", lw=1, alpha=0.3) + axes[0].set_title("{}, population {}".format( + list(counties)[args.county], truth["population"])) + time = np.array([truth["start_date"] + datetime.timedelta(days=t) + for t in range(model.duration + args.forecast)]) + + # Plot forecasted series. + num_samples = samples["R0"].size(0) + for name, ax in zip(["new_cases", "new_deaths"], axes): + pred = samples[name][..., model.duration:] + median = pred.median(dim=0).values + p05 = pred.kthvalue(int(round(0.5 + 0.05 * num_samples)), dim=0).values + p95 = pred.kthvalue(int(round(0.5 + 0.95 * num_samples)), dim=0).values + ax.fill_between(time[model.duration:], p05, p95, color="red", alpha=0.3, + label="90% CI") + ax.plot(time[model.duration:], median, "r-", label="median") + ax.plot(time, truth[name], "k--", label="truth") + ax.set_yscale("log") + ax.set_ylim(0.5, None) + ax.set_ylabel(f"{name.replace('_', ' ')} / day") + ax.legend(loc="upper left") + + # Plot the latent time series. + ax = axes[2] + for name, color in zip(["E", "I"], ["red", "blue"]): + value = samples[name] + median = value.median(dim=0).values + p05 = value.kthvalue(int(round(0.5 + 0.05 * num_samples)), dim=0).values + p95 = value.kthvalue(int(round(0.5 + 0.95 * num_samples)), dim=0).values + ax.fill_between(time, p05, p95, color=color, alpha=0.3) + ax.plot(time, median, color=color, label=name, lw=1) + if name in truth: + ax.plot(time, truth[name], color=color, linestyle="--") + ax.set_yscale("log") + ax.set_ylim(0.5, None) + ax.set_ylabel("# people") + ax.legend(loc="best") + + # Plot parameter time series. + for name, ax in zip(["Rt", "rho"], axes[3:]): + value = samples[name] + median = value.median(dim=0).values + p05 = value.kthvalue(int(round(0.5 + 0.05 * num_samples)), dim=0).values + p95 = value.kthvalue(int(round(0.5 + 0.95 * num_samples)), dim=0).values + ax.fill_between(time, p05, p95, color="red", alpha=0.3, label="90% CI") + ax.plot(time, median, "r-", label="median") + if name in truth: + ax.plot(time, truth[name], "k--", label="truth") + ax.axhline(1, color="black", linestyle=":", lw=1, alpha=0.3) + ax.set_ylabel(name) + ax.legend(loc="best") + axes[3].set_ylim(0, None) + axes[4].set_ylim(0, 1) + + ax.set_xlim(time[0], time[-1]) + locator = mdates.AutoDateLocator(minticks=5, maxticks=15) + formatter = mdates.ConciseDateFormatter(locator) + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(formatter) + plt.tight_layout() + plt.subplots_adjust(hspace=0) + + return samples + + +def evaluate(args, truth, model, samples): + metrics = [("mae", eval_mae), ("rmse", eval_rmse), ("crps", eval_crps)] + result = {} + for key, pred in samples.items(): + if key in ("new_cases", "new_deaths"): + pred = pred[..., model.duration:] + + result[key] = {} + result[key]["mean"] = pred.mean().item() + result[key]["std"] = pred.std(dim=0).mean().item() + + if key in ("new_cases", "new_deaths"): + true = truth[key][..., model.duration:] + for metric, fn in metrics: + result[key][metric] = fn(pred, true) + + # Print estimated values. + covariates = [(name, value.squeeze()) + for name, value in sorted(samples.items()) + if value[0].numel() == 1] + for name, value in covariates: + mean = value.mean().item() + std = value.std().item() + if name in truth: + true = truth[name] + logging.info(f"{name}: true = {true:0.3g}, pred = {mean:0.3g} \u00B1 {std:0.3g}") + else: + logging.info(f"{name} = {mean:0.3g} \u00B1 {std:0.3g}") + + if args.plot and args.infer == "mcmc": + # Plot pairwise joint distributions for selected variables. + import matplotlib.pyplot as plt + N = len(covariates) + fig, axes = plt.subplots(N, N, figsize=(8, 8), sharex="col", sharey="row") + for i in range(N): + axes[i][0].set_ylabel(covariates[i][0]) + axes[0][i].set_xlabel(covariates[i][0]) + axes[0][i].xaxis.set_label_position("top") + for j in range(N): + ax = axes[i][j] + ax.set_xticks(()) + ax.set_yticks(()) + ax.scatter(covariates[j][1], -covariates[i][1], + lw=0, color="darkblue", alpha=0.3) + plt.tight_layout() + plt.subplots_adjust(wspace=0, hspace=0) + + return result + + +def main(args): + pyro.enable_validation(__debug__) + pyro.set_rng_seed(args.rng_seed + 20200619) + + result = {"file": __file__, "args": args, "argv": sys.argv} + + truth = load_data(args) + if args.generate: + # Generate data with same population and duration as real data. + model = Model(args, truth["population"], + [None] * len(truth["new_cases"]), + [None] * len(truth["new_cases"])) + truth.update(model.generate()) + logging.info("Synthetic: {}".format(", ".join( + f"{name}={value:0.3g}" + for name, value in sorted(truth.items()) + if isinstance(value, torch.Tensor) and value.numel() == 1))) + result["data"] = { + "population": truth["population"], + "total_cases": truth["new_cases"].sum().item(), + "total_deaths": truth["new_deaths"].sum().item(), + "max_cases": truth["new_cases"].max().item(), + "max_deaths": truth["new_deaths"].max().item(), + } + + t0 = default_timer() + + model = Model(args, truth["population"], + truth["new_cases"][:-args.forecast], + truth["new_deaths"][:-args.forecast]) + infer = {"mcmc": infer_mcmc, "svi": infer_svi}[args.infer] + result["infer"] = infer(args, model) + + t1 = default_timer() + + samples = predict(args, model, truth) + + t2 = default_timer() + + result["evaluate"] = evaluate(args, truth, model, samples) + result["times"] = {"infer": t1 - t0, "predict": t2 - t1} + result["rusage"] = resource.getrusage(resource.RUSAGE_SELF) + logging.info("DONE") + return result + + +class Parser(argparse.ArgumentParser): + def __init__(self): + super().__init__(description="CompartmentalModel experiments") + self.add_argument("--county", default=0, type=int) + self.add_argument("--generate", action="store_true") + self.add_argument("--start-date", default="2/1/20") + self.add_argument("--intervene-date", default="3/1/20") + self.add_argument("--report-lag", type=int, default=5) + self.add_argument("--forecast", default=14, type=int) + self.add_argument("--recovery-time", default=14.0, type=float) + self.add_argument("--incubation-time", default=5.5, type=float) + self.add_argument("--infer", default="svi") + self.add_argument("--mcmc", action="store_const", const="mcmc", dest="infer") + self.add_argument("--svi", action="store_const", const="svi", dest="infer") + self.add_argument("--haar-full-mass", default=10, type=int) + self.add_argument("--guide-rank", type=int) + self.add_argument("--init-scale", default=0.01, type=float) + self.add_argument("--num-samples", default=200, type=int) + self.add_argument("--smc-particles", default=1024, type=int) + self.add_argument("--svi-steps", default=5000, type=int) + self.add_argument("--svi-particles", default=32, type=int) + self.add_argument("--learning-rate", default=0.1, type=float) + self.add_argument("--learning-rate-decay", default=0.01, type=float) + self.add_argument("--betas", default="0.8,0.99") + self.add_argument("--warmup-steps", type=int) + self.add_argument("--num-chains", default=2, type=int) + self.add_argument("--max-tree-depth", default=5, type=int) + self.add_argument("--rng-seed", default=0, type=int) + self.add_argument("--num-bins", default=1, type=int) + self.add_argument("--double", action="store_true", default=True) + self.add_argument("--single", action="store_false", dest="double") + self.add_argument("--cuda", action="store_true") + self.add_argument("--jit", action="store_true", default=True) + self.add_argument("--nojit", action="store_false", dest="jit") + self.add_argument("--plot", action="store_true") + + def parse_args(self, *args, **kwargs): + args = super().parse_args(*args, **kwargs) + + assert args.forecast > 0 + + args.betas = tuple(map(float, args.betas.split(","))) + + # Parse dates. + for name, value in args.__dict__.items(): + if name.endswith("_date"): + value = datetime.datetime.strptime(value, "%m/%d/%y") + setattr(args, name, value) + + if args.warmup_steps is None: + args.warmup_steps = int(round(0.4 * args.num_samples)) + + if args.double: + if args.cuda: + torch.set_default_tensor_type(torch.cuda.DoubleTensor) + else: + torch.set_default_dtype(torch.float64) + elif args.cuda: + torch.set_default_tensor_type(torch.cuda.FloatTensor) + + return args + + +if __name__ == "__main__": + assert pyro.__version__.startswith('1.3.1') + args = Parser().parse_args() + + args.plot = True # DEBUG + if args.plot: + main(args) + import matplotlib.pyplot as plt + plt.show() + else: + # Cache output. + outfile = get_filename(__file__, args) + if not os.path.exists(outfile): + result = main(args) + with open(outfile, "wb") as f: + pickle.dump(result, f) + logging.info("Saved {}".format(outfile)) diff --git a/2020-06-compartmental/uni_synth.py b/2020-06-compartmental/uni_synth.py new file mode 100644 index 0000000..5bb3ccd --- /dev/null +++ b/2020-06-compartmental/uni_synth.py @@ -0,0 +1,244 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import logging +import math +import os +import pickle +import resource +import sys +from timeit import default_timer + +import torch + +import pyro +from pyro.contrib.epidemiology.models import ( + HeterogeneousSIRModel, + OverdispersedSEIRModel, + OverdispersedSIRModel, + SimpleSEIRModel, + SimpleSIRModel, + SuperspreadingSEIRModel, + SuperspreadingSIRModel +) +from pyro.contrib.forecast.evaluate import eval_crps, eval_mae, eval_rmse +from pyro.infer.mcmc.util import summary + +from util import get_filename + +fmt = '%(process)d %(message)s' +logging.getLogger("pyro").handlers[0].setFormatter(logging.Formatter(fmt)) +logging.basicConfig(format=fmt, level=logging.INFO) + + +def Model(args, data): + """Dispatch between different model classes.""" + if args.heterogeneous: + assert args.incubation_time == 0 + assert args.overdispersion == 0 + return HeterogeneousSIRModel(args.population, args.recovery_time, data) + elif args.incubation_time > 0: + assert args.incubation_time > 1 + if args.concentration < math.inf: + return SuperspreadingSEIRModel(args.population, args.incubation_time, + args.recovery_time, data) + elif args.overdispersion > 0: + return OverdispersedSEIRModel(args.population, args.incubation_time, + args.recovery_time, data) + else: + return SimpleSEIRModel(args.population, args.incubation_time, + args.recovery_time, data) + else: + if args.concentration < math.inf: + return SuperspreadingSIRModel(args.population, args.recovery_time, data) + elif args.overdispersion > 0: + return OverdispersedSIRModel(args.population, args.recovery_time, data) + else: + return SimpleSIRModel(args.population, args.recovery_time, data) + + +def generate_data(args): + extended_data = [None] * (args.duration + args.forecast) + model = Model(args, extended_data) + logging.info("Simulating from a {}".format(type(model).__name__)) + for attempt in range(100): + truth = model.generate({"R0": args.R0, + "rho": args.response_rate, + "k": args.concentration, + "od": args.overdispersion}) + obs = truth["obs"][:args.duration] + new_I = truth.get("S2I", truth.get("E2I")) + + obs_sum = int(obs.sum()) + new_I_sum = int(new_I[:args.duration].sum()) + assert 0 <= args.min_obs_portion < args.max_obs_portion <= 1 + min_obs = int(math.ceil(args.min_obs_portion * args.population)) + max_obs = int(math.floor(args.max_obs_portion * args.population)) + if min_obs <= obs_sum <= max_obs: + logging.info("Observed {:d}/{:d} infections:\n{}".format( + obs_sum, new_I_sum, " ".join(str(int(x)) for x in obs))) + return truth + + if obs_sum < min_obs: + raise ValueError("Failed to generate >={} observations. " + "Try decreasing --min-obs-portion (currently {})." + .format(min_obs, args.min_obs_portion)) + else: + raise ValueError("Failed to generate <={} observations. " + "Try increasing --max-obs-portion (currently {})." + .format(max_obs, args.max_obs_portion)) + + +def _item(x): + if isinstance(x, torch.Tensor): + x = x.reshape(-1).median().item() + elif isinstance(x, dict): + for key, value in x.items(): + x[key] = _item(value) + return x + + +def infer_mcmc(args, model): + parallel = args.num_chains > 1 + + mcmc = model.fit_mcmc(heuristic_num_particles=args.smc_particles, + warmup_steps=args.warmup_steps, + num_samples=args.num_samples, + num_chains=args.num_chains, + mp_context="spawn" if parallel else None, + max_tree_depth=args.max_tree_depth, + arrowhead_mass=args.arrowhead_mass, + num_quant_bins=args.num_bins, + haar=args.haar, + haar_full_mass=args.haar_full_mass, + jit_compile=args.jit) + + result = summary(mcmc._samples) + result = _item(result) + return result + + +def infer_svi(args, model): + losses = model.fit_svi(heuristic_num_particles=args.smc_particles, + num_samples=args.num_samples, + num_steps=args.svi_steps, + num_particles=args.svi_particles, + haar=args.haar, + jit=args.jit) + + return {"loss_initial": losses[0], "loss_final": losses[-1]} + + +def evaluate(args, truth, model, samples): + metrics = [("mae", eval_mae), ("rmse", eval_rmse), ("crps", eval_crps)] + result = {} + for key, pred in samples.items(): + if key == "obs": + pred = pred[..., args.duration:] + + result[key] = {} + result[key]["mean"] = pred.mean().item() + result[key]["std"] = pred.std(dim=0).mean().item() + + if key in truth: + true = truth[key] + if key == "obs": + true = true[..., args.duration:] + for metric, fn in metrics: + result[key][metric] = fn(pred, true) + + return result + + +def main(args): + pyro.enable_validation(__debug__) + pyro.set_rng_seed(args.rng_seed + 20200617) + + result = {"file": __file__, "args": args, "argv": sys.argv} + + truth = generate_data(args) + + t0 = default_timer() + + model = Model(args, data=truth["obs"][:args.duration]) + infer = {"mcmc": infer_mcmc, "svi": infer_svi}[args.infer] + result["infer"] = infer(args, model) + + t1 = default_timer() + + samples = model.predict(forecast=args.forecast) + + t2 = default_timer() + + result["evaluate"] = evaluate(args, truth, model, samples) + result["times"] = {"infer": t1 - t0, "predict": t2 - t1} + result["rusage"] = resource.getrusage(resource.RUSAGE_SELF) + logging.info("DONE") + return result + + +class Parser(argparse.ArgumentParser): + def __init__(self): + super().__init__(description="CompartmentalModel experiments") + self.add_argument("--population", default=1000, type=float) + self.add_argument("--min-obs-portion", default=0.1, type=float) + self.add_argument("--max-obs-portion", default=0.3, type=float) + self.add_argument("--duration", default=20, type=int) + self.add_argument("--forecast", default=10, type=int) + self.add_argument("--R0", default=1.5, type=float) + self.add_argument("--recovery-time", default=7.0, type=float) + self.add_argument("--incubation-time", default=0.0, type=float) + self.add_argument("--concentration", default=math.inf, type=float) + self.add_argument("--response-rate", default=0.5, type=float) + self.add_argument("--overdispersion", default=0., type=float) + self.add_argument("--heterogeneous", action="store_true") + self.add_argument("--infer", default="mcmc") + self.add_argument("--mcmc", action="store_const", const="mcmc", dest="infer") + self.add_argument("--svi", action="store_const", const="svi", dest="infer") + self.add_argument("--haar", action="store_true") + self.add_argument("--nohaar", action="store_const", const=False, dest="haar") + self.add_argument("--haar-full-mass", default=10, type=int) + self.add_argument("--num-samples", default=200, type=int) + self.add_argument("--smc-particles", default=1024, type=int) + self.add_argument("--svi-steps", default=5000, type=int) + self.add_argument("--svi-particles", default=32, type=int) + self.add_argument("--warmup-steps", type=int) + self.add_argument("--num-chains", default=2, type=int) + self.add_argument("--max-tree-depth", default=5, type=int) + self.add_argument("--arrowhead-mass", action="store_true") + self.add_argument("--rng-seed", default=0, type=int) + self.add_argument("--num-bins", default=1, type=int) + self.add_argument("--double", action="store_true", default=True) + self.add_argument("--single", action="store_false", dest="double") + self.add_argument("--cuda", action="store_true") + self.add_argument("--jit", action="store_true", default=True) + self.add_argument("--nojit", action="store_false", dest="jit") + self.add_argument("--verbose", action="store_true") + + def parse_args(self, *args, **kwargs): + args = super().parse_args(*args, **kwargs) + args.population = int(args.population) # to allow e.g. --population=1e6 + if args.warmup_steps is None: + args.warmup_steps = args.num_samples + if args.double: + if args.cuda: + torch.set_default_tensor_type(torch.cuda.DoubleTensor) + else: + torch.set_default_dtype(torch.float64) + elif args.cuda: + torch.set_default_tensor_type(torch.cuda.FloatTensor) + return args + + +if __name__ == "__main__": + assert pyro.__version__.startswith('1.3.1') + args = Parser().parse_args() + + # Cache output. + outfile = get_filename(__file__, args) + if not os.path.exists(outfile): + result = main(args) + with open(outfile, "wb") as f: + pickle.dump(result, f) + logging.info("Saved {}".format(outfile)) diff --git a/2020-06-compartmental/util.py b/2020-06-compartmental/util.py new file mode 100644 index 0000000..26e6add --- /dev/null +++ b/2020-06-compartmental/util.py @@ -0,0 +1,24 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +import os +from hashlib import sha1 + +ROOT = os.path.dirname(os.path.abspath(__file__)) +DATA = os.path.join(ROOT, "data") +RESULTS = os.path.join(ROOT, "results") + +# Ensure directories exist. +for path in [DATA, RESULTS]: + if not os.path.exists(path): + try: + os.makedirs(path) + except FileExistsError: + pass + + +def get_filename(script, args): + unique = script, sorted(args.__dict__.items()) + fingerprint = sha1(str(unique).encode()).hexdigest() + cachefile = os.path.join(RESULTS, fingerprint + ".pkl") + return cachefile