Commit d2e2af66 by Chris Jewell

### Bivariate MCMC algorithm

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