Commit 548066ad authored by Chris Jewell's avatar Chris Jewell
Browse files

Re-written connectivity matrix calculations

parent 66841afc
......@@ -290,7 +290,7 @@ if __name__ == "__main__":
tf.random.set_seed(2)
current_state = [
np.array([0.65, 0.0, 0.0], dtype=DTYPE),
np.array([0.2, 0.0, 0.0], dtype=DTYPE),
np.zeros(
model.model["xi"](0.0).event_shape[-1]
# + model.model["beta3"]().event_shape[-1]
......
......@@ -113,9 +113,11 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
# ),
# 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),
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),
......@@ -137,9 +139,11 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
def transition_rate_fn(t, state):
C = tf.convert_to_tensor(covariates["C"], dtype=DTYPE)
C = tf.linalg.set_diag(
C + tf.transpose(C), tf.zeros(C.shape[0], dtype=DTYPE)
)
C = tf.linalg.set_diag(C, tf.zeros(C.shape[0], dtype=DTYPE))
Cstar = C + tf.transpose(C)
Cstar = tf.linalg.set_diag(Cstar, -tf.reduce_sum(C, axis=-2))
W = tf.constant(np.squeeze(covariates["W"]), dtype=DTYPE)
N = tf.constant(np.squeeze(covariates["N"]), dtype=DTYPE)
......@@ -164,7 +168,7 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
state[..., 2]
+ beta2
* commute_volume
* tf.linalg.matvec(C, state[..., 2] / tf.squeeze(N))
* tf.linalg.matvec(Cstar, state[..., 2] / tf.squeeze(N))
)
infec_rate = (
infec_rate / tf.squeeze(N) + 0.000000001
......@@ -221,9 +225,11 @@ def next_generation_matrix_fn(covar_data, param):
L = L - tf.reduce_mean(L, axis=0)
C = tf.convert_to_tensor(covar_data["C"], dtype=DTYPE)
C = tf.linalg.set_diag(
C + tf.transpose(C), tf.zeros(C.shape[0], dtype=DTYPE)
)
C = tf.linalg.set_diag(C, -tf.reduce_sum(C, axis=-2))
C = tf.linalg.set_diag(C, tf.zeros(C.shape[0], dtype=DTYPE))
Cstar = C + tf.transpose(C)
Cstar = tf.linalg.set_diag(Cstar, -tf.reduce_sum(C, axis=-2))
W = tf.constant(covar_data["W"], dtype=DTYPE)
N = tf.constant(covar_data["N"], dtype=DTYPE)
......@@ -242,8 +248,8 @@ def next_generation_matrix_fn(covar_data, param):
beta = tf.math.exp(xi + xB)
ngm = beta[:, tf.newaxis] * (
tf.eye(C.shape[0], dtype=state.dtype)
+ param["beta2"] * commute_volume * C / N[tf.newaxis, :]
tf.eye(Cstar.shape[0], dtype=state.dtype)
+ param["beta2"] * commute_volume * Cstar / N[tf.newaxis, :]
)
ngm = (
ngm
......
......@@ -16,9 +16,8 @@ import model_spec
def make_within_rate_fns(covariates, beta2):
C = tf.convert_to_tensor(covariates["C"], dtype=model_spec.DTYPE)
C = tf.linalg.set_diag(
C + tf.transpose(C), tf.zeros(C.shape[0], dtype=model_spec.DTYPE)
)
C = tf.linalg.set_diag(C, tf.zeros(C.shape[0], dtype=model_spec.DTYPE))
W = tf.convert_to_tensor(
tf.squeeze(covariates["W"]), dtype=model_spec.DTYPE
)
......@@ -27,13 +26,21 @@ def make_within_rate_fns(covariates, beta2):
)
def within_fn(t, state):
rate = state[..., 2]
w_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, W.shape[0] - 1)
commute_volume = tf.gather(W, w_idx)
rate = state[..., 2] - beta2 * state[
..., 2
] / N * commute_volume * tf.reduce_sum(C, axis=-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)
rate = beta2 * commute_volume * tf.linalg.matvec(C, state[..., 2] / N)
rate = (
beta2
* commute_volume
* tf.linalg.matvec(C + tf.transpose(C), state[..., 2] / N)
)
return rate
return within_fn, between_fn
......
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