covid_ode.py 1.08 KB
 Chris Jewell committed Mar 01, 2020 1 2 3 4 5 6 7 8 9 10 11 12 13 14 ``````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]) `````` Chris Jewell committed Mar 02, 2020 15 16 17 ``````param = {'beta': 0.0002, 'nu': 0.14, 'gamma': 0.14} `````` Chris Jewell committed Mar 01, 2020 18 19 `````` def h(t, state): `````` Chris Jewell committed Mar 02, 2020 20 `````` print(state) `````` Chris Jewell committed Mar 01, 2020 21 22 23 24 25 26 27 28 29 30 31 32 33 34 `````` 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): `````` Chris Jewell committed Mar 02, 2020 35 `````` return tode.DormandPrince(first_step_size=1., max_num_steps=5000).solve(rates, t_init, state_init, solution_times=t) `````` Chris Jewell committed Mar 01, 2020 36 37 38 `````` solution_times = np.arange(0., 365., 1.) `````` Chris Jewell committed Mar 02, 2020 39 ``````print("Running...", flush=True) `````` Chris Jewell committed Mar 01, 2020 40 ``````result = solve_ode(rates=h, t_init=0., state_init=state, t=solution_times) `````` Chris Jewell committed Mar 02, 2020 41 ``````print("Done") `````` Chris Jewell committed Mar 01, 2020 42 ``print(result)``