Commit 5a434328 authored by Chris Jewell's avatar Chris Jewell
Browse files

Added post-lockdown parameter

parent 90fb9795
......@@ -47,6 +47,7 @@ class CovidUK:
N: np.float64,
date_range: list,
holidays: list,
lockdown: list,
time_step: np.int64):
"""Represents a CovidUK ODE model
......@@ -96,6 +97,8 @@ class CovidUK:
self.m_select = np.int64((self.times >= holidays[0]) &
(self.times < holidays[1]))
self.lockdown_select = np.int64((self.times >= lockdown[0]) &
(self.times < lockdown[1]))
self.max_t = self.m_select.shape[0] - 1
def create_initial_state(self, init_matrix=None):
......@@ -132,8 +135,10 @@ class CovidUKODE(CovidUK):
t_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, self.max_t)
m_switch = tf.gather(self.m_select, t_idx)
commute_volume = tf.pow(tf.gather(self.W, t_idx), param['omega'])
lockdown = tf.gather(self.lockdown_select, t_idx)
beta = tf.where(lockdown == 0, param['beta1'], param['beta1']*param['beta3'])
infec_rate = param['beta1'] * (
infec_rate = beta * (
tf.gather(self.M.matvec(I), m_switch) +
param['beta2'] * self.Kbar * commute_volume * self.C.matvec(I / self.N_sum))
infec_rate = S * infec_rate / self.N
......@@ -224,7 +229,6 @@ class CovidUKStochastic(CovidUK):
shape=[state.shape[0],
state.shape[1],
state.shape[1]])
return rate_matrix
return h
......
......@@ -18,7 +18,8 @@ def sanitise_settings(par_dict):
d = {'inference_period': np.array(par_dict['inference_period'], dtype=np.datetime64),
'prediction_period': np.array(par_dict['prediction_period'], dtype=np.datetime64),
'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']]),
'lockdown': np.array([np.datetime64(date) for date in par_dict['lockdown']])}
return d
......@@ -123,7 +124,7 @@ def extract_locs(in_file: str, out_file: str, loc: list):
la_names = f['la_names'][:].astype(str)
la_loc = np.isin(la_names, loc)
extract = f['prediction'][:, :, :, la_loc]
extract = f['prediction'][:, :, la_loc, :]
save_sims(f['date'][:], extract, f['la_names'][la_loc],
f['age_names'][la_loc], out_file)
......
......@@ -76,6 +76,7 @@ if __name__ == '__main__':
N=N,
date_range=[date_range[0]-np.timedelta64(1,'D'), date_range[1]],
holidays=settings['holiday'],
lockdown=settings['lockdown'],
time_step=int(settings['time_step']))
seeding = seed_areas(N, n_names) # Seed 40-44 age group, 30 seeds by popn size
......@@ -84,16 +85,19 @@ if __name__ == '__main__':
def logp(par):
p = param
p['beta1'] = par[0]
p['gamma'] = par[1]
p['I0'] = par[2]
p['r'] = par[3]
p['beta3'] = par[1]
p['gamma'] = par[2]
p['I0'] = par[3]
p['r'] = par[4]
beta_logp = tfd.Gamma(concentration=tf.constant(1., dtype=DTYPE), rate=tf.constant(1., dtype=DTYPE)).log_prob(p['beta1'])
beta3_logp = tfd.Gamma(concentration=tf.constant(20., dtype=DTYPE),
rate=tf.constant(20., dtype=DTYPE)).log_prob(p['beta3'])
gamma_logp = tfd.Gamma(concentration=tf.constant(100., dtype=DTYPE), rate=tf.constant(400., dtype=DTYPE)).log_prob(p['gamma'])
I0_logp = tfd.Gamma(concentration=tf.constant(1.5, dtype=DTYPE), rate=tf.constant(0.05, dtype=DTYPE)).log_prob(p['I0'])
r_logp = tfd.Gamma(concentration=tf.constant(0.1, dtype=DTYPE), rate=tf.constant(0.1, dtype=DTYPE)).log_prob(p['gamma'])
t, sim, solve = simulator.simulate(p, state_init)
y_logp = covid19uk_logp(y_incr, sim, 0.1, p['r'])
logp = beta_logp + gamma_logp + I0_logp + r_logp + tf.reduce_sum(y_logp)
logp = beta_logp + beta3_logp + gamma_logp + I0_logp + r_logp + tf.reduce_sum(y_logp)
return logp
def trace_fn(_, pkr):
......@@ -104,7 +108,7 @@ if __name__ == '__main__':
unconstraining_bijector = [tfb.Exp()]
initial_mcmc_state = np.array([0.05, 0.25, 1.0, 50.], dtype=np.float64) # beta1, gamma, I0
initial_mcmc_state = np.array([0.05, 1.0, 0.25, 1.0, 50.], dtype=np.float64) # beta1, gamma, I0
print("Initial log likelihood:", logp(initial_mcmc_state))
@tf.function(autograph=False, experimental_compile=True)
......@@ -123,7 +127,7 @@ if __name__ == '__main__':
joint_posterior = tf.zeros([0] + list(initial_mcmc_state.shape), dtype=DTYPE)
scale = np.diag([0.1, 0.1, 0.1, 1.])
scale = np.diag([0.1, 0.1, 0.1, 0.1, 1.])
overall_start = time.perf_counter()
num_covariance_estimation_iterations = 50
......@@ -162,5 +166,5 @@ if __name__ == '__main__':
plt.show()
print(f"Posterior mean: {np.mean(joint_posterior, axis=0)}")
with open('pi_beta_2020-03-15.pkl', 'wb') as f:
with open('pi_beta_2020-03-29.pkl', 'wb') as f:
pkl.dump(joint_posterior, f)
......@@ -11,10 +11,11 @@ data:
parameter:
beta1: 0.1 # R0 2.4
beta2: 0.33 # Contact with commuters 1/3rd of the time
beta3: 1.0 # lockdown vs normal
omega: 1.0 # Non-linearity parameter for commuting volume
nu: 0.25 # E -> I transition rate
gamma: 0.25 # I -> R transition rate
I0: 0.0 # number of individuals at start of epidemic
I0: 1.0 # number of individuals at start of epidemic
r: 30.0 # Negative binomial overdispersion parameter
settings:
......@@ -24,6 +25,9 @@ settings:
holiday:
- 2020-03-23
- 2020-10-01
lockdown:
- 2020-03-23
- 2020-10-01
time_step: 1.
prediction_period:
- 2020-02-19
......
......@@ -56,7 +56,7 @@ if __name__ == '__main__':
y = case_reports['CumCases'].to_numpy()
y_incr = y[1:] - y[-1:]
with open('pi_beta_2020-03-15.pkl', 'rb') as f:
with open('pi_beta_2020-03-29.pkl', 'rb') as f:
pi_beta = pkl.load(f)
# Predictive distribution of epidemic spread
......@@ -71,17 +71,19 @@ if __name__ == '__main__':
date_range=[settings['prediction_period'][0],
settings['prediction_period'][1]],
holidays=settings['holiday'],
lockdown=settings['lockdown'],
time_step=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
def prediction(beta, gamma, I0, r_):
def prediction(beta, beta3, gamma, I0, r_):
sims = tf.TensorArray(tf.float64, size=beta.shape[0])
R0 = tf.TensorArray(tf.float64, size=beta.shape[0])
for i in tf.range(beta.shape[0]):
p = param
p['beta1'] = beta[i]
p['beta3'] = beta3[i]
p['gamma'] = gamma[i]
p['I0'] = I0[i]
p['r'] = r_[i]
......@@ -93,14 +95,14 @@ if __name__ == '__main__':
draws = pi_beta.numpy()[np.arange(5000, pi_beta.shape[0], 30), :]
with tf.device('/CPU:0'):
sims, R0 = prediction(draws[:, 0], draws[:, 1], draws[:, 2], draws[:, 3])
sims, R0 = prediction(draws[:, 0], draws[:, 1], draws[:, 2], draws[:, 3], draws[:, 4])
sims = tf.stack(sims) # shape=[n_sims, n_times, n_metapops, n_states]
save_sims(simulator.times, sims, la_names, age_groups, 'pred_2020-03-23.h5')
save_sims(simulator.times, sims, la_names, age_groups, 'pred_2020-03-29.h5')
dub_time = [doubling_time(simulator.times, sim, '2020-03-01', '2020-04-01') for sim in sims.numpy()]
plot_prediction(settings['prediction_period'], sims, case_reports)
plot_case_incidence(settings['prediction_period'], sims)
plot_prediction(settings['prediction_period'], sims, case_reports)
plot_case_incidence(settings['prediction_period'], sims.numpy())
# R0
......
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