Commit 996a4afc authored by Chris Jewell's avatar Chris Jewell
Browse files

Implemented household sim

parent 750c80c7
......@@ -31,7 +31,7 @@ def chain_binomial_simulate(hazard_fn, state, start, end, time_step, stoichiomet
propagate = chain_binomial_propagate(hazard_fn, time_step, stoichiometry)
times = tf.range(start, end, time_step)
output = tf.TensorArray(tf.float32, size=times.shape[0])
output = tf.TensorArray(tf.float64, size=times.shape[0])
output = output.write(0, state)
for i in tf.range(1, times.shape[0]):
......
......@@ -8,6 +8,25 @@ import numpy as np
from covid.impl.chainbinom_simulate import chain_binomial_simulate
def power_iteration(A, tol=1e-3):
b_k = tf.random.normal([A.shape[1], 1])
epsilon = 1.
i = 0
while tf.greater(epsilon, tol):
b_k1 = tf.matmul(A, b_k)
b_k1_norm = tf.linalg.norm(b_k1)
b_k_new = b_k1 / b_k1_norm
epsilon = tf.reduce_sum(tf.pow(b_k_new-b_k, 2))
b_k = b_k_new
i += 1
return b_k, i
#@tf.function
def rayleigh_quotient(A, b):
b = tf.reshape(b, [b.shape[0], 1])
numerator = tf.matmul(tf.transpose(b), tf.matmul(A, b))
denominator = tf.matmul(tf.transpose(b), b)
return numerator / denominator
class CovidUK:
def __init__(self, K, T, W):
......@@ -19,19 +38,18 @@ class CovidUK:
[0, -1, 1, 0],
[0, 0, -1, 1]]
@tf.function
def h(self, state):
state = tf.unstack(state, axis=0)
S, E, I, R = state
hazard_rates = tf.stack([
self.param['beta1'] * tf.dot(self.T, tf.dot(self.K, I)),
self.param['beta1'] * tf.dot(self.T, tf.dot(self.K, I))/self.K.shape[0],
self.param['nu'],
self.param['gamma']
])
return hazard_rates
@tf.function
def sample(self, initial_state, time_lims, param):
self.param = param
return chain_binomial_simulate(self.h, initial_state, time_lims[0],
......@@ -64,34 +82,42 @@ class Homogeneous:
class CovidUKODE:
def __init__(self, K, T, N, n_lads):
def __init__(self, K, T, N):
"""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 K: a MxM matrix of age group mixing in term time
:param T: a n_ladsxn_lads matrix of inter-LAD commuting
:param N: a vector of population sizes in each LAD
:param n_lads: the number of LADS
"""
self.n_ages = K.shape[0]
self.n_lads = T.shape[0]
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])
eye = tf.linalg.LinearOperatorIdentity(self.n_lads)
self.K = tf.linalg.LinearOperatorKronecker([eye, self.K])
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
self.N = tf.constant(N, dtype=tf.float32)
N_matrix = tf.reshape(self.N, [self.n_lads, self.n_ages])
N_sum = tf.reduce_sum(N_matrix, axis=1)
N_sum = N_sum[:, None] * tf.ones([1, self.n_ages])
self.N_sum = tf.reshape(N_sum, [-1])
def make_h(self, param):
def h_fn(t, state):
state = tf.unstack(state, axis=0)
S, E, I, R = state
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
infec_rate = param['beta1'] * tf.linalg.matvec(self.K, I)
infec_rate += param['beta1'] * param['beta2'] * self.Kbar * tf.linalg.matvec(self.T, I / self.N_sum)
infec_rate = S / self.N * infec_rate
dS = -infec_rate
dE = infec_rate - param['nu'] * E
......@@ -105,17 +131,37 @@ class CovidUKODE:
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])
I = np.zeros(self.N.shape, dtype=np.float32)
I[149*17+10] = 30. # Middle-aged in Surrey
else:
np.testing.assert_array_equal(init_matrix.shape, [self.n_lads, self.n_ages],
err_msg=f"init_matrix does not have shape [<num lads>,<num ages>] \
({self.n_lads},{self.n_ages})")
I = init_matrix.flatten()
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.arange(start=t_start, stop=t_end, step=t_step)
print(f"Running simulation with times {t_start}-{t_end}:{t_step}")
t0 = 0.
t1 = (t_end - t_start) / np.timedelta64(1, 'D')
t = np.arange(start=t0, stop=t1, step=t_step)[1:]
solver = tode.DormandPrince()
results = solver.solve(ode_fn=h, initial_time=t_start, initial_state=state_init, solution_times=t)
results = solver.solve(ode_fn=h, initial_time=t0, initial_state=state_init, solution_times=t)
return results.times, results.states
def ngm(self, param):
infec_rate = param['beta1'] * self.K.to_dense()
infec_rate += param['beta1'] * param['beta2'] * self.Kbar * self.T.to_dense() / self.N_sum[None, :]
ngm = infec_rate / param['gamma']
return ngm
def eval_R0(self, param, tol=1e-8):
ngm = self.ngm(param)
# Dominant eigen value by power iteration
dom_eigen_vec, i = power_iteration(ngm, tol=tol)
R0 = rayleigh_quotient(ngm, dom_eigen_vec)
return tf.squeeze(R0), i
......@@ -2,6 +2,7 @@
import numpy
import pyreadr as pyr
import numpy as np
def load_age_mixing(rds_file: str):
......@@ -10,7 +11,9 @@ def load_age_mixing(rds_file: str):
: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]
K = list(raw.values())[0]
age_groups = K.columns
return K.to_numpy(dtype=np.float32), age_groups
def load_mobility_matrix(rds_file: str):
......@@ -24,7 +27,7 @@ def load_mobility_matrix(rds_file: str):
colnames = df.columns
mobility_matrix = df.pivot(index='Workplace', columns='Residence', values=colnames[0])
mobility_matrix[mobility_matrix.isna()] = 0.
return mobility_matrix
return mobility_matrix.to_numpy(dtype=np.float32), mobility_matrix.index.to_numpy()
def load_population(rds_file: str):
......@@ -35,4 +38,4 @@ def load_population(rds_file: str):
raw = pyr.read_r(rds_file)
df = list(raw.values())[0]
df = df.sort_values(by=['LA.code', 'age'])
return df
return df['n'].to_numpy(dtype=np.float32), df[['name', 'Area.name.2']]
import optparse
import numpy as np
import yaml
import time
import h5py
import matplotlib.pyplot as plt
import tensorflow as tf
import yaml
from covid.model import CovidUKODE
from covid.rdata import *
......@@ -10,18 +12,149 @@ from covid.rdata import *
def sanitise_parameter(par_dict):
"""Sanitises a dictionary of parameters"""
par = ['beta','nu','gamma']
par = ['beta1', 'beta2', 'nu','gamma']
d = {key: np.float64(par_dict[key]) for key in par}
return d
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'])}
d = {'start': np.datetime64(par_dict['start']),
'end': np.datetime64(par_dict['end']),
'time_step': float(par_dict['time_step']),
'holiday': np.array([np.datetime64(date) for date in par_dict['holiday']])}
return d
def seed_areas(N, names, age_group=8, num_la=152, num_age=17, n_seed=30.):
areas = ['Inner London',
'Outer London',
'West Midlands (Met County)',
'Greater Manchester (Met County)']
names_matrix = names['Area.name.2'].to_numpy().reshape([num_la, num_age])
seed_areas = np.in1d(names_matrix[:, age_group], areas)
N_matrix = N.reshape([num_la, num_age]) # LA x Age
pop_size_sub = N_matrix[seed_areas, age_group] # Gather
n = np.round(n_seed * pop_size_sub / pop_size_sub.sum())
seeding = np.zeros_like(N_matrix)
seeding[seed_areas, age_group] = n # Scatter
return seeding
def sum_age_groups(sim):
infec = sim[:, 2, :]
infec = infec.reshape([infec.shape[0], 152, 17])
infec_uk = infec.sum(axis=2)
return infec_uk
def sum_la(sim):
infec = sim[:, 2, :]
infec = infec.reshape([infec.shape[0], 152, 17])
infec_uk = infec.sum(axis=1)
return infec_uk
def sum_total_removals(sim):
remove = sim[:, 3, :]
return remove.sum(axis=1)
def final_size(sim):
remove = sim[:, 3, :]
remove = remove.reshape([remove.shape[0], 152, 17])
fs = remove[-1, :, :].sum(axis=0)
return fs
def write_hdf5(filename, param, t, sim):
with h5py.File(filename, "w") as f:
dset_sim = f.create_dataset("simulation", sim.shape, dtype='f')
dset_sim[:] = sim
dset_t = f.create_dataset("time", t.shape, dtype='f')
dset_t[:] = t
grp_param = f.create_group("parameter")
for k, v in param.items():
d_beta = grp_param.create_dataset(k, [1], dtype='f')
d_beta[()] = v
def plot_total_curve(sim):
infec_uk = sum_la(sim)
infec_uk = infec_uk.sum(axis=1)
removals = sum_total_removals(sim)
times = np.datetime64('2020-02-20') + np.arange(removals.shape[0])
plt.plot(times, infec_uk, 'r-', label='Infected')
plt.plot(times, removals, 'b-', label='Removed')
plt.title('UK total cases')
plt.xlabel('Date')
plt.ylabel('Num infected or removed')
plt.legend()
def plot_by_age(sim, labels, t0=np.datetime64('2020-02-20'), ax=None):
if ax is None:
ax = plt.figure().gca()
infec_uk = sum_la(sim)
total_uk = infec_uk.mean(axis=1)
t = t0 + np.arange(infec_uk.shape[0])
colours = plt.cm.viridis(np.linspace(0., 1., infec_uk.shape[1]))
for i in range(infec_uk.shape[1]):
ax.plot(t, infec_uk[:, i], 'r-', alpha=0.4, color=colours[i], label=labels[i])
ax.plot(t, total_uk, '-', color='black', label='Mean')
return ax
def plot_by_la(sim, labels,t0=np.datetime64('2020-02-20'), ax=None):
if ax is None:
ax = plt.figure().gca()
infec_uk = sum_age_groups(sim)
total_uk = infec_uk.mean(axis=1)
t = t0 + np.arange(infec_uk.shape[0])
colours = plt.cm.viridis(np.linspace(0., 1., infec_uk.shape[1]))
for i in range(infec_uk.shape[1]):
ax.plot(t, infec_uk[:, i], 'r-', alpha=0.4, color=colours[i], label=labels[i])
ax.plot(t, total_uk, '-', color='black', label='Mean')
return ax
def draw_figs(sim, N):
# Attack rate
N = N.reshape([152, 17]).sum(axis=0)
fs = final_size(sim)
attack_rate = fs / N
print("Attack rate:", attack_rate)
print("Overall attack rate: ", np.sum(fs)/np.sum(N))
# Total UK epidemic curve
plot_total_curve(sim)
plt.xticks(rotation=45, horizontalalignment="right")
plt.savefig('total_uk_curve.pdf')
plt.show()
# TotalUK epidemic curve by age-group
fig, ax = plt.subplots(1, 2, figsize=[24, 12])
plot_by_la(sim, la_names, ax=ax[0])
plot_by_age(sim, age_groups, ax=ax[1])
ax[1].legend()
plt.xticks(rotation=45, horizontalalignment="right")
fig.autofmt_xdate()
plt.savefig('la_age_infec_curves.pdf')
plt.show()
# Plot attack rate
plt.figure(figsize=[4, 2])
plt.plot(age_groups, attack_rate, 'o-')
plt.xticks(rotation=90)
plt.title('Age-specific attack rate')
plt.savefig('age_attack_rate.pdf')
plt.show()
if __name__ == '__main__':
parser = optparse.OptionParser()
......@@ -31,22 +164,52 @@ if __name__ == '__main__':
with open(options.config, 'r') as ymlfile:
config = yaml.load(ymlfile)
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)
K_tt, age_groups = load_age_mixing(config['data']['age_mixing_matrix_term'])
K_hh, _ = load_age_mixing(config['data']['age_mixing_matrix_hol'])
T, la_names = load_mobility_matrix(config['data']['mobility_matrix'])
np.fill_diagonal(T, 0.)
N = load_population(config['data']['population_size'])['n'].to_numpy(dtype=np.float32)
N, n_names = load_population(config['data']['population_size'])
param = sanitise_parameter(config['parameter'])
settings = sanitise_settings(config['settings'])
model = CovidUKODE(K, T, N, T.shape[0])
state_init = model.create_initial_state()
# Straight, no school closures
model_term = CovidUKODE(K_tt, T, N)
model_holiday = CovidUKODE(K_hh/2., T/2., N)
seeding = seed_areas(N, n_names) # Seed 40-44 age group, 30 seeds by popn size
state_init = model_term.create_initial_state(init_matrix=seeding)
print('R_term=', model_term.eval_R0(param))
print('R_holiday=', model_holiday.eval_R0(param))
# School holidays and closures
@tf.function
def simulate():
t0, sim_0 = model_term.simulate(param, state_init,
settings['start'], settings['holiday'][0],
settings['time_step'])
t1, sim_1 = model_holiday.simulate(param, sim_0[-1, :, :],
settings['holiday'][0], settings['holiday'][1],
settings['time_step'])
t2, sim_2 = model_term.simulate(param, sim_1[-1, :, :],
settings['holiday'][1], settings['end'],
settings['time_step'])
t = tf.concat([t0, t1, t2], axis=0)
sim = tf.concat([tf.expand_dims(state_init, axis=0), sim_0, sim_1, sim_2], axis=0)
return t, sim
start = time.perf_counter()
t, sim = model.simulate(param, state_init,
settings['start'], settings['end'],
settings['time_step'])
t, sim = simulate()
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
print(f'Complete in {end-start} seconds')
draw_figs(sim.numpy(), N)
if 'simulation' in config['output']:
write_hdf5(config['output']['simulation'], param, t, sim)
print(f"Complete in {end-start} seconds")
\ No newline at end of file
"""Household simulation for SPI-M"""
import time
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf
from covid.model import CovidUK
# Household size distribution in the UK, taken from ONS, 2019
hh_size_distr = pd.DataFrame({'size': [1, 2, 3, 4, 5, 6], 'freq': [6608, 8094, 3918, 3469, 1186, 463]})
# Draw households
def draw_households(n: int, dist: pd.DataFrame) -> np.array:
probs = dist['freq'] / np.sum(dist['freq'])
hh_sizes = np.random.choice(dist['size'], n, p=probs)
#hh_vec = np.repeat(np.arange(hh_sizes.shape[0]), hh_sizes)
return hh_sizes
# Simulator
class HHstochastic(CovidUK):
def __init__(self, N):
self.stoichiometry = tf.constant([[-1, 1, 0, 0],
[0, -1, 1, 0],
[0, 0, -1, 1]], dtype=tf.float64)
self.N = N
self.popsize = tf.reduce_sum(N)
def h(self, state):
state = tf.unstack(state, axis=0)
S, E, I, R = state
eye = tf.eye(I.shape[1], dtype=tf.float64)
infec_rate = tf.linalg.tensordot(I, eye, axes=[[1], [1]])
infec_rate /= self.N
infec_rate *= self.param['omega']
infec_rate += tf.reduce_sum(I, axis=1, keepdims=True) / self.popsize
infec_rate *= self.param['beta']
hazard_rates = tf.stack([
infec_rate,
tf.constant(np.full(S.shape, self.param['nu']), dtype=tf.float64),
tf.constant(np.full(S.shape, self.param['gamma']), dtype=tf.float64)
])
return hazard_rates
if __name__=='__main__':
hh = draw_households(100, hh_size_distr)
K = (hh[:, None] == hh[None, :])
nsim = 120
model = HHstochastic(hh.astype(np.float64))
init_state = np.stack([np.broadcast_to(hh, [nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]]),
np.zeros([nsim, hh.shape[0]])]).astype(np.float64)
init_state[2, :, 0] = 1.
init_state[0, :, 0] = init_state[0, :, 0] - 1.
start = time.perf_counter()
t, sim = model.sample(init_state, [0., 365.], {'beta': 0.3, 'omega': 1.5, 'nu': 0.25, 'gamma': 0.25})
end = time.perf_counter()
print(f"Completed in {end-start} seconds")
r = tf.reduce_sum(sim[:, 2, :, :], axis=2)
plt.plot(t, r, 'r-', alpha=0.2)
plt.show()
# Covid_ODE model configuration
data:
age_mixing_matrix: data/polymod_normal_df.rds
age_mixing_matrix_term: data/polymod_normal_df.rds
age_mixing_matrix_hol: data/polymod_no_school_df.rds
mobility_matrix: data/movement.rds
population_size: data/pop.rds
parameter:
beta: 0.09
beta1: 0.04 # R0 2.4
beta2: 0.33 # Contact with commuters 1/3rd of the time
nu: 0.25
gamma: 0.65
gamma: 0.25
settings:
start: 0.
end: 365.
start: 2020-02-04
end: 2020-12-30
holiday:
- 2020-04-06
- 2020-05-01
time_step: 1.
output:
simulation: covid_ode.sim
simulation: covid_ode.hd5
import numpy as np
from scipy import optimize as opt
from covid.rdata import *
from covid.model import CovidUKODE
if __name__=='__main__':
K, _ = load_age_mixing('data/polymod_normal_df.rds')
T, _ = load_mobility_matrix('data/movement.rds')
np.fill_diagonal(T, 0.)
N, _ = load_population('data/pop.rds')
model = CovidUKODE(K, T, N)
print("R0 (beta=0.036) = ", model.eval_R0({'beta1': 0.036, 'beta2': 0.33, 'gamma': 0.25}))
R0 = 2.4
def opt_fn(beta, R0):
r0, i = model.eval_R0({'beta1': beta, 'beta2': 0.33, 'gamma': 0.25}, 1e-9)
return (r0 - R0)**2
res = opt.minimize_scalar(opt_fn, args=R0)
print(res['x'])
import tensorflow as tf
import numpy as np
from covid.model import Homogeneous
from covid.rdata import *
from covid.model import CovidUKODE
sw = tf.summary.create_file_writer('tf_log')
K = load_age_mixing('data/polymod_normal_df.rds').to_numpy(dtype=np.float32)
T = load_mobility_matrix('data/movement.rds').to_numpy(dtype=np.float32)
np.fill_diagonal(T, 0.)
N = load_population('data/pop.rds')['n'].to_numpy(dtype=np.float32)
model = Homogeneous()
popsize = 2500
state = np.array([np.full([popsize], 999.),
np.full([popsize], 0.),
np.full([popsize], 1.),
np.full([popsize], 0.)], dtype=np.float32)
model = CovidUKODE(K, T, N, T.shape[0])
print("Running...", flush=True, sep='')
tf.summary.trace_on(graph=True, profiler=True)
t, sim_state = model.sample(state,
[0., 10.],
{'beta': 0.2, 'nu': 0.14, 'gamma':0.14})
print("Done", flush=True)
print(sim_state)
with sw.as_default():
tf.summary.trace_export('profile', step=0, profiler_outdir='tf_log')
sw.flush()
sw.close()
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