def body(state_, occults_):
state_t1 = tf.roll(state_, shift=-1, axis=-2)
neg_state_idx = tf.where(state_t1 < 0)
first_neg_state_idx = tf.gather(
neg_state_idx,
tf.concat(
[
[[0]],
tf.where(neg_state_idx[:-1, 0] - neg_state_idx[1:, 0]) + 1,
],
axis=0,
),
)
mask = tf.scatter_nd(
first_neg_state_idx,
tf.ones([first_neg_state_idx.shape[0], 1], dtype=state_t1.dtype),
state_t1.shape,
)
delta_occults = tf.einsum("mts,xs->mtx", state_t1 * mask,
stoichiometry)
new_occults = tf.clip_by_value(occults_ - delta_occults,
clip_value_min=0.0,
clip_value_max=1.0e6)
new_state = compute_state(init_state, events + new_occults,
stoichiometry)
return new_state, new_occults
def test_compute_state(self):
init_state = tf.constant([[99, 1, 0], [100, 0, 0], [100, 0, 0]],
dtype=tf.float32)
events = tf.constant(
[
[[3, 0], [3, 2], [0, 2]],
[[1, 0], [2, 1], [1, 2]],
[[5, 0], [0, 5], [0, 0]],
],
dtype=tf.float32,
)
stoichiometry = tf.constant([[-1, 1, 0], [0, -1, 1]], dtype=tf.float32)
actual = compute_state(init_state, events, stoichiometry)
expected = np.array(
[
[[99, 1, 0], [96, 4, 0], [93, 5, 2]],
[[100, 0, 0], [99, 1, 0], [97, 2, 1]],
[[100, 0, 0], [95, 5, 0], [95, 0, 5]],
],
dtype=np.float32,
)
np.testing.assert_array_equal(expected, actual)
def r_fn(args):
theta_, xi_, 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
par = dict(beta1=theta_[0], beta2=theta_[1], gamma=theta_[2], xi=xi_)
ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par)
ngm = ngm_fn(t, state)
return ngm
def regularize_occults(events, occults, init_state, stoichiometry):
"""Regularizes an occult matrix such that counting
processes are valid
:param events: a [M, T, X] events tensor
:param occults: a [M, T, X] occults tensor
:param init_state: a [M, S] initial state tensor
:param stoichiometry: a [X, S] stoichiometry matrix
:returns: an tuple containing updated (state, occults) tensors
"""
from covid.impl.util import compute_state
def body(state_, occults_):
state_t1 = tf.roll(state_, shift=-1, axis=-2)
neg_state_idx = tf.where(state_t1 < 0)
first_neg_state_idx = tf.gather(
neg_state_idx,
tf.concat(
[
[[0]],
tf.where(neg_state_idx[:-1, 0] - neg_state_idx[1:, 0]) + 1,
],
axis=0,
),
)
mask = tf.scatter_nd(
first_neg_state_idx,
tf.ones([first_neg_state_idx.shape[0], 1], dtype=state_t1.dtype),
state_t1.shape,
)
delta_occults = tf.einsum("mts,xs->mtx", state_t1 * mask,
stoichiometry)
new_occults = tf.clip_by_value(occults_ - delta_occults,
clip_value_min=0.0,
clip_value_max=1.0e6)
new_state = compute_state(init_state, events + new_occults,
stoichiometry)
return new_state, new_occults
def cond(state_, _):
return tf.reduce_any(state_ < 0)
state = compute_state(init_state, events + occults, stoichiometry)
new_state, new_occults = tf.while_loop(cond, body, (state, occults))
return new_state, new_occults
def discrete_markov_log_prob(events, init_state, init_step, time_delta,
hazard_fn, stoichiometry):
"""Calculates an unnormalised log_prob function for a discrete time epidemic model.
:param events: a `[M, T, X]` batch of transition events for metapopulation M,
times `T`, and transitions `X`.
:param init_state: a vector of shape `[M, S]` the initial state of the epidemic for
`M` metapopulations and `S` states
:param init_step: the initial time step, as an offset to `range(events.shape[-2])`
:param time_delta: the size of the time step.
:param hazard_fn: a function that takes a state and returns a matrix of transition
rates.
:param stoichiometry: a `[X, S]` matrix describing the state update for each
transition.
:return: a scalar log probability for the epidemic.
"""
num_meta = events.shape[-3]
num_times = events.shape[-2]
num_events = events.shape[-1]
num_states = stoichiometry.shape[-1]
state_timeseries = compute_state(init_state, events,
stoichiometry) # MxTxS
tms_timeseries = tf.transpose(state_timeseries, perm=(1, 0, 2))
def fn(elems):
return hazard_fn(*elems)
tx_coords = transition_coords(stoichiometry)
rates = tf.vectorized_map(fn=fn,
elems=[tf.range(num_times), tms_timeseries])
rate_matrix = make_transition_matrix(rates, tx_coords,
tms_timeseries.shape)
probs = approx_expm(rate_matrix * time_delta)
# [T, M, S, S] to [M, T, S, S]
probs = tf.transpose(probs, perm=(1, 0, 2, 3))
event_matrix = make_transition_matrix(events, tx_coords,
[num_meta, num_times, num_states])
event_matrix = tf.linalg.set_diag(
event_matrix, state_timeseries - tf.reduce_sum(event_matrix, axis=-1))
logp = tfd.Multinomial(
tf.cast(state_timeseries, dtype=tf.float32),
probs=tf.cast(probs, dtype=tf.float32),
name="log_prob",
).log_prob(tf.cast(event_matrix, dtype=tf.float32))
return tf.cast(tf.reduce_sum(logp), dtype=events.dtype)
# Load in covariate data
covar_data = model_spec.read_covariates(config["data"])
# 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"])
# Single imputation of censored data
events = model_spec.impute_censored_events(cases)
# Initial conditions S(0), E(0), I(0), R(0) are calculated
# by calculating the state at the beginning of the inference period
state = compute_state(
initial_state=tf.concat(
[covar_data["N"], tf.zeros_like(events[:, 0, :])], axis=-1),
events=events,
stoichiometry=model_spec.STOICHIOMETRY,
)
start_time = state.shape[1] - cases.shape[1]
initial_state = state[:, start_time, :]
events = events[:, start_time:, :]
# Build model and sample
full_probability_model = model_spec.CovidUK(
covariates=covar_data,
initial_state=initial_state,
initial_step=0,
num_steps=80,
)
seir = full_probability_model.model["seir"](beta1=0.35,
beta2=0.65,
# Load posterior file
posterior = h5py.File(
config["output"]["posterior"],
"r",
rdcc_nbytes=1024**3,
rdcc_nslots=1e6,
)
# Pre-determined thinning of posterior (better done in MCMC?)
idx = slice(posterior["samples/theta"].shape[0]) # range(6000, 10000, 10)
theta = posterior["samples/theta"][idx]
xi = posterior["samples/xi"][idx]
events = posterior["samples/events"][idx]
init_state = posterior["initial_state"][:]
state_timeseries = compute_state(init_state, events,
model_spec.STOICHIOMETRY)
# Build model
model = model_spec.CovidUK(
covar_data,
initial_state=init_state,
initial_step=0,
num_steps=events.shape[1],
)
ngms = calc_R_it(theta, xi, 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 = np.stack([q, np.quantile(rt, q)], axis=-1)