Commit 750c80c7 authored by Chris Jewell's avatar Chris Jewell
Browse files

Implemented country-wide ODE model.

parent 3e5268ca
......@@ -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
"""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
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
# 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
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