Skip to content
GitLab
Commit 043fa4f5 authored by Chris Jewell's avatar Chris Jewell
Browse files

Working MCMC for known event times.

parent 2d423d9e
Hide whitespace changes
Inline Side-by-side
covid/impl/discrete_markov.py
View file @ 043fa4f5
......@@ -77,11 +77,13 @@ def discrete_markov_log_prob(events, init_state, hazard_fn, time_step):
states = tf.concat([[init_state], tf.reduce_sum(events, axis=-2)], axis=-3)[:-1]
t = tf.range(states.shape[-3])
rate_matrix = hazard_fn(t, states)
rate_matrix = tf.linalg.set_diag(rate_matrix,
-tf.reduce_sum(rate_matrix, axis=-1))
markov_transition = tf.linalg.expm(rate_matrix*time_step)
log_mt = tf.math.log(markov_transition)
idx = events > 0. # Todo: probably should check for x>0 when p==0.
logp = tf.reduce_sum(events[idx] * log_mt[idx])
return logp
def log_prob_t(a, elems):
t, event, state = elems
rate_matrix = hazard_fn(t, state)
rate_matrix = tf.linalg.set_diag(rate_matrix,
-tf.reduce_sum(rate_matrix, axis=-1))
markov_transition = tf.linalg.expm(rate_matrix*time_step)
logp = tfd.Multinomial(state, probs=markov_transition).log_prob(event)
return a + tf.reduce_sum(logp)
return tf.foldl(log_prob_t, (t, events, states), initializer=tf.constant(0., events.dtype))
covid/model.py
View file @ 043fa4f5
......@@ -4,9 +4,10 @@ import tensorflow_probability as tfp
from tensorflow_probability.python.internal import dtype_util
import numpy as np
from covid.impl.util import make_transition_rate_matrix
from covid.rdata import load_mobility_matrix, load_population, load_age_mixing
from covid.pydata import load_commute_volume
from covid.impl.discrete_markov import discrete_markov_simulation
from covid.impl.discrete_markov import discrete_markov_simulation, discrete_markov_log_prob
tode = tfp.math.ode
tla = tf.linalg
......@@ -236,27 +237,18 @@ class CovidUKStochastic(CovidUK):
t_idx = tf.clip_by_value(tf.cast(t, tf.int64), 0, self.max_t)
m_switch = tf.gather(self.m_select, t_idx)
commute_volume = tf.pow(tf.gather(self.W, t_idx), param['omega'])
lockdown = tf.gather(self.lockdown_select, t_idx)
beta = tf.where(lockdown == 0, param['beta1'], param['beta1'] * param['beta3'])
infec_rate = param['beta1'] * (
tf.gather(self.M.matvec(state[:, 2]), m_switch) +
param['beta2'] * self.Kbar * commute_volume * self.C.matvec(state[:, 2] / self.N_sum))
infec_rate = beta * (
tf.gather(self.M.matvec(state[..., 2]), m_switch) +
param['beta2'] * self.Kbar * commute_volume * self.C.matvec(state[..., 2] / self.N_sum))
infec_rate = infec_rate / self.N # Vector of length nc
ei = tf.broadcast_to([param['nu']], shape=[state.shape[0]]) # Vector of length nc
ir = tf.broadcast_to([param['gamma']], shape=[state.shape[0]]) # Vector of length nc
# Scatter rates into a [nc, ns, ns] tensor
n = state.shape[0]
b = tf.stack([tf.range(n),
tf.zeros(n, dtype=tf.int32),
tf.ones(n, dtype=tf.int32)], axis=-1)
indices = tf.stack([b, b + [0, 1, 1], b + [0, 2, 2]], axis=-2)
# Un-normalised rate matrix (diag is 0 here)
rate_matrix = tf.scatter_nd(indices=indices,
updates=tf.stack([infec_rate, ei, ir], axis=-1),
shape=[state.shape[0],
state.shape[1],
state.shape[1]]) # Tensor of dim [nc, ns, ns]
ei = tf.broadcast_to([tf.convert_to_tensor(param['nu'])], shape=[state.shape[0]]) # Vector of length nc
ir = tf.broadcast_to([tf.convert_to_tensor(param['gamma'])], shape=[state.shape[0]]) # Vector of length nc
rate_matrix = make_transition_rate_matrix([infec_rate, ei, ir], [[0, 1], [1, 2], [2, 3]], state)
return rate_matrix
return h
......@@ -273,3 +265,13 @@ class CovidUKStochastic(CovidUK):
t, sim = discrete_markov_simulation(hazard, state_init, np.float64(0.),
np.float64(self.times.shape[0]), self.time_step)
return t, sim
def log_prob(self, y, param, state_init):
"""Calculates the log probability of observing epidemic events y
:param y: a list of tensors. The first is of shape [n_times] containing times,
the second is of shape [n_times, n_states, n_states] containing event matrices.
:param param: a list of parameters
:returns: a scalar giving the log probability of the epidemic
"""
hazard = self.make_h(param)
return discrete_markov_log_prob(y, state_init, hazard, self.time_step)
covid_stochastic.py
View file @ 043fa4f5
import optparse
import time
import pickle as pkl
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
import numpy as np
import matplotlib.pyplot as plt
import yaml
......@@ -11,6 +16,21 @@ from covid.util import sanitise_parameter, sanitise_settings, seed_areas
DTYPE = np.float64
def random_walk_mvnorm_fn(covariance, name=None):
"""Returns callable that adds Multivariate Normal noise to the input"""
covariance = covariance + tf.eye(covariance.shape[0], dtype=tf.float64) * 1.e-9
scale_tril = tf.linalg.cholesky(covariance)
rv = tfp.distributions.MultivariateNormalTriL(loc=tf.zeros(covariance.shape[0], dtype=tf.float64),
scale_tril=scale_tril)
def _fn(state_parts, seed):
with tf.name_scope(name or 'random_walk_mvnorm_fn'):
new_state_parts = [rv.sample() + state_part for state_part in state_parts]
return new_state_parts
return _fn
def sum_age_groups(sim):
infec = sim[:, 2, :]
infec = infec.reshape([infec.shape[0], 152, 17])
......@@ -199,3 +219,5 @@ if __name__ == '__main__':
fig_uk.gca().grid(True)
plt.show()
with open('stochastic_sim.pkl', 'wb') as f:
pkl.dump({'events': upd, 'state_init': state_init}, f)
mcmc.py
View file @ 043fa4f5
"""Inference on stochastic models"""
import optparse
import pickle as pkl
import time
import pickle as pkl
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
import numpy as np
import matplotlib.pyplot as plt
import yaml
from tensorflow_probability import bijectors as tfb
from tensorflow_probability import distributions as tfd
from covid.model import CovidUKODE, covid19uk_logp, load_data
from covid.util import *
from covid.model import CovidUKStochastic, load_data
from covid.util import sanitise_parameter, sanitise_settings, seed_areas
DTYPE = np.float64
def random_walk_mvnorm_fn(covariance, name=None):
"""Returns callable that adds Multivariate Normal noise to the input"""
covariance = covariance + tf.eye(covariance.shape[0], dtype=tf.float64) * 1.e-9
......@@ -28,72 +33,67 @@ def random_walk_mvnorm_fn(covariance, name=None):
return _fn
if __name__ == '__main__':
parser = optparse.OptionParser()
parser.add_option("--config", "-c", dest="config", default="ode_config.yaml",
help="configuration file")
options, args = parser.parse_args()
with open(options.config, 'r') as ymlfile:
config = yaml.load(ymlfile)
param = sanitise_parameter(config['parameter'])
settings = sanitise_settings(config['settings'])
data = load_data(config['data'], settings)
parser = optparse.OptionParser()
parser.add_option("--config", "-c", dest="config", default="ode_config.yaml",
help="configuration file")
options, args = parser.parse_args()
with open(options.config, 'r') as ymlfile:
config = yaml.load(ymlfile)
case_reports = pd.read_csv(config['data']['reported_cases'])
case_reports['DateVal'] = pd.to_datetime(case_reports['DateVal'])
case_reports = case_reports[case_reports['DateVal'] >= '2020-02-19']
date_range = [case_reports['DateVal'].min(), case_reports['DateVal'].max()]
y = case_reports['CumCases'].to_numpy()
y_incr = np.round((y[1:] - y[:-1]) * 0.8)
param = sanitise_parameter(config['parameter'])
settings = sanitise_settings(config['settings'])
simulator = CovidUKODE(M_tt=data['M_tt'],
M_hh=data['M_hh'],
C=data['C'],
N=data['pop']['n'].to_numpy(),
W=data['W'].to_numpy(),
date_range=[date_range[0] - np.timedelta64(1, 'D'), date_range[1]],
holidays=settings['holiday'],
lockdown=settings['lockdown'],
time_step=int(settings['time_step']))
data = load_data(config['data'], settings, DTYPE)
seeding = seed_areas(data['pop']['n'].to_numpy(), data['pop']['Area.name.2']) # Seed 40-44 age group, 30 seeds by popn size
state_init = simulator.create_initial_state(init_matrix=seeding)
model = CovidUKStochastic(M_tt=data['M_tt'],
M_hh=data['M_hh'],
C=data['C'],
N=data['pop']['n'].to_numpy(),
W=data['W'],
date_range=settings['inference_period'],
holidays=settings['holiday'],
lockdown=settings['lockdown'],
time_step=1.)
with open('stochastic_sim.pkl', 'rb') as f:
sim = pkl.load(f)
events = sim['events']
state_init = sim['state_init']
def logp(par):
p = param
p['beta1'] = par[0]
p['beta3'] = par[1]
p['gamma'] = par[2]
p['I0'] = par[3]
p['r'] = par[4]
beta_logp = tfd.Gamma(concentration=tf.constant(1., dtype=DTYPE), rate=tf.constant(1., dtype=DTYPE)).log_prob(p['beta1'])
beta_logp = tfd.Gamma(concentration=tf.constant(1., dtype=DTYPE),
rate=tf.constant(1., dtype=DTYPE)).log_prob(p['beta1'])
beta3_logp = tfd.Gamma(concentration=tf.constant(200., dtype=DTYPE),
rate=tf.constant(200., dtype=DTYPE)).log_prob(p['beta3'])
gamma_logp = tfd.Gamma(concentration=tf.constant(100., dtype=DTYPE), rate=tf.constant(400., dtype=DTYPE)).log_prob(p['gamma'])
I0_logp = tfd.Gamma(concentration=tf.constant(1.5, dtype=DTYPE), rate=tf.constant(0.05, dtype=DTYPE)).log_prob(p['I0'])
r_logp = tfd.Gamma(concentration=tf.constant(0.1, dtype=DTYPE), rate=tf.constant(0.1, dtype=DTYPE)).log_prob(p['gamma'])
state_init = simulator.create_initial_state(init_matrix=seeding * p['I0'])
t, sim, solve = simulator.simulate(p, state_init)
y_logp = covid19uk_logp(y_incr, sim, 0.1, p['r'])
logp = beta_logp + beta3_logp + gamma_logp + I0_logp + r_logp + tf.reduce_sum(y_logp)
gamma_logp = tfd.Gamma(concentration=tf.constant(100., dtype=DTYPE),
rate=tf.constant(400., dtype=DTYPE)).log_prob(p['gamma'])
y_logp = model.log_prob(events, p, state_init)
logp = beta_logp + beta3_logp + gamma_logp + y_logp
return logp
def trace_fn(_, pkr):
return (
pkr.inner_results.log_accept_ratio,
pkr.inner_results.accepted_results.target_log_prob,
pkr.inner_results.accepted_results.step_size)
unconstraining_bijector = [tfb.Exp()]
initial_mcmc_state = np.array([0.05, 1.0, 0.25, 1.0, 50.], dtype=np.float64) # beta1, gamma, I0
initial_mcmc_state = np.array([0.05, 0.5, 0.25], dtype=np.float64) # beta1, gamma, I0
print("Initial log likelihood:", logp(initial_mcmc_state))
@tf.function(autograph=False, experimental_compile=True)
@tf.function #(experimental_compile=True)
def sample(n_samples, init_state, scale, num_burnin_steps=0):
return tfp.mcmc.sample_chain(
num_results=n_samples,
......@@ -109,7 +109,7 @@ if __name__ == '__main__':
joint_posterior = tf.zeros([0] + list(initial_mcmc_state.shape), dtype=DTYPE)
scale = np.diag([0.1, 0.1, 0.1, 0.1, 1.])
scale = np.diag([0.1, 0.1, 0.1])
overall_start = time.perf_counter()
num_covariance_estimation_iterations = 50
......@@ -148,5 +148,5 @@ if __name__ == '__main__':
plt.show()
print(f"Posterior mean: {np.mean(joint_posterior, axis=0)}")
with open('pi_beta_2020-03-29.pkl', 'wb') as f:
pkl.dump(joint_posterior, f)
with open('stochastic_posterior.pkl', 'wb') as f:
pkl.dump(joint_posterior, f)
\ No newline at end of file
Supports Markdown
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