Implemented Tier 2 and 3 model

covid/data.py

@@ -141,6 +141,7 @@ def read_tier_restriction_data(
     tier_restriction_csv, lad19cd_lookup, date_low, date_high
 ):
     data = pd.read_csv(tier_restriction_csv)
+    data.loc[:, "date"] = pd.to_datetime(data["date"])
 
     # Group merged ltlas
     london = ["City of London", "Westminster"]
@@ -166,6 +167,7 @@ def read_tier_restriction_data(
     lad19cd = lad19cd_lookup["lad19cd"].sort_values().unique()
     new_index = pd.MultiIndex.from_product([dates, lad19cd])
     data = data.reindex(new_index, fill_value=0.0)
+    warn(f"Tier summary: {np.mean(data, axis=0)}")
 
     # Pack into [T, M, V] array.
     arr_data = data.to_xarray().to_array()

inference.py

@@ -110,10 +110,10 @@ if __name__ == "__main__":
     def logp(block0, block1, events):
         return model.log_prob(
             dict(
-                beta1=block1[0],
                 beta2=block0[0],
-                beta3=block1[1:3],
                 gamma=block0[1],
+                beta1=block1[0],
+                beta3=block1[1:3],
                 xi=block1[3:],
                 seir=events,
             )
@@ -234,7 +234,7 @@ if __name__ == "__main__":
         return recurse(f, results)
 
     # Build MCMC algorithm here.  This will be run in bursts for memory economy
-    @tf.function  # (autograph=False, experimental_compile=True)
+    @tf.function(autograph=False, experimental_compile=True)
     def sample(n_samples, init_state, previous_results=None):
         with tf.name_scope("main_mcmc_sample_loop"):
 
@@ -320,10 +320,14 @@ if __name__ == "__main__":
 
     output_results = [
         posterior.create_dataset(
-            "results/theta", (NUM_SAVED_SAMPLES, 3), dtype=DTYPE,
+            "results/theta",
+            (NUM_SAVED_SAMPLES, 3),
+            dtype=DTYPE,
         ),
         posterior.create_dataset(
-            "results/xi", (NUM_SAVED_SAMPLES, 3), dtype=DTYPE,
+            "results/xi",
+            (NUM_SAVED_SAMPLES, 3),
+            dtype=DTYPE,
         ),
         posterior.create_dataset(
             "results/move/S->E",

model_spec.py

@@ -86,7 +86,7 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
         return tfd.Sample(
             tfd.Normal(
                 loc=tf.constant(0.0, dtype=DTYPE),
-                scale=tf.constant(1000.0, dtype=DTYPE),
+                scale=tf.constant(100.0, dtype=DTYPE),
             ),
             sample_shape=2,
         )
@@ -176,17 +176,17 @@ def CovidUK(covariates, initial_state, initial_step, num_steps, priors):
 
 
 def next_generation_matrix_fn(covar_data, param):
-    """The next generation matrix calculates the force of infection from 
-       individuals in metapopulation i to all other metapopulations j during
-       a typical infectious period (1/gamma). i.e.
-         
-         \[ A_{ij} = S_j * \beta_1 ( 1 + \beta_2 * w_t * C_{ij} / N_i) / N_j / gamma \]
-
-       :param covar_data: a dictionary of covariate data
-       :param param: a dictionary of parameters
-       :returns: a function taking arguments `t` and `state` giving the time and
-                 epidemic state (SEIR) for which the NGM is to be calculated.  This
-                 function in turn returns an MxM next generation matrix.
+    """The next generation matrix calculates the force of infection from
+    individuals in metapopulation i to all other metapopulations j during
+    a typical infectious period (1/gamma). i.e.
+
+      \[ A_{ij} = S_j * \beta_1 ( 1 + \beta_2 * w_t * C_{ij} / N_i) / N_j / gamma \]
+
+    :param covar_data: a dictionary of covariate data
+    :param param: a dictionary of parameters
+    :returns: a function taking arguments `t` and `state` giving the time and
+              epidemic state (SEIR) for which the NGM is to be calculated.  This
+              function in turn returns an MxM next generation matrix.
     """
 
     def fn(t, state):

pyproject.toml

@@ -17,7 +17,10 @@ xlrd = "^1.2.0"
 tqdm = "^4.50.2"
 openpyxl = "^3.0.5"
 h5py = "^2.10.0"
+tf-nightly = "2.4.0.dev20201021"
 gemlib = {git = "http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib.git"}
+xarray = "^0.16.1"
+seaborn = "^0.11.0"
 
 [tool.poetry.dev-dependencies]
 ipython = "^7.18.1"