Commit 107186d1 authored by Chris Jewell's avatar Chris Jewell
Browse files

Implemented weekday effect for I->R

CHANGES:

1. Implemented log-linear model for weekday effect in I->R hazard;
2. Updated summary.py and within_between.py to match;
3. Implemented linear random walk for new gamma0 parameter.
4. Implemented per-parameter data structure in output HDF5 file.
parent 7cf93af6
......@@ -161,7 +161,7 @@ def read_tier_restriction_data(
# Re-index
data.index = pd.MultiIndex.from_frame(data[["date", "lad19cd"]])
data = data[["tier_2", "tier_3"]]
data = data[["tier_2", "tier_3", "national_lockdown"]]
data = data[~data.index.duplicated()]
dates = pd.date_range(date_low, date_high - pd.Timedelta(1, "D"))
lad19cd = lad19cd_lookup["lad19cd"].sort_values().unique()
......
......@@ -2,9 +2,9 @@
# Enqueues COVID-19 pipelines
CASES_FILE="data/Anonymised Combined Line List 20201109.csv"
DATE_LOW="2020-08-14"
DATE_HIGH="2020-11-06"
CASES_FILE="data/Anonymised Combined Line List 20201117.csv"
DATE_LOW="2020-08-21"
DATE_HIGH="2020-11-13"
TEMPLATE_CONFIG=template_config.yaml
......
......@@ -5,7 +5,6 @@ import os
from time import perf_counter
import tqdm
import yaml
import h5py
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
......@@ -108,10 +107,11 @@ if __name__ == "__main__":
return model.log_prob(
dict(
beta2=block0[0],
gamma=block0[1],
gamma0=block0[1],
gamma1=block0[2],
beta1=block1[0],
beta3=block1[1:3],
xi=block1[3:],
beta3=[0.0, 0.0, 0.0], # block1[1:4],
xi=block1[1:], # block1[4:],
seir=events,
)
)
......@@ -134,7 +134,10 @@ if __name__ == "__main__":
* 1e-1,
covariance_burnin=200,
),
bijector=tfp.bijectors.Exp(),
bijector=tfp.bijectors.Blockwise(
bijectors=[tfp.bijectors.Exp(), tfp.bijectors.Identity()],
block_sizes=[1, 2],
),
name=name,
)
......@@ -211,13 +214,13 @@ if __name__ == "__main__":
results_dict = {}
res0 = results.inner_results
results_dict["theta"] = {
results_dict["block0"] = {
"is_accepted": res0[0].inner_results.is_accepted,
"target_log_prob": res0[
0
].inner_results.accepted_results.target_log_prob,
}
results_dict["xi"] = {
results_dict["block1"] = {
"is_accepted": res0[1].is_accepted,
"target_log_prob": res0[1].accepted_results.target_log_prob,
}
......@@ -254,8 +257,8 @@ if __name__ == "__main__":
gibbs_schema = GibbsKernel(
target_log_prob_fn=logp,
kernel_list=[
(0, make_blk0_kernel(init_state[0].shape, "theta")),
(1, make_blk1_kernel(init_state[1].shape, "xi")),
(0, make_blk0_kernel(init_state[0].shape, "block0")),
(1, make_blk1_kernel(init_state[1].shape, "block1")),
(2, make_event_multiscan_kernel),
],
name="gibbs0",
......@@ -287,10 +290,16 @@ if __name__ == "__main__":
tf.random.set_seed(2)
current_state = [
np.array([0.65, 0.48], dtype=DTYPE),
np.zeros(model.model["xi"](0.0).event_shape[-1] + 3, dtype=DTYPE),
np.array([0.65, 0.0, 0.0], dtype=DTYPE),
np.zeros(
model.model["xi"](0.0).event_shape[-1]
# + model.model["beta3"]().event_shape[-1]
+ 1,
dtype=DTYPE,
),
events,
]
print("Initial logpi:", logp(*current_state))
# Output file
samples, results, _ = sample(1, current_state)
......@@ -299,18 +308,23 @@ if __name__ == "__main__":
os.path.expandvars(config["output"]["results_dir"]),
config["output"]["posterior"],
),
sample_dict={"theta": (samples[0], (NUM_BURST_SAMPLES, 1)),
"xi": (samples[1], (NUM_BURST_SAMPLES, 1)),
"events": (samples[2], (NUM_BURST_SAMPLES, 64, 64, 1)),
},
sample_dict={
"beta2": (samples[0][:, 0], (NUM_BURST_SAMPLES,)),
"gamma0": (samples[0][:, 1], (NUM_BURST_SAMPLES,)),
"gamma1": (samples[0][:, 2], (NUM_BURST_SAMPLES,)),
"beta1": (samples[1][:, 0], (NUM_BURST_SAMPLES,)),
"xi": (
samples[1][:, 1:],
(NUM_BURST_SAMPLES, samples[1].shape[1] - 1),
),
"events": (samples[2], (NUM_BURST_SAMPLES, 64, 64, 1)),
},
results_dict=results,
num_samples=NUM_SAVED_SAMPLES,
)
posterior._file.create_dataset("initial_state", data=initial_state)
posterior._file.create_dataset("config", data=yaml.dump(config))
print("Initial logpi:", logp(*current_state))
# We loop over successive calls to sample because we have to dump results
# to disc, or else end OOM (even on a 32GB system).
# with tf.profiler.experimental.Profile("/tmp/tf_logdir"):
......@@ -329,7 +343,14 @@ if __name__ == "__main__":
start = perf_counter()
posterior.write_samples(
{"theta": samples[0], "xi": samples[1], "events": samples[2]},
{
"beta2": samples[0][:, 0],
"gamma0": samples[0][:, 1],
"gamma1": samples[0][:, 2],
"beta1": samples[1][:, 0],
"xi": samples[1][:, 1:],
"events": samples[2],
},
first_dim_offset=i * NUM_BURST_SAMPLES,
)
posterior.write_results(results, first_dim_offset=i * NUM_BURST_SAMPLES)
......@@ -339,13 +360,13 @@ if __name__ == "__main__":
print(
"Acceptance theta:",
tf.reduce_mean(
tf.cast(results["theta"]["is_accepted"], tf.float32)
tf.cast(results["block0"]["is_accepted"], tf.float32)
),
)
print(
"Acceptance xi:",
tf.reduce_mean(
tf.cast(results["theta"]["is_accepted"], tf.float32),
tf.cast(results["block1"]["is_accepted"], tf.float32),
),
)
print(
......@@ -374,9 +395,9 @@ if __name__ == "__main__":
)
print(
f"Acceptance theta: {posterior['results/theta/is_accepted'][:].mean()}"
f"Acceptance theta: {posterior['results/block0/is_accepted'][:].mean()}"
)
print(f"Acceptance xi: {posterior['results/xi/is_accepted'][:].mean()}")
print(f"Acceptance xi: {posterior['results/block1/is_accepted'][:].mean()}")
print(
f"Acceptance move S->E: {posterior['results/move/S->E/is_accepted'][:].mean()}"
)
......
"""Implements the COVID SEIR model as a TFP Joint Distribution"""
import pandas as pd
import geopandas as gp
import numpy as np
import tensorflow as tf
......@@ -40,11 +41,14 @@ def read_covariates(paths, date_low, date_high):
date_low,
date_high,
)
weekday = pd.date_range(date_low, date_high).weekday < 5
return dict(
C=mobility.to_numpy().astype(DTYPE),
W=commute_volume.to_numpy().astype(DTYPE),
N=popsize.to_numpy().astype(DTYPE),
L=tier_restriction.astype(DTYPE),
weekday=weekday.astype(DTYPE),
)
......@@ -88,7 +92,7 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
loc=tf.constant(0.0, dtype=DTYPE),
scale=tf.constant(100.0, dtype=DTYPE),
),
sample_shape=2,
sample_shape=covariates["L"].shape[-1],
)
def xi(beta1):
......@@ -102,24 +106,35 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
index_points=idx_pts[:, tf.newaxis],
)
def gamma():
return tfd.Gamma(
concentration=tf.constant(
priors["gamma"]["concentration"], dtype=DTYPE
),
rate=tf.constant(priors["gamma"]["rate"], dtype=DTYPE),
def gamma0():
# return tfd.Gamma(
# concentration=tf.constant(
# priors["gamma"]["concentration"], dtype=DTYPE
# ),
# rate=tf.constant(priors["gamma"]["rate"], dtype=DTYPE),
# )
return tfd.Normal(loc=tf.constant(0.0, dtype=DTYPE),
scale=tf.constant(100.0, dtype=DTYPE),
)
def gamma1():
return tfd.Normal(
loc=tf.constant(0.0, dtype=DTYPE),
scale=tf.constant(100.0, dtype=DTYPE),
)
def seir(beta2, beta3, xi, gamma):
def seir(beta2, beta3, xi, gamma0, gamma1):
beta2 = tf.convert_to_tensor(beta2, DTYPE)
beta3 = tf.convert_to_tensor(beta3, DTYPE)
xi = tf.convert_to_tensor(xi, DTYPE)
gamma = tf.convert_to_tensor(gamma, DTYPE)
gamma0 = tf.convert_to_tensor(gamma0, DTYPE)
gamma1 = tf.convert_to_tensor(gamma1, DTYPE)
L = tf.convert_to_tensor(covariates["L"], DTYPE)
L = L - tf.reduce_mean(L, axis=0)
weekday = tf.convert_to_tensor(covariates["weekday"], DTYPE)
weekday = weekday - tf.reduce_mean(weekday, axis=-1)
def transition_rate_fn(t, state):
C = tf.convert_to_tensor(covariates["C"], dtype=DTYPE)
C = tf.linalg.set_diag(
......@@ -140,6 +155,11 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
Lt = tf.gather(L, L_idx)
xB = tf.linalg.matvec(Lt, beta3)
weekday_idx = tf.clip_by_value(
tf.cast(t, tf.int64), 0, weekday.shape[0] - 1
)
weekday_t = tf.gather(weekday, weekday_idx)
infec_rate = tf.math.exp(xi_ + xB) * (
state[..., 2]
+ beta2
......@@ -154,7 +174,8 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
[NU], shape=[state.shape[0]]
) # Vector of length nc
ir = tf.broadcast_to(
[gamma], shape=[state.shape[0]]
[tf.math.exp(gamma0 + gamma1 * weekday_t)],
shape=[state.shape[0]],
) # Vector of length nc
return [infec_rate, ei, ir]
......@@ -170,7 +191,13 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
return tfd.JointDistributionNamed(
dict(
beta1=beta1, beta2=beta2, beta3=beta3, xi=xi, gamma=gamma, seir=seir
beta1=beta1,
beta2=beta2,
beta3=beta3,
xi=xi,
gamma0=gamma0,
gamma1=gamma1,
seir=seir,
)
)
......@@ -214,14 +241,14 @@ def next_generation_matrix_fn(covar_data, param):
beta = tf.math.exp(xi + xB)
ngm = beta[tf.newaxis, :] * (
ngm = beta[:, tf.newaxis] * (
tf.eye(C.shape[0], dtype=state.dtype)
+ param["beta2"] * commute_volume * C / N[tf.newaxis, :]
)
ngm = (
ngm
* state[..., 0][..., tf.newaxis]
/ (N[:, tf.newaxis] * param["gamma"])
/ (N[:, tf.newaxis] * tf.math.exp(param["gamma0"]))
)
return ngm
......
......@@ -339,7 +339,7 @@ tfp-nightly = "0.12.0.dev20201027"
type = "git"
url = "http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib.git"
reference = "master"
resolved_reference = "d4ac2adfacec9efcd326e6a7030ef560ef1d559e"
resolved_reference = "eddb85f7bfb3155777dc475eb467582f8caee9a5"
[[package]]
name = "geopandas"
......
......@@ -25,7 +25,7 @@ GIS_TEMPLATE = "data/UK2019mod_pop.gpkg"
# Reproduction number calculation
def calc_R_it(theta, xi, events, init_state, covar_data):
def calc_R_it(param, events, init_state, covar_data):
"""Calculates effective reproduction number for batches of metapopulations
:param theta: a tensor of batched theta parameters [B] + theta.shape
:param xi: a tensor of batched xi parameters [B] + xi.shape
......@@ -34,31 +34,41 @@ def calc_R_it(theta, xi, events, init_state, covar_data):
:param covar_data: the covariate data
:return a batched vector of R_it estimates
"""
print("Theta shape: ", theta.shape)
def r_fn(args):
theta_, xi_, events_ = args
beta1_, beta2_, xi_, gamma0_, events_ = args
t = events_.shape[-2] - 1
state = compute_state(init_state, events_, model_spec.STOICHIOMETRY)
state = tf.gather(state, t - 1, axis=-2) # State on final inference day
state = tf.gather(state, t, axis=-2) # State on final inference day
par = dict(
beta1=xi_[0],
beta2=theta_[0],
beta3=xi_[1:3],
gamma=theta_[1],
xi=xi_[3:],
beta1=beta1_,
beta2=beta2_,
beta3=tf.constant(
[0.0, 0.0, 0.0], dtype=model_spec.DTYPE
), # xi_[1:4],
gamma0=gamma0_,
xi=xi_,
)
ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par)
ngm = ngm_fn(t, state)
return ngm
return tf.vectorized_map(r_fn, elems=(theta, xi, events))
return tf.vectorized_map(
r_fn,
elems=(
param["beta1"],
param["beta2"],
param["xi"],
param["gamma0"],
events,
),
)
@tf.function
def predicted_incidence(theta, xi, init_state, init_step, num_steps, priors):
def predicted_incidence(param, init_state, init_step, num_steps, priors):
"""Runs the simulation forward in time from `init_state` at time `init_time`
for `num_steps`.
:param theta: a tensor of batched theta parameters [B] + theta.shape
......@@ -72,9 +82,16 @@ def predicted_incidence(theta, xi, init_state, init_step, num_steps, priors):
"""
def sim_fn(args):
theta_, xi_, init_ = args
beta1_, beta2_, xi_, gamma0_, gamma1_, init_ = args
par = dict(beta1=xi_[0], beta2=theta_[0], beta3=xi_[1:3], gamma=theta_[1], xi=xi_[3:])
par = dict(
beta1=beta1_,
beta2=beta2_,
beta3=[0.0, 0.0, 0.0],
gamma0=gamma0_,
gamma1=gamma1_,
xi=xi_,
)
model = model_spec.CovidUK(
covar_data,
......@@ -88,7 +105,14 @@ def predicted_incidence(theta, xi, init_state, init_step, num_steps, priors):
events = tf.map_fn(
sim_fn,
elems=(theta, xi, init_state),
elems=(
param["beta1"],
param["beta2"],
param["xi"],
param["gamma0"],
param["gamma1"],
init_state,
),
fn_output_signature=(tf.float64),
)
return events
......@@ -134,11 +158,13 @@ if __name__ == "__main__":
)
# Load posterior file
posterior_path = os.path.join(
config["output"]["results_dir"], config["output"]["posterior"]
)
print("Using posterior:", posterior_path)
posterior = h5py.File(
os.path.expandvars(
os.path.join(
config["output"]["results_dir"], config["output"]["posterior"]
)
posterior_path,
),
"r",
rdcc_nbytes=1024 ** 3,
......@@ -147,8 +173,13 @@ if __name__ == "__main__":
# Pre-determined thinning of posterior (better done in MCMC?)
idx = range(6000, 10000, 10)
theta = posterior["samples/theta"][idx]
xi = posterior["samples/xi"][idx]
param = dict(
beta1=posterior["samples/beta1"][idx],
beta2=posterior["samples/beta2"][idx],
xi=posterior["samples/xi"][idx],
gamma0=posterior["samples/gamma0"][idx],
gamma1=posterior["samples/gamma1"][idx],
)
events = posterior["samples/events"][idx]
init_state = posterior["initial_state"][:]
state_timeseries = compute_state(
......@@ -164,11 +195,13 @@ if __name__ == "__main__":
priors=config["mcmc"]["prior"],
)
ngms = calc_R_it(theta, xi, events, init_state, covar_data)
ngms = calc_R_it(param, events, init_state, covar_data)
b, _ = power_iteration(ngms)
rt = rayleigh_quotient(ngms, b)
q = np.arange(0.05, 1.0, 0.05)
rt_quantiles = pd.DataFrame({"Rt": np.quantile(rt, q)}, index=q).T.to_excel(
rt_quantiles = pd.DataFrame(
{"Rt": np.quantile(rt, q, axis=-1)}, index=q
).T.to_excel(
os.path.join(
config["output"]["results_dir"], config["output"]["national_rt"]
),
......@@ -178,8 +211,7 @@ if __name__ == "__main__":
# Note a 4 day recording lag in the case timeseries data requires that
# now = state_timeseries.shape[-2] + 4
prediction = predicted_incidence(
theta,
xi,
param,
init_state=state_timeseries[..., -1, :],
init_step=state_timeseries.shape[-2] - 1,
num_steps=70,
......@@ -233,7 +265,7 @@ if __name__ == "__main__":
ltla = gp.read_file(GIS_TEMPLATE, layer="UK2019mod_pop_xgen")
ltla = ltla[ltla["lad19cd"].str.startswith("E")] # England only, for now.
ltla = ltla.sort_values("lad19cd")
rti = tf.reduce_sum(ngms, axis=-1)
rti = tf.reduce_sum(ngms, axis=-2)
geosummary(
ltla,
......
......@@ -3,7 +3,7 @@
data:
mobility_matrix: data/mergedflows.csv
population_size: data/c2019modagepop.csv
commute_volume: data/201106_OFF_SEN_COVID19_road_traffic_national_table.xlsx
commute_volume: data/201113_OFF_SEN_COVID19_road_traffic_national_table.xlsx
reported_cases: data/Anonymised Combined Line List 20201109.csv
case_date_type: specimen
pillar: both
......@@ -35,7 +35,7 @@ settings:
- 2020-08-01
mcmc:
dmax: 21
dmax: 84
nmax: 50
m: 1
occult_nmax: 15
......
......@@ -13,7 +13,7 @@ from gemlib.util import compute_state
import model_spec
def make_within_rate_fns(covariates, theta, xi):
def make_within_rate_fns(covariates, beta2):
C = tf.convert_to_tensor(covariates["C"], dtype=model_spec.DTYPE)
C = tf.linalg.set_diag(
......@@ -26,39 +26,30 @@ def make_within_rate_fns(covariates, theta, xi):
tf.squeeze(covariates["N"]), dtype=model_spec.DTYPE
)
beta2 = theta[0]
def within_fn(t, state):
beta = tf.math.exp(xi)
rate = state[..., 2]
return rate
def between_fn(t, state):
w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, W.shape[0] - 1)
commute_volume = tf.gather(W, w_idx)
beta = tf.math.exp(xi)
rate = (
beta
* beta2
* commute_volume
* tf.linalg.matvec(C, state[..., 2] / N)
)
rate = beta2 * commute_volume * tf.linalg.matvec(C, state[..., 2] / N)
return rate
return within_fn, between_fn
# @tf.function
def calc_pressure_components(covariates, theta, xi, state):
def calc_pressure_components(covariates, beta2, state):
def atomic_fn(args):
theta_, xi_, state_ = args
within_fn, between_fn = make_within_rate_fns(covariates, theta_, xi_)
beta2_, state_ = args
within_fn, between_fn = make_within_rate_fns(covariates, beta2_)
within = within_fn(covariates["W"].shape[0], state_)
between = between_fn(covariates["W"].shape[0], state_)
total = within + between
return within / total, between / total
return tf.vectorized_map(atomic_fn, elems=(theta, xi, state))
return tf.vectorized_map(atomic_fn, elems=(beta2, state))
args = cli_args()
......@@ -89,14 +80,13 @@ posterior = h5py.File(
# Pre-determined thinning of posterior (better done in MCMC?)
idx = range(6000, 10000, 10)
theta = posterior["samples/theta"][idx]
xi = posterior["samples/xi"][idx]
beta2 = posterior["samples/beta2"][idx]
events = posterior["samples/events"][idx]
init_state = posterior["initial_state"][:]
state_timeseries = compute_state(init_state, events, model_spec.STOICHIOMETRY)
within, between = calc_pressure_components(
covar_data, theta, xi[:, -1], state_timeseries[..., -1, :]
covar_data, beta2, state_timeseries[..., -1, :]
)
gpkg = gp.read_file(
......
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