covid_ode.py 1015 Bytes
 Chris Jewell committed Mar 01, 2020 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 ``````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)``````