Commit 4c0c31d0 authored by Chris Jewell's avatar Chris Jewell
Browse files

Initial commit

parents
"""Functions for chain binomial simulation."""
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
def update_state(update, state, stoichiometry):
update = tf.expand_dims(update, 1) # Rx1xN
update *= tf.expand_dims(stoichiometry, -1) # RxSx1
update = tf.reduce_sum(update, axis=0) # SxN
return state + update
def chain_binomial_propagate(h, time_step, stoichiometry):
def propagate_fn(state):
rates = h(state)
probs = 1 - tf.exp(-rates*time_step) # RxN
update = tfd.Binomial(state[:-1], probs=probs).sample() # RxN
update = tf.expand_dims(update, 1) # Rx1xN
upd_shape = tf.concat([stoichiometry.shape, tf.fill([tf.rank(state)-1], 1)], axis=0)
update *= tf.reshape(stoichiometry, upd_shape) # RxSx1
update = tf.reduce_sum(update, axis=0)
state = state + update
return state
return propagate_fn
#@tf.function
def chain_binomial_simulate(hazard_fn, state, start, end, time_step, stoichiometry):
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 = output.write(0, state)
for i in tf.range(1, times.shape[0]):
state = propagate(state)
output = output.write(i, state)
sim = output.gather(tf.range(times.shape[0]))
return times, sim
"""Functions for infection rates"""
import tensorflow as tf
from covid.impl.chainbinom_simulate import chain_binomial_simulate
class CovidUK:
def __init__(self, K, T, W):
self.K = K
self.T = T
self.W = W
self.stoichiometry = [[-1, 1, 0, 0],
[0, -1, 1, 0],
[0, 0, -1, 1]]
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['nu'],
self.param['gamma']
])
return hazard_rates
def sample(self, initial_state, time_lims, param):
self.param = param
return chain_binomial_simulate(self.h, initial_state, time_lims[0],
time_lims[1], 1., self.stoichiometry)
class Homogeneous:
def __init__(self):
self.stoichiometry = tf.constant([[-1, 1, 0, 0],
[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['nu'] * tf.ones_like(I),
self.param['gamma'] * tf.ones_like(I)
])
return hazard_rates
def sample(self, initial_state, time_lims, param):
self.param = param
return chain_binomial_simulate(self.h, initial_state, time_lims[0],
time_lims[1], 1., self.stoichiometry)
import unittest
import tensorflow as tf
from covid.impl.chainbinom_simulate import chain_binomial_propagate
class TestPropagate(unittest.TestCase):
def test_1d_shape(self):
state = tf.constant([[18], [19], [20]], dtype=tf.float32)
stoichiometry = tf.constant([[-1, 1, 0], [0, -1, 1]], dtype=tf.float32)
def h(state):
return tf.stack([[0.1], [0.1]])
propagate_fn = chain_binomial_propagate(h, 1, stoichiometry)
new_state = propagate_fn(state)
self.assertListEqual(list(new_state.shape), [3, 1])
def test_2d_shape(self):
state = tf.constant([[20, 10],
[15, 5],
[0, 0]], dtype=tf.float32)
stoichiometry = tf.constant([[-1, 1, 0], [0, -1, 1]], dtype=tf.float32)
def h(state):
return tf.stack([[0.1, 0.1],
[0.1, 0.1]])
propagate_fn = chain_binomial_propagate(h, 1, stoichiometry)
new_state = propagate_fn(state)
self.assertListEqual(list(new_state.shape), [3, 2])
def test_3d_shape(self):
state = tf.constant([[[20, 19], [18, 17]],
[[16, 15], [14, 13]],
[[12, 11], [10, 9]]], dtype=tf.float32)
stoichiometry = tf.constant([[-1, 1, 0], [0, -1, 1]], dtype=tf.float32)
def h(state):
return tf.stack([[[0.1, 0.1], [0.1, 0.1]],
[[0.1, 0.1], [0.1, 0.1]]])
propagate_fn = chain_binomial_propagate(h, 1, stoichiometry)
new_state = propagate_fn(state)
self.assertListEqual(list(new_state.shape), [3, 2, 2])
if __name__ == '__main__':
unittest.main()
import unittest
import tensorflow as tf
from covid.impl.chainbinom_simulate import chain_binomial_simulate
class TestSimulator(unittest.TestCase):
def test_homogenousSIR(self):
state = tf.constant([[999], [1], [0]], dtype=tf.float32)
stoichiometry = tf.constant([[-1, 1, 0],
[0, -1, 1]], dtype=tf.float32)
def h(state):
return tf.stack([[0.4 * state[1, 0]/tf.reduce_sum(state)],
[0.14]])
sim = chain_binomial_simulate(h, state, 0., 365., 1, stoichiometry)
print(sim)
if __name__ == '__main__':
unittest.main()
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