Commit 7883c91b authored by Chris Jewell's avatar Chris Jewell
Browse files

Added LAD-level random effect

parent ffd59342
......@@ -56,7 +56,7 @@ def _get_window_sizes(num_adaptation_steps):
return first_window_size, slow_window_size, last_window_size
@tf.function(jit_compile=False)
@tf.function(autograph=False, jit_compile=False)
def _fast_adapt_window(
num_draws,
joint_log_prob_fn,
......@@ -114,7 +114,7 @@ def _fast_adapt_window(
return draws, trace, step_size, weighted_running_variance
@tf.function(jit_compile=False)
@tf.function(autograph=False, jit_compile=False)
def _slow_adapt_window(
num_draws,
joint_log_prob_fn,
......@@ -183,7 +183,7 @@ def _slow_adapt_window(
)
@tf.function(jit_compile=False)
@tf.function(autograph=False, jit_compile=False)
def _fixed_window(
num_draws,
joint_log_prob_fn,
......@@ -266,13 +266,19 @@ def trace_results_fn(_, results):
def draws_to_dict(draws):
num_locs = draws[1].shape[1]
num_times = draws[1].shape[2]
return {
"psi": draws[0][:, 0],
"beta_area": draws[0][:, 1],
"gamma0": draws[0][:, 2],
"gamma1": draws[0][:, 3],
"alpha_0": draws[0][:, 4],
"alpha_t": draws[0][:, 5:],
"sigma_space": draws[0][:, 1],
"beta_area": draws[0][:, 2],
"gamma0": draws[0][:, 3],
"gamma1": draws[0][:, 4],
"alpha_0": draws[0][:, 5],
"alpha_t": draws[0][:, 6 : (6 + num_times - 1)],
"spatial_effect": draws[0][
:, (6 + num_times - 1) : (6 + num_times - 1 + num_locs)
],
"seir": draws[1],
}
......@@ -356,7 +362,8 @@ def run_mcmc(
current_state = [s[-1] for s in draws]
draws[0] = param_bijector.inverse(draws[0])
posterior.write_samples(
draws_to_dict(draws), first_dim_offset=offset,
draws_to_dict(draws),
first_dim_offset=offset,
)
posterior.write_results(trace, first_dim_offset=offset)
offset += first_window_size
......@@ -388,7 +395,8 @@ def run_mcmc(
current_state = [s[-1] for s in draws]
draws[0] = param_bijector.inverse(draws[0])
posterior.write_samples(
draws_to_dict(draws), first_dim_offset=offset,
draws_to_dict(draws),
first_dim_offset=offset,
)
posterior.write_results(trace, first_dim_offset=offset)
offset += window_num_draws
......@@ -408,7 +416,8 @@ def run_mcmc(
current_state = [s[-1] for s in draws]
draws[0] = param_bijector.inverse(draws[0])
posterior.write_samples(
draws_to_dict(draws), first_dim_offset=offset,
draws_to_dict(draws),
first_dim_offset=offset,
)
posterior.write_results(trace, first_dim_offset=offset)
offset += last_window_size
......@@ -434,10 +443,12 @@ def run_mcmc(
current_state = [state_part[-1] for state_part in draws]
draws[0] = param_bijector.inverse(draws[0])
posterior.write_samples(
draws_to_dict(draws), first_dim_offset=offset,
draws_to_dict(draws),
first_dim_offset=offset,
)
posterior.write_results(
trace, first_dim_offset=offset,
trace,
first_dim_offset=offset,
)
offset += config["num_burst_samples"]
......@@ -502,8 +513,9 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
tfb.Softplus(low=dtype_util.eps(DTYPE)),
tfb.Identity(),
tfb.Identity(),
tfb.Identity(),
],
block_sizes=[1, 3, events.shape[1]],
block_sizes=[2, 4, events.shape[1] - 1, events.shape[0]],
)
)
......@@ -512,11 +524,17 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
return model.log_prob(
dict(
psi=params[0],
beta_area=params[1],
gamma0=params[2],
gamma1=params[3],
alpha_0=params[4],
alpha_t=params[5:],
sigma_space=params[1],
beta_area=params[2],
gamma0=params[3],
gamma1=params[4],
alpha_0=params[5],
alpha_t=params[6 : (6 + events.shape[1] - 1)],
spatial_effect=params[
(6 + events.shape[1] - 1) : (
6 + events.shape[1] - 1 + events.shape[0]
)
],
seir=events,
)
) + param_bij.inverse_log_det_jacobian(
......@@ -530,8 +548,12 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
current_chain_state = [
tf.concat(
[
np.array([0.1, 0.0, 0.0, 0.0], dtype=DTYPE),
np.full(events.shape[1], -1.75, dtype=DTYPE,),
np.array([0.0, 0.0, 0.0, 0.0, 0.0], dtype=DTYPE),
np.full(
events.shape[1] + events.shape[0],
0.0,
dtype=DTYPE,
),
],
axis=0,
),
......@@ -553,7 +575,8 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
)
posterior._file.create_dataset("initial_state", data=initial_state)
posterior._file.create_dataset(
"time", data=np.array(dates).astype(str).astype(h5py.string_dtype()),
"time",
data=np.array(dates).astype(str).astype(h5py.string_dtype()),
)
print(f"Acceptance theta: {posterior['results/hmc/is_accepted'][:].mean()}")
......
......@@ -26,6 +26,15 @@ TIME_DELTA = 1.0
NU = tf.constant(0.28, dtype=DTYPE) # E->I rate assumed known.
def _compute_adjacency_matrix(geom, names, tol=200):
mat = geom.apply(lambda x: geom.distance(x) < tol)
return xarray.DataArray(
mat.to_numpy().astype(DTYPE),
coords=[names, names],
dims=["location_dest", "location_src"],
)
def gather_data(config):
"""Loads covariate data
......@@ -43,9 +52,12 @@ def gather_data(config):
commute_volume = read_traffic_flow(
config["commute_volume"], date_low=date_low, date_high=date_high
)
geo = gp.read_file(config["geopackage"])
geo = gp.read_file(config["geopackage"], layer="UK2019mod_pop_xgen")
geo = geo.sort_values("lad19cd")
geo = geo[geo["lad19cd"].isin(locations["lad19cd"])]
adjacency = _compute_adjacency_matrix(geo.geometry, geo["lad19cd"], 200)
area = xarray.DataArray(
geo.area,
name="area",
......@@ -68,6 +80,7 @@ def gather_data(config):
C=mobility.astype(DTYPE),
W=commute_volume.astype(DTYPE),
N=popsize.astype(DTYPE),
adjacency=adjacency,
weekday=weekday.astype(DTYPE),
area=area.astype(DTYPE),
locations=xarray.DataArray(
......@@ -132,6 +145,7 @@ def CovidUK(covariates, initial_state, initial_step, num_steps):
)
def alpha_t():
"""Time-varying force of infection"""
return tfd.MultivariateNormalDiag(
loc=tf.constant(0.0, dtype=DTYPE),
scale_diag=tf.fill(
......@@ -139,6 +153,33 @@ def CovidUK(covariates, initial_state, initial_step, num_steps):
),
)
def sigma_space():
"""Variance of CAR prior on space"""
return tfd.HalfNormal(scale=tf.constant(0.1, dtype=DTYPE))
def spatial_effect(sigma_space):
# W = covariates["adjacency"]
# W = tf.linalg.set_diag(W, tf.zeros(W.shape[0], dtype=W.dtype))
# Dw = tf.linalg.diag(tf.reduce_sum(W, axis=-1))
# rho = 0.5
# tf.print("Sigma space:", sigma_space)
# precision = (
# 1.0 / sigma_space * (Dw - rho * W)
# + tf.eye(Dw.shape[0], dtype=DTYPE) * 1e-3
# )
# precision_factor = tf.linalg.cholesky(precision)
# return tfp.experimental.distributions.MultivariateNormalPrecisionFactorLinearOperator(
# loc=tf.constant(0.0, DTYPE),
# precision_factor=tf.linalg.LinearOperatorFullMatrix(
# precision_factor
# ),
# )
return tfd.MultivariateNormalDiag(
loc=tf.constant(0.0, dtype=DTYPE),
scale_diag=tf.ones(covariates["adjacency"].shape[0], dtype=DTYPE),
)
def gamma0():
return tfd.Normal(
loc=tf.constant(0.0, dtype=DTYPE),
......@@ -151,7 +192,16 @@ def CovidUK(covariates, initial_state, initial_step, num_steps):
scale=tf.constant(100.0, dtype=DTYPE),
)
def seir(psi, beta_area, alpha_0, alpha_t, gamma0, gamma1):
def seir(
psi,
beta_area,
alpha_0,
alpha_t,
spatial_effect,
sigma_space,
gamma0,
gamma1,
):
psi = tf.convert_to_tensor(psi, DTYPE)
beta_area = tf.convert_to_tensor(beta_area, DTYPE)
alpha_t = tf.convert_to_tensor(alpha_t, DTYPE)
......@@ -199,8 +249,7 @@ def CovidUK(covariates, initial_state, initial_step, num_steps):
),
),
)
eta = alpha_t_ + beta_area * log_area
eta = alpha_t_ + beta_area * log_area + sigma_space * spatial_effect
infec_rate = tf.math.exp(eta) * (
state[..., 2]
+ psi
......@@ -236,6 +285,8 @@ def CovidUK(covariates, initial_state, initial_step, num_steps):
beta_area=beta_area,
psi=psi,
alpha_t=alpha_t,
sigma_space=sigma_space,
spatial_effect=spatial_effect,
gamma0=gamma0,
gamma1=gamma1,
seir=seir,
......
Markdown is supported
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