Create a stochastic version of the Covid model CovidUKStochastic. See covid_stochastic.py.

covid/impl/chainbinom_simulate.py

@@ -40,11 +40,11 @@ def chain_binomial_simulate(hazard_fn, state, start, end, time_step):
     output = tf.TensorArray(tf.float64, size=times.shape[0])
     output = output.write(0, state)
 
-    for i in tf.range(1, times.shape[0]):
+    for i in range(1, times.shape[0]):
         state = propagate(i, state)
         output = output.write(i, state)
 
-    sim = output.gather(tf.range(times.shape[0]))
+    sim = output.gather(tf.range(1, times.shape[0]))
     return times, sim
 
 

covid/model.py

@@ -216,7 +216,7 @@ class CovidUKStochastic(CovidUK):
             return rates
         return h
 
-    @tf.function
+    @tf.function(experimental_compile=True)
     def simulate(self, param, state_init):
         """Runs a simulation from the epidemic model
 

covid_stochastic.py

@@ -179,25 +179,20 @@ if __name__ == '__main__':
                               C=C,
                               N=N,
                               W=W,
-                              date_range=settings['prediction_period'],
+                              date_range=settings['inference_period'],
                               holidays=settings['holiday'],
                               time_step=1.)
 
     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'):
+    with tf.device('/CPU:0'):
         start = time.perf_counter()
         t, sim = model.simulate(param, state_init)
         end = time.perf_counter()
         print(f'Complete in {end - start} seconds')
 
-
     # Plotting functions
-    dates = settings['start'] + t.numpy().astype(np.timedelta64)
-    dt = doubling_time(dates, sim.numpy(), '2020-03-01', '2020-03-31')
-    print(f"Doubling time: {dt}")
-
     fig_attack = plt.figure()
     fig_uk = plt.figure()
 

ode_config.yaml

@@ -9,7 +9,7 @@ data:
   reported_cases: data/DailyConfirmedCases_2020-03-20.csv
 
 parameter:
-  beta1: 0.04  # R0 2.4
+  beta1: 0.1  # R0 2.4
   beta2: 0.33  # Contact with commuters 1/3rd of the time
   omega: 1.0
   nu: 0.25
@@ -18,7 +18,7 @@ parameter:
 settings:
   inference_period:
     - 2020-02-19
-    - 2020-04-01
+    - 2020-09-01
   holiday:
     - 2020-03-23
     - 2020-10-01