covid_ode.py 1.08 KB
Newer Older
Chris Jewell's avatar
Chris Jewell committed
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's avatar
Chris Jewell committed
15
16
17
param = {'beta': 0.0002, 'nu': 0.14, 'gamma': 0.14}


Chris Jewell's avatar
Chris Jewell committed
18
19

def h(t, state):
Chris Jewell's avatar
Chris Jewell committed
20
    print(state)
Chris Jewell's avatar
Chris Jewell committed
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's avatar
Chris Jewell committed
35
    return tode.DormandPrince(first_step_size=1., max_num_steps=5000).solve(rates, t_init, state_init, solution_times=t)
Chris Jewell's avatar
Chris Jewell committed
36
37
38

solution_times = np.arange(0., 365., 1.)

Chris Jewell's avatar
Chris Jewell committed
39
print("Running...", flush=True)
Chris Jewell's avatar
Chris Jewell committed
40
result = solve_ode(rates=h, t_init=0., state_init=state, t=solution_times)
Chris Jewell's avatar
Chris Jewell committed
41
print("Done")
Chris Jewell's avatar
Chris Jewell committed
42
print(result)