Implemented country-wide ODE model.

covid/model.py

@@ -2,11 +2,13 @@
 
 import tensorflow as tf
 import tensorflow_probability as tfp
+
 tode = tfp.math.ode
 import numpy as np
 
 from covid.impl.chainbinom_simulate import chain_binomial_simulate
 
+
 class CovidUK:
     def __init__(self, K, T, W):
         self.K = K
@@ -38,17 +40,16 @@ class CovidUK:
 
 class Homogeneous:
     def __init__(self):
-
         self.stoichiometry = tf.constant([[-1, 1, 0, 0],
-                                         [0, -1, 1, 0],
-                                         [0, 0, -1, 1]], dtype=tf.float32)
+                                          [0, -1, 1, 0],
+                                          [0, 0, -1, 1]], dtype=tf.float32)
 
     def h(self, state):
         state = tf.unstack(state, axis=0)
         S, E, I, R = state
 
         hazard_rates = tf.stack([
-            self.param['beta'] * I/tf.reduce_sum(state),
+            self.param['beta'] * I / tf.reduce_sum(state),
             self.param['nu'] * tf.ones_like(I),
             self.param['gamma'] * tf.ones_like(I)
         ])
@@ -63,7 +64,7 @@ class Homogeneous:
 
 class CovidUKODE:
 
-    def __init__(self, K, T, N, n_lads=152):
+    def __init__(self, K, T, N, n_lads):
         """Represents a CovidUK ODE model
 
         :param K: a MxM matrix of age group mixing
@@ -71,23 +72,26 @@ class CovidUKODE:
         :param N: a vector of population sizes in each LAD
         :param n_lads: the number of LADS
         """
-        K = tf.linalg.LinearOperatorFullMatrix(K)
-        eye = tf.linalg.LinearOperatorDiag(n_lads)
-        self.K = tf.linalg.LinearOperatorKronecker([K, eye])
+        self.Kbar = tf.reduce_mean(K)
+        self.K = tf.linalg.LinearOperatorFullMatrix(K)
+        eye = tf.linalg.LinearOperatorIdentity(n_lads)
+        self.K = tf.linalg.LinearOperatorKronecker([self.K, eye])
 
-        T = tf.linalg.LinearOperatorFullMatrix(T)
-        shp = tf.linalg.LinearOperatorFullMatrix(np.ones([n_lads, n_lads]))
-        self.T = tf.linalg.LinearOperatorKronecker([T, shp])
+        self.T = tf.linalg.LinearOperatorFullMatrix(T)
+        shp = tf.linalg.LinearOperatorFullMatrix(np.ones_like(K, dtype=np.float32))
+        self.T = tf.linalg.LinearOperatorKronecker([self.T, shp])
+
+        self.N = N
+        self.n_lads = n_lads
 
     def make_h(self, param):
-        @tf.function
-        def h_fn(state):
+        def h_fn(t, state):
             state = tf.unstack(state, axis=0)
             S, E, I, R = state
 
-            infec_rate = tf.matvec(self.K, I)
-            infec_rate += tf.reduce_mean(self.K) * tf.matvec(self.T, I/self.N)
-            infec_rate  = infec_rate * S / self.N
+            infec_rate = tf.linalg.matvec(self.K, I)
+            infec_rate += self.Kbar * tf.linalg.matvec(self.T, I / self.N)
+            infec_rate = param['beta'] * S / self.N * infec_rate
 
             dS = -infec_rate
             dE = infec_rate - param['nu'] * E
@@ -96,12 +100,22 @@ class CovidUKODE:
 
             df = tf.stack([dS, dE, dI, dR])
             return df
+
         return h_fn
 
+    def create_initial_state(self, init_matrix=None):
+        if init_matrix is None:
+            I = np.ones(self.N.shape, dtype=np.float32)
+            S = self.N - I
+            E = np.zeros(self.N.shape, dtype=np.float32)
+            R = np.zeros(self.N.shape, dtype=np.float32)
+            return np.stack([S, E, I, R])
+
     @tf.function
     def simulate(self, param, state_init, t_start, t_end, t_step=1.):
         h = self.make_h(param)
-        t = np.linspace(t_start, t_end, t_step)
-        sim = tode.DormandPrince(first_step_size=1., max_num_steps=5000).solve(h, t_start, state_init,
-                                                                               solution_times=t)
-        return sim
\ No newline at end of file
+        t = np.arange(start=t_start, stop=t_end, step=t_step)
+        print(f"Running simulation with times {t_start}-{t_end}:{t_step}")
+        solver = tode.DormandPrince()
+        results = solver.solve(ode_fn=h, initial_time=t_start, initial_state=state_init, solution_times=t)
+        return results.times, results.states

covid/rdata.py

+"""Loads R data.frame data structures"""
+
+import numpy
+import pyreadr as pyr
+
+
+def load_age_mixing(rds_file: str):
+    """Loads age mixing matrix from R.
+
+    :param rds_file: a .rds file containing an R data.frame with mixing matrix
+    """
+    raw = pyr.read_r(rds_file)
+    return list(raw.values())[0]
+
+
+def load_mobility_matrix(rds_file: str):
+    """Loads mobility COO from RDS file.
+
+     :param rds_file: a .rds file containing an R data.frame with columns 'Residence',
+     and 'Workplace' indicating matrix coordinates, and first column containing the value
+     """
+    raw = pyr.read_r(rds_file)
+    df = list(raw.values())[0]
+    colnames = df.columns
+    mobility_matrix = df.pivot(index='Workplace', columns='Residence', values=colnames[0])
+    mobility_matrix[mobility_matrix.isna()] = 0.
+    return mobility_matrix
+
+
+def load_population(rds_file: str):
+    """Loads population data from RDS file.
+
+    :param rds_file: and RDS file containing a data.frame with columns 'age', 'LA.code', 'n'
+    """
+    raw = pyr.read_r(rds_file)
+    df = list(raw.values())[0]
+    df = df.sort_values(by=['LA.code', 'age'])
+    return df

covid_ode.py

+import optparse
 import numpy as np
-import tensorflow as tf
-import tensorflow_probability as tfp
-tode = tfp.math.ode
+import yaml
+import time
+import matplotlib.pyplot as plt
 
+from covid.model import CovidUKODE
+from covid.rdata import *
 
-popsize = 2500
-state = np.array([np.full([popsize], 999.),
-                  np.full([popsize], 0.),
-                  np.full([popsize], 1.),
-                  np.full([popsize], 0.)], dtype=np.float64)
 
-K = np.random.uniform(size=[popsize, popsize])
+def sanitise_parameter(par_dict):
+    """Sanitises a dictionary of parameters"""
+    par = ['beta','nu','gamma']
+    d = {key: np.float64(par_dict[key]) for key in par}
+    return d
 
-param = {'beta': 0.0002, 'nu': 0.14, 'gamma': 0.14}
 
+def sanitise_settings(par_dict):
+    d = {'start': np.float64(par_dict['start']),
+         'end': np.float64(par_dict['end']),
+         'time_step': np.float64(par_dict['time_step'])}
+    return d
 
 
-def h(t, state):
-    print(state)
-    state = tf.unstack(state, axis=0)
-    S, E, I, R = state
+if __name__ == '__main__':
 
-    infec_rate = param['beta'] * S * tf.linalg.matvec(K, I)
-    dS = -infec_rate
-    dE = infec_rate - param['nu'] * E
-    dI = param['nu'] * E - param['gamma'] * I
-    dR = param['gamma'] * I
+    parser = optparse.OptionParser()
+    parser.add_option("--config", "-c", dest="config",
+                      help="configuration file")
+    options, args = parser.parse_args()
+    with open(options.config, 'r') as ymlfile:
+        config = yaml.load(ymlfile)
 
-    df = tf.stack([dS, dE, dI, dR])
-    return df
+    K = load_age_mixing(config['data']['age_mixing_matrix']).to_numpy(dtype=np.float32)
+    T = load_mobility_matrix(config['data']['mobility_matrix']).to_numpy(dtype=np.float32)
+    np.fill_diagonal(T, 0.)
+    N = load_population(config['data']['population_size'])['n'].to_numpy(dtype=np.float32)
 
-@tf.function
-def solve_ode(rates, t_init, state_init, t):
-    return tode.DormandPrince(first_step_size=1., max_num_steps=5000).solve(rates, t_init, state_init, solution_times=t)
+    param = sanitise_parameter(config['parameter'])
+    settings = sanitise_settings(config['settings'])
 
-solution_times = np.arange(0., 365., 1.)
+    model = CovidUKODE(K, T, N, T.shape[0])
 
-print("Running...", flush=True)
-result = solve_ode(rates=h, t_init=0., state_init=state, t=solution_times)
-print("Done")
-print(result)
\ No newline at end of file
+    state_init = model.create_initial_state()
+    start = time.perf_counter()
+    t, sim = model.simulate(param, state_init,
+                            settings['start'], settings['end'],
+                            settings['time_step'])
+    end = time.perf_counter()
+    print(f"Complete in {end-start} seconds")
+    plt.plot(t, sim[:, 2, 0], 'r-')
+    plt.show()
\ No newline at end of file

ode_config.yaml

+# Covid_ODE model configuration
+
+data:
+  age_mixing_matrix: data/polymod_normal_df.rds
+  mobility_matrix: data/movement.rds
+  population_size: data/pop.rds
+
+parameter:
+  beta: 0.09
+  nu: 0.25
+  gamma: 0.65
+
+settings:
+  start: 0.
+  end: 365.
+  time_step: 1.
+
+output:
+  simulation: covid_ode.sim