Commit 7c8712ef authored by Chris Jewell's avatar Chris Jewell
Browse files

Added ODE example

parent b31f2d11
......@@ -26,7 +26,6 @@ def chain_binomial_propagate(h, time_step, stoichiometry):
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)
......
......@@ -14,6 +14,7 @@ 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
......@@ -25,6 +26,7 @@ class CovidUK:
])
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],
......@@ -49,6 +51,7 @@ class Homogeneous:
])
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],
......
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
tode = tfp.math.ode
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])
param = {'beta': 0.2, 'nu': 0.14, 'gamma': 0.14}
@tf.function
def h(t, 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
dR = param['gamma'] * I
df = tf.stack([dS, dE, dI, dR])
return df
@tf.function
def solve_ode(rates, t_init, state_init, t):
return tode.DormandPrince().solve(rates, t_init, state_init, solution_times=t)
solution_times = np.arange(0., 365., 1.)
result = solve_ode(rates=h, t_init=0., state_init=state, t=solution_times)
print(result)
\ No newline at end of file
import tensorboard as tb
import tensorflow as tf
import numpy as np
from covid.model import Homogeneous
sw = tf.summary.create_file_writer('tf_log')
model = Homogeneous()
popsize = 10
popsize = 2500
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., 365.],
[0., 5.],
{'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