Added ODE model.

covid/model.py

 """Functions for infection rates"""
 
 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
 
@@ -56,3 +59,49 @@ class Homogeneous:
         self.param = param
         return chain_binomial_simulate(self.h, initial_state, time_lims[0],
                                        time_lims[1], 1., self.stoichiometry)
+
+
+class CovidUKODE:
+
+    def __init__(self, K, T, N, n_lads=152):
+        """Represents a CovidUK ODE model
+
+        :param K: a MxM matrix of age group mixing
+        :param T: a n_ladsxn_lads matrix of inter-LAD connectivity
+        :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])
+
+        T = tf.linalg.LinearOperatorFullMatrix(T)
+        shp = tf.linalg.LinearOperatorFullMatrix(np.ones([n_lads, n_lads]))
+        self.T = tf.linalg.LinearOperatorKronecker([T, shp])
+
+    def make_h(self, param):
+        @tf.function
+        def h_fn(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
+
+            dS = -infec_rate
+            dE = infec_rate - param['nu'] * E
+            dI = param['nu'] * E - param['gamma'] * I
+            dR = param['gamma'] * I
+
+            df = tf.stack([dS, dE, dI, dR])
+            return df
+        return h_fn
+
+    @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

covid_ode.py

@@ -12,15 +12,16 @@ state = np.array([np.full([popsize], 999.),
 
 K = np.random.uniform(size=[popsize, popsize])
 
-param = {'beta': 0.2, 'nu': 0.14, 'gamma': 0.14}
+param = {'beta': 0.0002, 'nu': 0.14, 'gamma': 0.14}
+
+
 
-@tf.function
 def h(t, state):
+    print(state)
     state = tf.unstack(state, axis=0)
     S, E, I, R = state
 
     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
@@ -31,9 +32,11 @@ def h(t, state):
 
 @tf.function
 def solve_ode(rates, t_init, state_init, t):
-    return tode.DormandPrince().solve(rates, t_init, state_init, solution_times=t)
+    return tode.DormandPrince(first_step_size=1., max_num_steps=5000).solve(rates, t_init, state_init, solution_times=t)
 
 solution_times = np.arange(0., 365., 1.)
 
+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

scratch_covid.py

@@ -11,12 +11,12 @@ state = np.array([np.full([popsize], 999.),
                   np.full([popsize], 0.),
                   np.full([popsize], 1.),
                   np.full([popsize], 0.)], dtype=np.float32)
-print(state)
+
 print("Running...", flush=True, sep='')
 tf.summary.trace_on(graph=True, profiler=True)
 
 t, sim_state = model.sample(state,
-                            [0., 5.],
+                            [0., 10.],
                             {'beta': 0.2, 'nu': 0.14, 'gamma':0.14})
 print("Done", flush=True)
 print(sim_state)