Commit 78a8c54b authored by Chris Jewell's avatar Chris Jewell
Browse files

COVID-19 UK space implementation with data ingestors for real data.

parent 557c256e
"""MCMC Test Rig for COVID-19 UK model"""
import optparse
# pylint: disable=E402
import argparse
import os
# Uncomment to block GPU use
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
from time import perf_counter
import tqdm
......@@ -10,18 +17,16 @@ import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from covid import config
from covid.impl.util import compute_state
from covid.impl.mcmc import UncalibratedLogRandomWalk, random_walk_mvnorm_fn
from covid.impl.event_time_mh import UncalibratedEventTimesUpdate
from covid.impl.occult_events_mh import UncalibratedOccultUpdate, TransitionTopology
from covid.impl.gibbs import DeterministicScanKernel, GibbsStep, flatten_results
from covid.impl.multi_scan_kernel import MultiScanKernel
from covid.data import read_phe_cases
import model_spec
# Uncomment to block GPU use
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
if tf.test.gpu_device_name():
print("Using GPU")
else:
......@@ -30,31 +35,42 @@ else:
tfd = tfp.distributions
tfb = tfp.bijectors
DTYPE = config.floatX
DTYPE = model_spec.DTYPE
if __name__ == "__main__":
# Read in settings
parser = optparse.OptionParser()
parser.add_option(
"--config",
parser = argparse.ArgumentParser()
parser.add_argument(
"-c",
dest="config",
"--config",
type=str,
default="example_config.yaml",
help="configuration file",
)
options, cmd_args = parser.parse_args()
print("Loading config file:", options.config)
args = parser.parse_args()
print("Loading config file:", args.config)
with open(options.config, "r") as f:
with open(args.config, "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
covar_data = model_spec.read_covariates(config["data"])
inference_period = [
np.datetime64(x) for x in config["settings"]["inference_period"]
]
covar_data = model_spec.read_covariates(
config["data"], date_low=inference_period[0], date_high=inference_period[1],
)
# We load in cases and impute missing infections first, since this sets the
# time epoch which we are analysing.
cases = model_spec.read_cases(config["data"]["reported_cases"])
cases = read_phe_cases(
config["data"]["reported_cases"],
date_low=inference_period[0],
date_high=inference_period[1],
date_type="report",
).astype(DTYPE)
# Impute censored events, return cases
events = model_spec.impute_censored_events(cases)
......@@ -67,7 +83,7 @@ if __name__ == "__main__":
# to set up a sensible initial state.
state = compute_state(
initial_state=tf.concat(
[covar_data["N"], tf.zeros_like(events[:, 0, :])], axis=-1
[covar_data["N"][:, tf.newaxis], tf.zeros_like(events[:, 0, :])], axis=-1
),
events=events,
stoichiometry=model_spec.STOICHIOMETRY,
......
......@@ -7,6 +7,7 @@ import tensorflow_probability as tfp
from covid.model import DiscreteTimeStateTransitionModel
from covid.util import impute_previous_cases
import covid.data as data
tfd = tfp.distributions
DTYPE = np.float64
......@@ -17,7 +18,7 @@ XI_FREQ = 14 # baseline transmission changes every 14 days
NU = tf.constant(0.5, dtype=DTYPE) # E->I rate assumed known.
def read_covariates(paths):
def read_covariates(paths, date_low, date_high):
"""Loads covariate data
:param paths: a dictionary of paths to data with keys {'mobility_matrix',
......@@ -25,9 +26,11 @@ def read_covariates(paths):
:returns: a dictionary of covariate information to be consumed by the model
{'C': commute_matrix, 'W': traffic_flow, 'N': population_size}
"""
mobility = pd.read_csv(paths["mobility_matrix"], index_col=0)
popsize = pd.read_csv(paths["population_size"], index_col=0)
commute_volume = pd.read_csv(paths["commute_volume"], index_col=0)
mobility = data.read_mobility(paths["mobility_matrix"])
popsize = data.read_population(paths["population_size"])
commute_volume = data.read_traffic_flow(
paths["commute_volume"], date_low=date_low, date_high=date_high
)
return dict(
C=mobility.to_numpy().astype(DTYPE),
......@@ -36,25 +39,18 @@ def read_covariates(paths):
)
def read_cases(path):
"""Loads case data from CSV file"""
cases_tidy = pd.read_csv(path)
cases_wide = cases_tidy.pivot(index="lad19cd", columns="date", values="cases")
return cases_wide
def impute_censored_events(cases):
"""Imputes censored S->E and E->I events using geometric
sampling algorithm in `impute_previous_cases`
There are application-specific magic numbers hard-coded below,
There are application-specific magic numbers hard-coded below,
which reflect experimentation to get the right lag between EI and
IR events, and SE and EI events respectively. These were chosen
by experimentation and examination of the resulting epidemic
by experimentation and examination of the resulting epidemic
trajectories.
:param cases: a MxT matrix of case numbers (I->R)
:returns: a MxTx3 tensor of events where the first two indices of
:returns: a MxTx3 tensor of events where the first two indices of
the right-most dimension contain the imputed event times.
"""
ei_events, lag_ei = impute_previous_cases(cases, 0.44)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment