Bivariate MCMC algorithm

covid/model.py

@@ -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

covid/util.py

@@ -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
 
 

covid_ode.py

@@ -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()

data/DailyConfirmedCases.csv

@@ -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

mcmc.py

@@ -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)
-
-

ode_config.yaml

@@ -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:

prediction.py

@@ -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)