Commit 264cea5c authored by Chris Jewell's avatar Chris Jewell
Browse files

Bugfix to naming in thin_samples.pkl

CHANGES:

1. `thin_samples['init_state']` --> `thin_samples['initial_state']`
2. `thin_samples['events']` --> `thin_samples['seir']`
3. add time dimension to next generation matrix
4. update downstream scripts to changes.
parent 37b30227
......@@ -276,10 +276,9 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
"gamma0": samples[0][:, 1],
"gamma1": samples[0][:, 2],
"sigma": samples[0][:, 3],
"beta3": tf.zeros([1, 5], dtype=DTYPE),
"beta1": samples[1][:, 0],
"xi": samples[1][:, 1:],
"events": samples[2],
"seir": samples[2],
},
results_dict=results,
num_samples=NUM_SAVED_SAMPLES,
......@@ -309,12 +308,9 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
"gamma0": samples[0][:, 1],
"gamma1": samples[0][:, 2],
"sigma": samples[0][:, 3],
"beta3": tf.zeros(
[samples[0].shape[0], 5], dtype=DTYPE
), # samples[0][:, 4:],
"beta1": samples[1][:, 0],
"xi": samples[1][:, 1:],
"events": samples[2],
"seir": samples[2],
},
first_dim_offset=i * NUM_BURST_SAMPLES,
)
......
......@@ -5,13 +5,12 @@ import numpy as np
import xarray
import tensorflow as tf
from covid import model_spec
from covid.util import copy_nc_attrs
from gemlib.util import compute_state
def calc_posterior_ngm(samples, covar_data):
def calc_posterior_ngm(samples, initial_state, times, 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
......@@ -20,21 +19,26 @@ def calc_posterior_ngm(samples, covar_data):
:param covar_data: the covariate data
:return a batched vector of R_it estimates
"""
times = tf.convert_to_tensor(times)
def r_fn(args):
par = tf.nest.pack_sequence_as(samples, args)
t = par['seir'].shape[-2] - 1
state = compute_state(
samples["init_state"], par['seir'], model_spec.STOICHIOMETRY
initial_state, par["seir"], model_spec.STOICHIOMETRY
)
state = tf.gather(state, t, axis=-2) # State on final inference day
del par["seir"]
del par['seir']
ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par)
ngm = ngm_fn(t, state)
return ngm
def fn(t):
state_ = tf.gather(
state, t, axis=-2
) # State on final inference day
ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par)
ngm = ngm_fn(t, state_)
return ngm
return tf.vectorized_map(fn, elems=times)
return tf.vectorized_map(
r_fn,
......@@ -49,16 +53,24 @@ def next_generation_matrix(input_files, output_file):
with open(input_files[1], "rb") as f:
samples = pkl.load(f)
initial_state = samples["initial_state"]
del samples["initial_state"]
times = [
samples["seir"].shape[-2] - 1,
]
# Compute ngm posterior
ngm = calc_posterior_ngm(samples, covar_data).numpy()
ngm = calc_posterior_ngm(samples, initial_state, times, covar_data)
ngm = xarray.DataArray(
ngm,
coords=[
np.arange(ngm.shape[0]),
covar_data.coords["time"][times],
covar_data.coords["location"],
covar_data.coords["location"],
],
dims=["iteration", "dest", "src"],
dims=["iteration", "time", "dest", "src"],
)
ngm = xarray.Dataset({"ngm": ngm})
......@@ -73,21 +85,19 @@ if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"-s",
"--samples",
"samples",
type=str,
description="A pickle file with MCMC samples",
required=True,
help="A pickle file with MCMC samples",
)
parser.add_argument(
"-d",
"--data",
type=str,
decription="A data glob pickle file",
require=True,
help="A data glob pickle file",
required=True,
)
parser.add_argument(
"-o", "--output", type=str, description="The output file", require=True
"-o", "--output", type=str, help="The output file", required=True
)
args = parser.parse_args()
......
......@@ -15,6 +15,7 @@ def overall_rt(next_generation_matrix, output_file):
ngms = xarray.open_dataset(
next_generation_matrix, group="posterior_predictive"
)["ngm"]
ngms = ngms[:, 0, :, :].drop("time")
b, _ = power_iteration(ngms)
rt = rayleigh_quotient(ngms, b)
q = np.arange(0.05, 1.0, 0.05)
......@@ -30,10 +31,11 @@ if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"input_file",
description="The input .pkl file containing the next generation matrix",
type=str,
help="The input .pkl file containing the next generation matrix",
)
parser.add_argument(
"output_file", description="The name of the output .xlsx file"
"output_file", type=str, help="The name of the output .xlsx file"
)
args = parser.parse_args()
......
......@@ -9,8 +9,10 @@ from covid import model_spec
from covid.util import copy_nc_attrs
from gemlib.util import compute_state
@tf.function
def predicted_incidence(posterior_samples, covar_data, init_step, num_steps):
def predicted_incidence(
posterior_samples, init_state, covar_data, init_step, num_steps
):
"""Runs the simulation forward in time from `init_state` at time `init_time`
for `num_steps`.
:param param: a dictionary of model parameters
......@@ -22,33 +24,36 @@ def predicted_incidence(posterior_samples, covar_data, init_step, num_steps):
"""
posterior_state = compute_state(
posterior_samples["init_state"],
init_state,
posterior_samples["seir"],
model_spec.STOICHIOMETRY,
)
posterior_samples['init_state_'] = posterior_state[..., init_step, :]
del posterior_samples['seir']
def sim_fn(args):
par = tf.nest.pack_sequence_as(posterior_samples, args)
init_ = par['init_state_']
del par['init_state_']
model = model_spec.CovidUK(
covar_data,
initial_state=init_,
initial_step=init_step,
num_steps=num_steps,
posterior_samples["new_init_state"] = posterior_state[..., init_step, :]
del posterior_samples["seir"]
@tf.function
def do_sim():
def sim_fn(args):
par = tf.nest.pack_sequence_as(posterior_samples, args)
init_ = par["new_init_state"]
del par["new_init_state"]
model = model_spec.CovidUK(
covar_data,
initial_state=init_,
initial_step=init_step,
num_steps=num_steps,
)
sim = model.sample(**par)
return sim["seir"]
return tf.map_fn(
sim_fn,
elems=tf.nest.flatten(posterior_samples),
fn_output_signature=(tf.float64),
)
sim = model.sample(**par)
return sim["seir"]
events = tf.map_fn(
sim_fn,
elems=tf.nest.flatten(posterior_samples),
fn_output_signature=(tf.float64),
)
return init_state, events
return posterior_samples["new_init_state"], do_sim()
def read_pkl(filename):
......@@ -60,7 +65,10 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
covar_data = xarray.open_dataset(data, group="constant_data")
cases = xarray.open_dataset(data, group="observations")
samples = read_pkl(posterior_samples)
initial_state = samples["initial_state"]
del samples["initial_state"]
if initial_step < 0:
initial_step = samples["seir"].shape[-2] + initial_step
......@@ -73,7 +81,7 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
)
estimated_init_state, predicted_events = predicted_incidence(
samples, covar_data, initial_step, num_steps
samples, initial_state, covar_data, initial_step, num_steps
)
prediction = xarray.DataArray(
......@@ -109,23 +117,21 @@ if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"-i", "--initial-step", type=int, default=0, description="Initial step"
)
parser.add_argument(
"-n", "--num-steps", type=int, default=1, description="Number of steps"
"-i", "--initial-step", type=int, default=0, help="Initial step"
)
parser.add_argument(
"data_pkl", type=str, description="Covariate data pickle"
"-n", "--num-steps", type=int, default=1, help="Number of steps"
)
parser.add_argument("data_pkl", type=str, help="Covariate data pickle")
parser.add_argument(
"posterior_samples_pkl",
type=str,
description="Posterior samples pickle",
help="Posterior samples pickle",
)
parser.add_argument(
"output_file",
type=str,
description="Output pkl file",
help="Output pkl file",
)
args = parser.parse_args()
......
......@@ -22,7 +22,7 @@ def rt(input_file, output_file):
ngm = xarray.open_dataset(input_file, group="posterior_predictive")["ngm"]
rt = np.sum(ngm, axis=-2)
rt = ngm.sum(dim="dest").isel(time=-1).drop("time")
rt_summary = mean_and_ci(rt, name="Rt")
exceed = np.mean(rt > 1.0, axis=0)
......
......@@ -6,24 +6,12 @@ import pickle as pkl
def thin_posterior(input_file, output_file, config):
thin_idx = range(config['start'], config['end'], config['by'])
idx = slice(config["start"], config["end"], config["by"])
f = h5py.File(input_file, "r", rdcc_nbytes=1024 ** 3, rdcc_nslots=1e6)
output_dict = dict(
beta1=f["samples/beta1"][thin_idx],
beta2=f["samples/beta2"][thin_idx],
beta3=f["samples/beta3"][
thin_idx,
],
sigma=f["samples/sigma"][
thin_idx,
],
xi=f["samples/xi"][thin_idx],
gamma0=f["samples/gamma0"][thin_idx],
gamma1=f["samples/gamma1"][thin_idx],
seir=f["samples/events"][thin_idx],
init_state=f["initial_state"][:],
)
output_dict = {k: v[idx] for k, v in f["samples"].items()}
output_dict["initial_state"] = f["initial_state"][:]
f.close()
with open(output_file, "wb") as f:
......@@ -36,26 +24,20 @@ def thin_posterior(input_file, output_file, config):
if __name__ == "__main__":
import yaml
import os
from covid.cli_arg_parse import cli_args
args = cli_args()
with open(args.config, "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
import argparse
def join_and_expand(prefix, filename):
return os.path.expandvars(os.path.join(prefix, filename))
input_file = join_and_expand(
config["output"]["results_dir"], config["output"]["posterior"]
parser = argparse.ArgumentParser()
parser.add_argument(
"-c", "--config", type=str, help="Configuration file", required=True
)
output_file = join_and_expand(
config["output"]["results_dir"], config["ThinPosterior"]["output_file"]
parser.add_argument(
"-o", "--output", type=str, help="Output pkl file", required=True
)
parser.add_argument("samples", type=str, help="Posterior HDF5 file")
args = parser.parse_args()
thin_idx = range(
config["ThinPosterior"]["thin_start"],
config["ThinPosterior"]["thin_end"],
config["ThinPosterior"]["thin_by"],
)
thin_posterior(input_file, output_file, thin_idx)
with open(args.config, "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
print("Config: ", config["ThinPosterior"])
thin_posterior(args.samples, args.output, config["ThinPosterior"])
......@@ -70,7 +70,7 @@ def within_between(input_files, output_file):
beta2 = samples["beta2"]
events = samples["seir"]
init_state = samples["init_state"]
init_state = samples["initial_state"]
state_timeseries = compute_state(
init_state, events, model_spec.STOICHIOMETRY
)
......
[tool.poetry]
name = "covid19uk"
version = "0.5.0"
version = "0.5.5"
description = "Spatial stochastic SEIR analysis of COVID-19 in the UK"
authors = ["Chris Jewell <c.jewell@lancaster.ac.uk>"]
license = "MIT"
......@@ -20,7 +20,7 @@ xlrd = "^1.2.0"
tqdm = "^4.50.2"
h5py = "^2.10.0"
gemlib = {git = "http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib.git"}
xarray = "^0.16.1"
xarray = {extras = ["netcdf4"], version = "^0.17.0"}
seaborn = "^0.11.0"
ruffus = "^2.8.4"
tensorflow = "^2.4.0"
......
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