Python Gekko ODE flat results

I’ve been playing with Gekko(https://gekko.readthedocs.io/en/latest/). To test it I implemented the Covid model from here https://julia.quantecon.org/continuous_time/seir_model.html. But the results are just flat graphs. It looks nothing like the result from the link. Can anyone see what I’m doing wrong?

import numpy as np
from gekko import GEKKO
import matplotlib.pyplot as plt

m = GEKKO(remote=False)    # create GEKKO model

# constants
y = 1/18
o = 1/5.2
R = 3

# create variables
i_0 = 1E-7
e_0 = 4.0 * i_0
s_0 = 1.0 - i_0 - e_0
r_0 = 0.0

s = m.Var(s_0)
e = m.Var(e_0)
i = m.Var(i_0)
r = m.Var(r_0)

# create equations
m.Equation(s.dt()==-y*R*s*i)
m.Equation(e.dt()==-y*R*s*i - o*e)
m.Equation(i.dt()==o*e - y*i)
m.Equation(r.dt()==y*i)

m.time = np.linspace(0,350)

# solve ODE
m.options.IMODE = 4
m.solve()

# plot results
plt.plot(m.time, s, label="s")
plt.plot(m.time, e, label="e")
plt.plot(m.time, i, label="i")
plt.plot(m.time, r, label="r")
plt.legend()
plt.show()

Solution

It is a mistake with this equation:

m.Equation(e.dt()==-y*R*s*i - o*e)

It should be:

m.Equation(e.dt()==y*R*s*i - o*e)

You also need these options to increase the accuracy because the problem has small values and Gekko uses 1e-6 tolerance by default:

# solve ODE
m.options.IMODE = 7
m.options.OTOL  = 1e-8
m.options.RTOL  = 1e-8
m.options.NODES = 3

If you don't want to refine the solver tolerance then you can also scale the equations.

sc = 1000
m.Equation(sc*s.dt()==-y*R*s*i*sc)
m.Equation(sc*e.dt()==(y*R*s*i - o*e)*sc)
m.Equation(sc*i.dt()==(o*e - y*i)*sc)
m.Equation(sc*r.dt()==y*i*sc)

Here is a working version of the script with a couple other modifications to improve the accuracy and solution speed:

import numpy as np
from gekko import GEKKO
import matplotlib.pyplot as plt

m = GEKKO(remote=False)    # create GEKKO model

# constants
y = 1/18
o = 1/5.2
R = 3

# create variables
i_0 = 1E-7
e_0 = 4.0 * i_0
s_0 = 1.0 - i_0 - e_0
r_0 = 0.0

s = m.Var(s_0)
e = m.Var(e_0)
i = m.Var(i_0)
r = m.Var(r_0)

# create equations
m.Equation(s.dt()==-y*R*s*i)
m.Equation(e.dt()==y*R*s*i - o*e)
m.Equation(i.dt()==o*e - y*i)
m.Equation(r.dt()==y*i)

m.time = np.linspace(0,350)
m.time = np.insert(m.time,1,np.logspace(-8,-1,10))
    
# solve ODE
m.options.IMODE = 7
m.options.OTOL  = 1e-8
m.options.RTOL  = 1e-8
m.options.NODES = 3
m.solve()

# plot results
plt.plot(m.time, s, label="s")
plt.plot(m.time, e, label="e")
plt.plot(m.time, i, label="i")
plt.plot(m.time, r, label="r")
plt.legend()
plt.show()

Here is another SEIR model for COVID-19 in Gekko along with a simple case study on optimization of healthcare resources.

import numpy as np
from gekko import GEKKO
import matplotlib.pyplot as plt

t_incubation = 5.1 
t_infective = 3.3 
R0 = 2.4 
N = 100000 

# initial number of infected and recovered individuals
e_initial = 1/N
i_initial = 0.00
r_initial = 0.00
s_initial = 1 - e_initial - i_initial - r_initial

alpha = 1/t_incubation
gamma = 1/t_infective
beta = R0*gamma

m = GEKKO()
u = m.FV(0)
s,e,i,r = m.Array(m.Var,4)
s.value = s_initial
e.value = e_initial
i.value = i_initial
s.value = s_initial
m.Equations([s.dt()==-(1-u)*beta * s * i,\
             e.dt()== (1-u)*beta * s * i - alpha * e,\
             i.dt()==alpha * e - gamma * i,\
             r.dt()==gamma*i])

t = np.linspace(0, 200, 101)
t = np.insert(t,1,[0.001,0.002,0.004,0.008,0.02,0.04,0.08,\
                   0.2,0.4,0.8])
m.time = t
m.options.IMODE=7
m.solve(disp=False)

# plot the data
plt.figure(figsize=(8,5))
plt.subplot(2,1,1)
plt.plot(m.time, s.value, color='blue', lw=3, label='Susceptible')
plt.plot(m.time, r.value, color='red',  lw=3, label='Recovered')
plt.ylabel('Fraction')
plt.legend()

plt.subplot(2,1,2)
plt.plot(m.time, i.value, color='orange', lw=3, label='Infective')
plt.plot(m.time, e.value, color='purple', lw=3, label='Exposed')
plt.ylim(0, 0.2)
plt.xlabel('Time (days)')
plt.ylabel('Fraction')
plt.legend()

plt.show()

The case study shows how to calculate optimal interaction parameters to meet healthcare constraints.