Fixed missing right bracket

covid/impl/chainbinom_simulate.py

@@ -32,7 +32,7 @@ def chain_binomial_propagate(h, time_step):
           probs = markov_transition[..., :, i]
           binom = tfd.Binomial(
               total_count=total_count,
-              probs=tf.clip_by_value(probs / (1. - prev_probs), 0., 1.)
+              probs=tf.clip_by_value(probs / (1. - prev_probs), 0., 1.))
           sample = binom.sample()
           counts = tf.concat([counts, sample[..., tf.newaxis]], axis=-1)
           total_count -= sample
@@ -47,7 +47,6 @@ def chain_binomial_propagate(h, time_step):
 def chain_binomial_simulate(hazard_fn, state, start, end, time_step):
     propagate = chain_binomial_propagate(hazard_fn, time_step)
     times = tf.range(start, end, time_step)
-    print(times.shape[0])
 
     output = tf.TensorArray(state.dtype, size=times.shape[0])
     output = output.write(0, state)

covid/model.py

@@ -101,7 +101,7 @@ class CovidUK:
     def create_initial_state(self, init_matrix=None):
         if init_matrix is None:
             I = np.zeros(self.N.shape, dtype=np.float64)
-            I[149*17+10] = 30. # Middle-aged in Surrey
+            I[149*17+10] = 30.  # Middle-aged in Surrey
         else:
             np.testing.assert_array_equal(init_matrix.shape, [self.n_lads, self.n_ages],
                                           err_msg=f"init_matrix does not have shape [<num lads>,<num ages>] \
@@ -224,7 +224,7 @@ class CovidUKStochastic(CovidUK):
             return rate_matrix
         return h
 
-    @tf.function(autograph=False)
+    @tf.function(autograph=False, experimental_compile=True)
     def simulate(self, param, state_init):
         """Runs a simulation from the epidemic model
 

covid_stochastic.py

@@ -32,7 +32,7 @@ def sum_total_removals(sim):
 
 
 def final_size(sim):
-    remove = sim[:, 3, :]
+    remove = sim[:, :, 3]
     remove = remove.reshape([remove.shape[0], 152, 17])
     fs = remove[-1, :, :].sum(axis=0)
     return fs
@@ -186,7 +186,7 @@ if __name__ == '__main__':
     seeding = seed_areas(N, n_names)  # Seed 40-44 age group, 30 seeds by popn size
     state_init = model.create_initial_state(init_matrix=seeding)
 
-    with tf.device('/CPU:0'):
+    with tf.device('CPU'):
         start = time.perf_counter()
         t, sim = model.simulate(param, state_init)
         end = time.perf_counter()

ode_config.yaml

@@ -18,7 +18,7 @@ parameter:
 settings:
   inference_period:
     - 2020-02-19
-    - 2020-09-01
+    - 2020-04-01
   holiday:
     - 2020-03-23
     - 2020-10-01