Added post-lockdown parameter

covid/model.py

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

covid/util.py

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

mcmc.py

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

ode_config.yaml

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

prediction.py

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