...
 
Commits (2)
......@@ -88,7 +88,7 @@ def dense_to_block_diagonal(A, n_blocks):
class CovidUKODE: # TODO: add background case importation rate to the UK, e.g. \epsilon term.
def __init__(self, M_tt, M_hh, C, N, start, end, holidays, t_step):
def __init__(self, M_tt, M_hh, C, N, start, end, holidays, bg_max_t, t_step):
"""Represents a CovidUK ODE model
:param K_tt: a MxM matrix of age group mixing in term time
......@@ -117,7 +117,8 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
self.times = np.arange(start, end, np.timedelta64(t_step, 'D'))
m_select = (np.less_equal(holidays[0], self.times) & np.less(self.times, holidays[1])).astype(np.int64)
self.m_select = tf.constant(m_select)
self.m_select = tf.constant(m_select, dtype=tf.int64)
self.bg_select = tf.constant(np.less(bg_max_t, self.times), dtype=tf.int64)
self.solver = tode.DormandPrince()
def make_h(self, param):
......@@ -125,13 +126,15 @@ class CovidUKODE: # TODO: add background case importation rate to the UK, e.g.
def h_fn(t, state):
state = tf.unstack(state, axis=0)
S, E, I, R = state
m_switch = tf.gather(self.m_select, tf.cast(t, tf.int64))
t = tf.clip_by_value(tf.cast(t, tf.int64), 0, self.m_select.shape[0]-1)
m_switch = tf.gather(self.m_select, t)
epsilon = param['epsilon'] * tf.cast(tf.gather(self.bg_select, t), tf.float32)
if m_switch == 0:
infec_rate = param['beta1'] * tf.linalg.matvec(self.M[0], I)
else:
infec_rate = param['beta1'] * tf.linalg.matvec(self.M[1], I)
infec_rate += param['beta1'] * param['beta2'] * self.Kbar * tf.linalg.matvec(self.C, I / self.N_sum)
infec_rate = S / self.N * infec_rate
infec_rate = S / self.N * (infec_rate + epsilon)
dS = -infec_rate
dE = infec_rate - param['nu'] * E
......
......@@ -2,9 +2,10 @@
import numpy as np
def sanitise_parameter(par_dict):
"""Sanitises a dictionary of parameters"""
par = ['beta1', 'beta2', 'nu', 'gamma']
par = ['epsilon', 'beta1', 'beta2', 'nu', 'gamma']
d = {key: np.float64(par_dict[key]) for key in par}
return d
......@@ -13,7 +14,8 @@ def sanitise_settings(par_dict):
d = {'start': np.datetime64(par_dict['start']),
'end': np.datetime64(par_dict['end']),
'time_step': float(par_dict['time_step']),
'holiday': np.array([np.datetime64(date) for date in par_dict['holiday']])}
'holiday': np.array([np.datetime64(date) for date in par_dict['holiday']]),
'bg_max_time': np.datetime64(par_dict['bg_max_time'])}
return d
......
......@@ -12,7 +12,7 @@ from covid.rdata import *
def sanitise_parameter(par_dict):
"""Sanitises a dictionary of parameters"""
par = ['beta1', 'beta2', 'nu', 'gamma']
par = ['epsilon', 'beta1', 'beta2', 'nu', 'gamma']
d = {key: np.float64(par_dict[key]) for key in par}
return d
......@@ -21,7 +21,8 @@ def sanitise_settings(par_dict):
d = {'start': np.datetime64(par_dict['start']),
'end': np.datetime64(par_dict['end']),
'time_step': float(par_dict['time_step']),
'holiday': np.array([np.datetime64(date) for date in par_dict['holiday']])}
'holiday': np.array([np.datetime64(date) for date in par_dict['holiday']]),
'bg_max_time': np.datetime64(par_dict['bg_max_time'])}
return d
......@@ -211,33 +212,17 @@ if __name__ == '__main__':
K1_hh[14:, :] *= cocooning_ratio
K1_hh[:, 14:] *= cocooning_ratio
model_term = CovidUKODE(K1_tt, T, N)
model_holiday = CovidUKODE(K1_hh, T, N)
model = CovidUKODE(K1_tt, K_hh, T, N, settings['start'], settings['end'], settings['holiday'],
settings['bg_max_time'], 1)
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, 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
state_init = model.create_initial_state(init_matrix=seeding)
print('R_term=', model.eval_R0(param))
#print('R_holiday=', model_holiday.eval_R0(param))
start = time.perf_counter()
t, sim = simulate()
t, sim, _ = model.simulate(param, state_init)
end = time.perf_counter()
print(f'Complete in {end - start} seconds')
......@@ -252,11 +237,10 @@ if __name__ == '__main__':
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")
# 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_uk.gca().legend(title="Contact ratio")
fig_attack.autofmt_xdate()
fig_uk.autofmt_xdate()
......
......@@ -42,4 +42,5 @@ DateVal,CMODateCount,CumCases
2020-03-11,83,456
2020-03-12,139,590
2020-03-13,207,797
2020-03-14,343,1140
2020-03-14,264,1061
2020-03-15,330,1391
......@@ -63,54 +63,72 @@ if __name__ == '__main__':
y_incr = np.round((y[1:] - y[:-1]) * 0.8)
simulator = CovidUKODE(K_tt, K_hh, T, N, date_range[0]-np.timedelta64(1,'D'),
date_range[1], settings['holiday'], int(settings['time_step']))
date_range[1], settings['holiday'], settings['bg_max_time'], int(settings['time_step']))
seeding = seed_areas(N, n_names) # Seed 40-44 age group, 30 seeds by popn size
state_init = simulator.create_initial_state(init_matrix=seeding)
#@tf.function
def logp(beta1):
def logp(epsilon, beta):
p = param
p['beta1'] = beta1
p['epsilon'] = epsilon
p['beta1'] = beta
epsilon_logp = tfd.Gamma(concentration=1., rate=1.).log_prob(p['epsilon'])
beta_logp = tfd.Gamma(concentration=1., rate=1.).log_prob(p['beta1'])
t, sim, solve = simulator.simulate(p, state_init)
y_logp = covid19uk_logp(y_incr, sim, 0.1)
return beta_logp + tf.reduce_sum(y_logp)
logp = epsilon_logp + beta_logp + tf.reduce_sum(y_logp)
return logp
unconstraining_bijector = [tfb.Log()]
initial_mcmc_state = [0.03]
print("Initial log likelihood:", logp(0.03))
def trace_fn(_, pkr):
return (
pkr.inner_results.log_accept_ratio,
pkr.inner_results.accepted_results.target_log_prob,
pkr.inner_results.accepted_results.step_size)
unconstraining_bijector = [tfb.Exp(), tfb.Exp()]
initial_mcmc_state = [0.00001, 0.03]
print("Initial log likelihood:", logp(*initial_mcmc_state))
@tf.function
def sample():
def sample(n_samples, step_size=[0.01, 0.1]):
return tfp.mcmc.sample_chain(
num_results=2000,
num_burnin_steps=500,
num_results=n_samples,
num_burnin_steps=0,
current_state=initial_mcmc_state,
kernel=tfp.mcmc.SimpleStepSizeAdaptation(
tfp.mcmc.TransformedTransitionKernel(
inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=logp,
step_size=0.0001,
step_size=step_size,
num_leapfrog_steps=5),
bijector=unconstraining_bijector),
num_adaptation_steps=400),
trace_fn=lambda _, pkr: pkr.inner_results.inner_results.is_accepted)
start = time.perf_counter()
pi_beta, accept = sample()
end = time.perf_counter()
print(f"Simulation complete in {end-start} seconds")
plt.plot(pi_beta[0])
num_adaptation_steps=500),
trace_fn=lambda _, pkr: pkr.inner_results.inner_results.accepted_results.step_size)
with tf.device("/CPU:0"):
start = time.perf_counter()
pi_beta, results = sample(1000, [0.1, 0.1])
# sd = tfp.stats.stddev(pi_beta, sample_axis=1)
# sd = sd / tf.reduce_sum(sd)
# new_step_size = [results[-1] * sd[0], results[-1] * sd[1]]
# print("New step size:",new_step_size)
# pi_beta, results = sample(1000, step_size=new_step_size)
end = time.perf_counter()
print(f"Simulation complete in {end-start} seconds")
print(results)
fig, ax = plt.subplots(1, 2)
ax[0].plot(pi_beta[0])
ax[1].plot(pi_beta[1])
plt.show()
print(f"Posterior mean: {np.mean(pi_beta[0])}")
with open('pi_beta_2020-03-15.pkl','wb') as f:
with open('pi_beta_2020-03-15.pkl', 'wb') as f:
pkl.dump(pi_beta, f)
#dates = settings['start'] + t.numpy().astype(np.timedelta64)
#plotting(dates, sim)
......@@ -8,18 +8,19 @@ data:
reported_cases: data/DailyConfirmedCases.csv
parameter:
beta1: 0.035 # R0 2.4
#beta1: 0.04
epsilon: 1.0e-6
beta1: 0.04 # R0 2.4
beta2: 0.33 # Contact with commuters 1/3rd of the time
nu: 0.25
gamma: 0.25
settings:
start: 2020-02-04
end: 2020-05-01
end: 2020-04-01
holiday:
- 2020-04-06
- 2020-04-17
bg_max_time: 2020-03-01
time_step: 1.
output:
......
......@@ -31,6 +31,7 @@ if __name__ == '__main__':
N, n_names = load_population(config['data']['population_size'])
param = sanitise_parameter(config['parameter'])
param['epsilon'] = 0.0
settings = sanitise_settings(config['settings'])
case_reports = pd.read_csv(config['data']['reported_cases'])
......@@ -47,11 +48,11 @@ if __name__ == '__main__':
date_range[1]+np.timedelta64(1,'D'),
np.timedelta64(1, 'D'))
simulator = CovidUKODE(K_tt, K_hh, T, N, date_range[0] - np.timedelta64(1, 'D'),
np.datetime64('2020-07-01'), settings['holiday'], 1)
np.datetime64('2020-05-01'), settings['holiday'], settings['bg_max_time'], 1)
seeding = seed_areas(N, n_names) # Seed 40-44 age group, 30 seeds by popn size
state_init = simulator.create_initial_state(init_matrix=seeding)
#@tf.function
@tf.function
def prediction(beta):
sims = tf.TensorArray(tf.float32, size=beta.shape[0])
for i in tf.range(beta.shape[0]):
......@@ -65,7 +66,7 @@ if __name__ == '__main__':
draws = pi_beta[0].numpy()[np.arange(0, pi_beta[0].shape[0], 20)]
sims = prediction(draws)
dates = np.arange(date_range[0]-np.timedelta64(1, 'D'), np.datetime64('2020-07-01'),
dates = np.arange(date_range[0]-np.timedelta64(1, 'D'), np.datetime64('2020-05-01'),
np.timedelta64(1, 'D'))
total_infected = tfs.percentile(tf.reduce_sum(sims[:, :, 1:3], axis=2), q=[2.5, 50, 97.5], axis=0)
removed = tfs.percentile(sims[:, :, 3], q=[2.5, 50, 97.5], axis=0)
......