Skip to content
GitLab
Commit 45ac7d72 authored by Chris Jewell's avatar Chris Jewell
Browse files

Age-specific attack rate analysis.

parent aaf591bc
Hide whitespace changes
Inline Side-by-side
covid/impl/chainbinom_simulate.py
View file @ 45ac7d72
......@@ -14,9 +14,11 @@ def update_state(update, state, stoichiometry):
def chain_binomial_propagate(h, time_step, stoichiometry):
def propagate_fn(state):
rates = h(state)
state_idx, rates = h(state)
probs = 1 - tf.exp(-rates*time_step) # RxN
update = tfd.Binomial(state[:-1], probs=probs).sample() # RxN
state_mult = tf.scatter_nd(state_idx[:, None], state,
shape=[state_idx.shape[0], state.shape[1], state.shape[2]])
update = tfd.Binomial(state_mult, probs=probs).sample() # RxN
update = tf.expand_dims(update, 1) # Rx1xN
upd_shape = tf.concat([stoichiometry.shape, tf.fill([tf.rank(state)-1], 1)], axis=0)
update *= tf.reshape(stoichiometry, upd_shape) # RxSx1
......
covid/model.py
View file @ 45ac7d72
"""Functions for infection rates"""
from warnings import warn
import tensorflow as tf
import tensorflow_probability as tfp
......@@ -49,7 +49,7 @@ class CovidUK:
])
return hazard_rates
#@tf.function
def sample(self, initial_state, time_lims, param):
self.param = param
return chain_binomial_simulate(self.h, initial_state, time_lims[0],
......@@ -98,7 +98,7 @@ class CovidUKODE:
eye = tf.linalg.LinearOperatorIdentity(self.n_lads)
self.K = tf.linalg.LinearOperatorKronecker([eye, self.K])
self.T = tf.linalg.LinearOperatorFullMatrix(T)
self.T = tf.linalg.LinearOperatorFullMatrix(T + tf.transpose(T))
shp = tf.linalg.LinearOperatorFullMatrix(np.ones_like(K, dtype=np.float32))
self.T = tf.linalg.LinearOperatorKronecker([self.T, shp])
......@@ -144,14 +144,15 @@ class CovidUKODE:
return np.stack([S, E, I, R])
@tf.function
def simulate(self, param, state_init, t_start, t_end, t_step=1.):
def simulate(self, param, state_init, t_start, t_end, t_step=1., solver_state=None):
h = self.make_h(param)
t0 = 0.
t1 = (t_end - t_start) / np.timedelta64(1, 'D')
t = np.arange(start=t0, stop=t1, step=t_step)[1:]
solver = tode.DormandPrince()
results = solver.solve(ode_fn=h, initial_time=t0, initial_state=state_init, solution_times=t)
return results.times, results.states
results = solver.solve(ode_fn=h, initial_time=t0, initial_state=state_init, solution_times=t,
previous_solver_internal_state=solver_state)
return results.times, results.states, results.solver_internal_state
def ngm(self, param):
infec_rate = param['beta1'] * self.K.to_dense()
......
covid_ode.py
View file @ 45ac7d72
......@@ -12,7 +12,7 @@ from covid.rdata import *
def sanitise_parameter(par_dict):
"""Sanitises a dictionary of parameters"""
par = ['beta1', 'beta2', 'nu','gamma']
par = ['beta1', 'beta2', 'nu', 'gamma']
d = {key: np.float64(par_dict[key]) for key in par}
return d
......@@ -82,7 +82,6 @@ def write_hdf5(filename, param, t, sim):
d_beta[()] = v
def plot_total_curve(sim):
infec_uk = sum_la(sim)
infec_uk = infec_uk.sum(axis=1)
......@@ -96,6 +95,13 @@ def plot_total_curve(sim):
plt.legend()
def plot_infec_curve(ax, sim, label, color):
infec_uk = sum_la(sim)
infec_uk = infec_uk.sum(axis=1)
times = np.datetime64('2020-02-20') + np.arange(infec_uk.shape[0])
ax.plot(times, infec_uk, '-', label=label, color=color)
def plot_by_age(sim, labels, t0=np.datetime64('2020-02-20'), ax=None):
if ax is None:
ax = plt.figure().gca()
......@@ -109,7 +115,7 @@ def plot_by_age(sim, labels, t0=np.datetime64('2020-02-20'), ax=None):
return ax
def plot_by_la(sim, labels,t0=np.datetime64('2020-02-20'), ax=None):
def plot_by_la(sim, labels, t0=np.datetime64('2020-02-20'), ax=None):
if ax is None:
ax = plt.figure().gca()
infec_uk = sum_age_groups(sim)
......@@ -128,7 +134,7 @@ def draw_figs(sim, N):
fs = final_size(sim)
attack_rate = fs / N
print("Attack rate:", attack_rate)
print("Overall attack rate: ", np.sum(fs)/np.sum(N))
print("Overall attack rate: ", np.sum(fs) / np.sum(N))
# Total UK epidemic curve
plot_total_curve(sim)
......@@ -155,6 +161,23 @@ def draw_figs(sim, N):
plt.show()
def doubling_time(t, sim, t1, t2):
t1 = np.where(t == np.datetime64(t1))[0]
t2 = np.where(t == np.datetime64(t2))[0]
delta = t2 - t1
r = sum_total_removals(sim)
q1 = r[t1]
q2 = r[t2]
return delta * np.log(2) / np.log(q2 / q1)
def plot_age_attack_rate(ax, sim, N, label, color):
Ns = N.reshape([152, 17]).sum(axis=0)
fs = final_size(sim.numpy())
attack_rate = fs / Ns
ax.plot(age_groups, attack_rate, 'o-', color=color, label=label)
if __name__ == '__main__':
parser = optparse.OptionParser()
......@@ -176,40 +199,69 @@ if __name__ == '__main__':
settings = sanitise_settings(config['settings'])
# Straight, no school closures
model_term = CovidUKODE(K_tt, T, N)
model_holiday = CovidUKODE(K_hh/2., T/2., N)
seeding = seed_areas(N, n_names) # Seed 40-44 age group, 30 seeds by popn size
state_init = model_term.create_initial_state(init_matrix=seeding)
print('R_term=', model_term.eval_R0(param))
print('R_holiday=', model_holiday.eval_R0(param))
# School holidays and closures
@tf.function
def simulate():
t0, sim_0 = model_term.simulate(param, state_init,
settings['start'], settings['holiday'][0],
settings['time_step'])
t1, sim_1 = model_holiday.simulate(param, sim_0[-1, :, :],
settings['holiday'][0], settings['holiday'][1],
settings['time_step'])
t2, sim_2 = model_term.simulate(param, sim_1[-1, :, :],
settings['holiday'][1], settings['end'],
settings['time_step'])
t = tf.concat([t0, t1, t2], axis=0)
sim = tf.concat([tf.expand_dims(state_init, axis=0), sim_0, sim_1, sim_2], axis=0)
return t, sim
# Effect of cocooning on elderly (70+)
def cocooning_model(cocooning_ratio=1.):
start = time.perf_counter()
t, sim = simulate()
end = time.perf_counter()
print(f'Complete in {end-start} seconds')
K1_tt = K_tt.copy()
K1_hh = K_hh.copy()
draw_figs(sim.numpy(), N)
K1_tt[14:, :] *= cocooning_ratio
K1_tt[:, 14:] *= cocooning_ratio
K1_hh[14:, :] *= cocooning_ratio
K1_hh[:, 14:] *= cocooning_ratio
if 'simulation' in config['output']:
write_hdf5(config['output']['simulation'], param, t, sim)
model_term = CovidUKODE(K1_tt, T, N)
model_holiday = CovidUKODE(K1_hh, T, N)
seeding = seed_areas(N, n_names) # Seed 40-44 age group, 30 seeds by popn size
state_init = model_term.create_initial_state(init_matrix=seeding)
print('R_term=', model_term.eval_R0(param))
print('R_holiday=', model_holiday.eval_R0(param))
# School holidays and closures
def simulate():
t0, sim_0, solve0 = model_term.simulate(param, state_init,
settings['start'], settings['holiday'][0],
settings['time_step'])
t1, sim_1, solve1 = model_holiday.simulate(param, sim_0[-1, :, :],
settings['holiday'][0], settings['holiday'][1],
settings['time_step'], solver_state=None)
t2, sim_2, _ = model_term.simulate(param, sim_1[-1, :, :],
settings['holiday'][1], settings['end'],
settings['time_step'], solver_state=None)
t = tf.concat([t0, t1 + t0[-1], t2 + t0[-1] + t1[-1]], axis=0)
sim = tf.concat([tf.expand_dims(state_init, axis=0), sim_0, sim_1, sim_2], axis=0)
return t, sim
start = time.perf_counter()
t, sim = simulate()
end = time.perf_counter()
print(f'Complete in {end - start} seconds')
dates = settings['start'] + t.numpy().astype(np.timedelta64)
dt = doubling_time(dates, sim.numpy(), '2020-03-01', '2020-04-01')
print(f"Doubling time: {dt}")
return t, sim
fig_attack = plt.figure()
fig_uk = plt.figure()
cocoon_ratios = [1.]
for i, r in enumerate(cocoon_ratios):
print(f"Simulation, r={r}", flush=True)
t, sim = cocooning_model(r)
plot_age_attack_rate(fig_attack.gca(), sim, N, f"{1 - r}", f"C{i}")
#fig_attack.gca().legend(title="Contact ratio")
plot_infec_curve(fig_uk.gca(), sim.numpy(), f"{1 - r}", f"C{i}")
#fig_uk.gca().legend(title="Contact ratio")
fig_attack.autofmt_xdate()
fig_uk.autofmt_xdate()
fig_attack.gca().grid(True)
fig_uk.gca().grid(True)
plt.show()
print(f"Complete in {end-start} seconds")
\ No newline at end of file
# if 'simulation' in config['output']:
# write_hdf5(config['output']['simulation'], param, t, sim)
household_sim.py
View file @ 45ac7d72
......@@ -23,15 +23,20 @@ def draw_households(n: int, dist: pd.DataFrame) -> np.array:
class HHstochastic(CovidUK):
def __init__(self, N):
self.stoichiometry = tf.constant([[-1, 1, 0, 0],
[0, -1, 1, 0],
[0, 0, -1, 1]], dtype=tf.float64)
self.stoichiometry = tf.constant([[-1, 1, 0, 0, 0], # S->E
[0, -1, 1, 0, 0], # E->I
[0, 0, -1, 1, 0], # I->D
[0, 0, -1, 0, 1], # I->R
[0, 0, 0, -1, 1]], dtype=tf.float64) # D->R
self.N = N
self.popsize = tf.reduce_sum(N)
def rules(self, state):
tf.
def h(self, state):
state = tf.unstack(state, axis=0)
S, E, I, R = state
S, E, I, D, R = state
eye = tf.eye(I.shape[1], dtype=tf.float64)
......@@ -42,27 +47,37 @@ class HHstochastic(CovidUK):
infec_rate *= self.param['beta']
hazard_rates = tf.stack([
infec_rate,
tf.constant(np.full(S.shape, self.param['nu']), dtype=tf.float64),
tf.constant(np.full(S.shape, self.param['gamma']), dtype=tf.float64)
infec_rate, # S->E
tf.constant(np.full(S.shape, self.param['nu']), dtype=tf.float64), # E->I
tf.constant(np.full(S.shape, self.param['alpha']), dtype=tf.float64), # I->D
self.param['rho'] * D, # S->R
tf.constant(np.full(S.shape, self.param['gamma']), dtype=tf.float64) # I->R
])
return hazard_rates
return tf.constant([0, 1, 2, 2, 3]), hazard_rates
if __name__=='__main__':
hh = draw_households(100, hh_size_distr)
K = (hh[:, None] == hh[None, :])
nsim = 1000
hh = draw_households(20, hh_size_distr)
nsim = 2
param = {'beta': 0.3, # Baseline infection
'omega': 1.5, # Additional transmission in house
'nu': 0.25, # Incubation rate
'alpha': 0.5, # Prodromal
'gamma': 0.5, # Symptomatic
'rho': 10e6} # Isolation for whole house
model = HHstochastic(hh.astype(np.float64))
init_state = np.stack([np.broadcast_to(hh, [nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]])]).astype(np.float64)
init_state = np.stack([np.broadcast_to(hh, [nsim, hh.shape[0]]), # S
np.zeros([nsim, hh.shape[0]]), # E
np.zeros([nsim, hh.shape[0]]), # I
np.zeros([nsim, hh.shape[0]]), # D
np.zeros([nsim, hh.shape[0]])]).astype(np.float64) # R
init_state[2, :, 0] = 1.
init_state[0, :, 0] = init_state[0, :, 0] - 1.
start = time.perf_counter()
t, sim = model.sample(init_state, [0., 365.], {'beta': 0.3, 'omega': 1.5, 'nu': 0.25, 'gamma': 0.25})
t, sim = model.sample(init_state, [0., 5.], param)
end = time.perf_counter()
print(f"Completed in {end-start} seconds")
......
ode_config.yaml
View file @ 45ac7d72
......@@ -7,17 +7,18 @@ data:
population_size: data/pop.rds
parameter:
beta1: 0.04 # R0 2.4
# beta1: 0.0347 # R0 2.4
beta1: 0.04
beta2: 0.33 # Contact with commuters 1/3rd of the time
nu: 0.25
gamma: 0.25
settings:
start: 2020-02-04
end: 2020-12-30
end: 2020-05-01
holiday:
- 2020-04-06
- 2020-05-01
- 2020-04-17
time_step: 1.
output:
......
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