Bugfix to naming in thin_samples.pkl

CHANGES:

1. `thin_samples['init_state']` --> `thin_samples['initial_state']`
2. `thin_samples['events']` --> `thin_samples['seir']`
3. add time dimension to next generation matrix
4. update downstream scripts to changes.

covid/tasks/inference.py

@@ -276,10 +276,9 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
             "gamma0": samples[0][:, 1],
             "gamma1": samples[0][:, 2],
             "sigma": samples[0][:, 3],
-            "beta3": tf.zeros([1, 5], dtype=DTYPE),
             "beta1": samples[1][:, 0],
             "xi": samples[1][:, 1:],
-            "events": samples[2],
+            "seir": samples[2],
         },
         results_dict=results,
         num_samples=NUM_SAVED_SAMPLES,
@@ -309,12 +308,9 @@ def mcmc(data_file, output_file, config, use_autograph=False, use_xla=True):
                 "gamma0": samples[0][:, 1],
                 "gamma1": samples[0][:, 2],
                 "sigma": samples[0][:, 3],
-                "beta3": tf.zeros(
-                    [samples[0].shape[0], 5], dtype=DTYPE
-                ),  # samples[0][:, 4:],
                 "beta1": samples[1][:, 0],
                 "xi": samples[1][:, 1:],
-                "events": samples[2],
+                "seir": samples[2],
             },
             first_dim_offset=i * NUM_BURST_SAMPLES,
         )

covid/tasks/next_generation_matrix.py

@@ -5,13 +5,12 @@ import numpy as np
 import xarray
 import tensorflow as tf
 
-
 from covid import model_spec
 from covid.util import copy_nc_attrs
 from gemlib.util import compute_state
 
 
-def calc_posterior_ngm(samples, covar_data):
+def calc_posterior_ngm(samples, initial_state, times, covar_data):
     """Calculates effective reproduction number for batches of metapopulations
     :param theta: a tensor of batched theta parameters [B] + theta.shape
     :param xi: a tensor of batched xi parameters [B] + xi.shape
@@ -20,21 +19,26 @@ def calc_posterior_ngm(samples, covar_data):
     :param covar_data: the covariate data
     :return a batched vector of R_it estimates
     """
+    times = tf.convert_to_tensor(times)
 
     def r_fn(args):
 
         par = tf.nest.pack_sequence_as(samples, args)
-        
-        t = par['seir'].shape[-2] - 1
+
         state = compute_state(
-            samples["init_state"], par['seir'], model_spec.STOICHIOMETRY
+            initial_state, par["seir"], model_spec.STOICHIOMETRY
         )
-        state = tf.gather(state, t, axis=-2)  # State on final inference day
+        del par["seir"]
 
-        del par['seir']
-        ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par)
-        ngm = ngm_fn(t, state)
-        return ngm
+        def fn(t):
+            state_ = tf.gather(
+                state, t, axis=-2
+            )  # State on final inference day
+            ngm_fn = model_spec.next_generation_matrix_fn(covar_data, par)
+            ngm = ngm_fn(t, state_)
+            return ngm
+
+        return tf.vectorized_map(fn, elems=times)
 
     return tf.vectorized_map(
         r_fn,
@@ -49,16 +53,24 @@ def next_generation_matrix(input_files, output_file):
     with open(input_files[1], "rb") as f:
         samples = pkl.load(f)
 
+    initial_state = samples["initial_state"]
+    del samples["initial_state"]
+
+    times = [
+        samples["seir"].shape[-2] - 1,
+    ]
+
     # Compute ngm posterior
-    ngm = calc_posterior_ngm(samples, covar_data).numpy()
+    ngm = calc_posterior_ngm(samples, initial_state, times, covar_data)
     ngm = xarray.DataArray(
         ngm,
         coords=[
             np.arange(ngm.shape[0]),
+            covar_data.coords["time"][times],
             covar_data.coords["location"],
             covar_data.coords["location"],
         ],
-        dims=["iteration", "dest", "src"],
+        dims=["iteration", "time", "dest", "src"],
     )
     ngm = xarray.Dataset({"ngm": ngm})
 
@@ -73,21 +85,19 @@ if __name__ == "__main__":
 
     parser = ArgumentParser()
     parser.add_argument(
-        "-s",
-        "--samples",
+        "samples",
         type=str,
-        description="A pickle file with MCMC samples",
-        required=True,
+        help="A pickle file with MCMC samples",
     )
     parser.add_argument(
         "-d",
         "--data",
         type=str,
-        decription="A data glob pickle file",
-        require=True,
+        help="A data glob pickle file",
+        required=True,
     )
     parser.add_argument(
-        "-o", "--output", type=str, description="The output file", require=True
+        "-o", "--output", type=str, help="The output file", required=True
     )
     args = parser.parse_args()
 

covid/tasks/overall_rt.py

@@ -15,6 +15,7 @@ def overall_rt(next_generation_matrix, output_file):
     ngms = xarray.open_dataset(
         next_generation_matrix, group="posterior_predictive"
     )["ngm"]
+    ngms = ngms[:, 0, :, :].drop("time")
     b, _ = power_iteration(ngms)
     rt = rayleigh_quotient(ngms, b)
     q = np.arange(0.05, 1.0, 0.05)
@@ -30,10 +31,11 @@ if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument(
         "input_file",
-        description="The input .pkl file containing the next generation matrix",
+        type=str,
+        help="The input .pkl file containing the next generation matrix",
     )
     parser.add_argument(
-        "output_file", description="The name of the output .xlsx file"
+        "output_file", type=str, help="The name of the output .xlsx file"
     )
 
     args = parser.parse_args()

covid/tasks/predict.py

@@ -9,8 +9,10 @@ from covid import model_spec
 from covid.util import copy_nc_attrs
 from gemlib.util import compute_state
 
-@tf.function
-def predicted_incidence(posterior_samples, covar_data, init_step, num_steps):
+
+def predicted_incidence(
+    posterior_samples, init_state, covar_data, init_step, num_steps
+):
     """Runs the simulation forward in time from `init_state` at time `init_time`
        for `num_steps`.
     :param param: a dictionary of model parameters
@@ -22,33 +24,36 @@ def predicted_incidence(posterior_samples, covar_data, init_step, num_steps):
     """
 
     posterior_state = compute_state(
-        posterior_samples["init_state"],
+        init_state,
         posterior_samples["seir"],
         model_spec.STOICHIOMETRY,
     )
-    posterior_samples['init_state_'] = posterior_state[..., init_step, :]
-    del posterior_samples['seir']
-    
-    def sim_fn(args):
-        par = tf.nest.pack_sequence_as(posterior_samples, args)
-        init_ = par['init_state_']
-        del par['init_state_']
-        
-        model = model_spec.CovidUK(
-            covar_data,
-            initial_state=init_,
-            initial_step=init_step,
-            num_steps=num_steps,
+    posterior_samples["new_init_state"] = posterior_state[..., init_step, :]
+    del posterior_samples["seir"]
+
+    @tf.function
+    def do_sim():
+        def sim_fn(args):
+            par = tf.nest.pack_sequence_as(posterior_samples, args)
+            init_ = par["new_init_state"]
+            del par["new_init_state"]
+
+            model = model_spec.CovidUK(
+                covar_data,
+                initial_state=init_,
+                initial_step=init_step,
+                num_steps=num_steps,
+            )
+            sim = model.sample(**par)
+            return sim["seir"]
+
+        return tf.map_fn(
+            sim_fn,
+            elems=tf.nest.flatten(posterior_samples),
+            fn_output_signature=(tf.float64),
         )
-        sim = model.sample(**par)
-        return sim["seir"]
 
-    events = tf.map_fn(
-        sim_fn,
-        elems=tf.nest.flatten(posterior_samples),
-        fn_output_signature=(tf.float64),
-    )
-    return init_state, events
+    return posterior_samples["new_init_state"], do_sim()
 
 
 def read_pkl(filename):
@@ -60,7 +65,10 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
 
     covar_data = xarray.open_dataset(data, group="constant_data")
     cases = xarray.open_dataset(data, group="observations")
+
     samples = read_pkl(posterior_samples)
+    initial_state = samples["initial_state"]
+    del samples["initial_state"]
 
     if initial_step < 0:
         initial_step = samples["seir"].shape[-2] + initial_step
@@ -73,7 +81,7 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
     )
 
     estimated_init_state, predicted_events = predicted_incidence(
-        samples, covar_data, initial_step, num_steps
+        samples, initial_state, covar_data, initial_step, num_steps
     )
 
     prediction = xarray.DataArray(
@@ -109,23 +117,21 @@ if __name__ == "__main__":
 
     parser = ArgumentParser()
     parser.add_argument(
-        "-i", "--initial-step", type=int, default=0, description="Initial step"
-    )
-    parser.add_argument(
-        "-n", "--num-steps", type=int, default=1, description="Number of steps"
+        "-i", "--initial-step", type=int, default=0, help="Initial step"
     )
     parser.add_argument(
-        "data_pkl", type=str, description="Covariate data pickle"
+        "-n", "--num-steps", type=int, default=1, help="Number of steps"
     )
+    parser.add_argument("data_pkl", type=str, help="Covariate data pickle")
     parser.add_argument(
         "posterior_samples_pkl",
         type=str,
-        description="Posterior samples pickle",
+        help="Posterior samples pickle",
     )
     parser.add_argument(
         "output_file",
         type=str,
-        description="Output pkl file",
+        help="Output pkl file",
     )
     args = parser.parse_args()
 

covid/tasks/summarize.py

@@ -22,7 +22,7 @@ def rt(input_file, output_file):
 
     ngm = xarray.open_dataset(input_file, group="posterior_predictive")["ngm"]
 
-    rt = np.sum(ngm, axis=-2)
+    rt = ngm.sum(dim="dest").isel(time=-1).drop("time")
     rt_summary = mean_and_ci(rt, name="Rt")
     exceed = np.mean(rt > 1.0, axis=0)
 

covid/tasks/thin_posterior.py

@@ -6,24 +6,12 @@ import pickle as pkl
 
 def thin_posterior(input_file, output_file, config):
 
-    thin_idx = range(config['start'], config['end'], config['by'])
+    idx = slice(config["start"], config["end"], config["by"])
 
     f = h5py.File(input_file, "r", rdcc_nbytes=1024 ** 3, rdcc_nslots=1e6)
-    output_dict = dict(
-        beta1=f["samples/beta1"][thin_idx],
-        beta2=f["samples/beta2"][thin_idx],
-        beta3=f["samples/beta3"][
-            thin_idx,
-        ],
-        sigma=f["samples/sigma"][
-            thin_idx,
-        ],
-        xi=f["samples/xi"][thin_idx],
-        gamma0=f["samples/gamma0"][thin_idx],
-        gamma1=f["samples/gamma1"][thin_idx],
-        seir=f["samples/events"][thin_idx],
-        init_state=f["initial_state"][:],
-    )
+
+    output_dict = {k: v[idx] for k, v in f["samples"].items()}
+    output_dict["initial_state"] = f["initial_state"][:]
     f.close()
 
     with open(output_file, "wb") as f:
@@ -36,26 +24,20 @@ def thin_posterior(input_file, output_file, config):
 if __name__ == "__main__":
 
     import yaml
-    import os
-    from covid.cli_arg_parse import cli_args
-
-    args = cli_args()
-    with open(args.config, "r") as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+    import argparse
 
-    def join_and_expand(prefix, filename):
-        return os.path.expandvars(os.path.join(prefix, filename))
-
-    input_file = join_and_expand(
-        config["output"]["results_dir"], config["output"]["posterior"]
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c", "--config", type=str, help="Configuration file", required=True
     )
-    output_file = join_and_expand(
-        config["output"]["results_dir"], config["ThinPosterior"]["output_file"]
+    parser.add_argument(
+        "-o", "--output", type=str, help="Output pkl file", required=True
     )
+    parser.add_argument("samples", type=str, help="Posterior HDF5 file")
+    args = parser.parse_args()
 
-    thin_idx = range(
-        config["ThinPosterior"]["thin_start"],
-        config["ThinPosterior"]["thin_end"],
-        config["ThinPosterior"]["thin_by"],
-    )
-    thin_posterior(input_file, output_file, thin_idx)
+    with open(args.config, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    print("Config: ", config["ThinPosterior"])
+    thin_posterior(args.samples, args.output, config["ThinPosterior"])

covid/tasks/within_between.py

@@ -70,7 +70,7 @@ def within_between(input_files, output_file):
 
     beta2 = samples["beta2"]
     events = samples["seir"]
-    init_state = samples["init_state"]
+    init_state = samples["initial_state"]
     state_timeseries = compute_state(
         init_state, events, model_spec.STOICHIOMETRY
     )

pyproject.toml

 [tool.poetry]
 name = "covid19uk"
-version = "0.5.0"
+version = "0.5.5"
 description = "Spatial stochastic SEIR analysis of COVID-19 in the UK"
 authors = ["Chris Jewell <c.jewell@lancaster.ac.uk>"]
 license = "MIT"
@@ -20,7 +20,7 @@ xlrd = "^1.2.0"
 tqdm = "^4.50.2"
 h5py = "^2.10.0"
 gemlib = {git = "http://fhm-chicas-code.lancs.ac.uk/GEM/gemlib.git"}
-xarray = "^0.16.1"
+xarray = {extras = ["netcdf4"], version = "^0.17.0"}
 seaborn = "^0.11.0"
 ruffus = "^2.8.4"
 tensorflow = "^2.4.0"