Commit 3e5268ca authored by Chris Jewell's avatar Chris Jewell
Browse files

Added ODE model.

parent 7c8712ef
"""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
......@@ -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
......@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment