covid_ode.py 1015 Bytes
Newer Older
Chris Jewell's avatar
Chris Jewell committed
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)