8.3 Spreading of Disease: A System of First Order ODEs
233
For all this to work, the complete numerical solution must be represented by
a two-dimensional array, created by u = zeros((N_t+1, m+1)). The first index
counts the time points and the second the components of the solution vector at one
time point. That is, u[n,i] corresponds to the mathematical quantity u
n
i . When we
use only one index, as in u[n], this is the same as u[n,:] and picks out all the
components in the solution at the time point with index n. Then the assignment
u[n+1] = ... becomes correct because it is actually an in-place assignment
u[n+1, :] = .... The nice feature of these facts is that the same piece of Python
code works for both a scalar ODE and a system of ODEs!
The ode_FE function for the vector ODE is placed in the file ode_system_FE.py
and was written as follows:
import numpy as np
import matplotlib.pyplot as plt
def ode_FE(f, U_0, dt, T):
N_t = int(round(T/dt))
# Ensure that any list/tuple returned from f_ is wrapped as array
f_ = lambda u, t: np.asarray(f(u, t))
u = np.zeros((N_t+1, len(U_0)))
t = np.linspace(0, N_t*dt, len(u))
u[0] = U_0
for n in range(N_t):
u[n+1] = u[n] + dt*f_(u[n], t[n])
return u, t
The line f_ = lambda ... needs an explanation. For a user, who just needs to
define the f in the ODE system, it is convenient to insert the various mathematical
expressions on the right-hand sides in a list and return that list. Obviously, we
could demand the user to convert the list to a numpy array, but it is so easy to do
a general such conversion in the ode_FE function as well. To make sure that the
result from f is indeed an array that can be used for array computing in the formula
u[n] + dt*f(u[n], t[n]), we introduce a new function f_ that calls the user’s
f and sends the results through the numpy function asarray, which ensures that its
argument is converted to a numpy array (if it is not already an array).
Note also the extra parenthesis required when calling zeros with two indices.
Let us show how the previous SIR model can be solved using the new general
ode_FE that can solve any vector ODE. The user’s f(u, t) function takes a vector
u, with three components corresponding to S, I , and R as argument, along with the
current time point t[n], and must return the values of the formulas of the right-hand
sides in the vector ODE. An appropriate implementation is
def f(u, t):
S, I, R = u
return [-beta*S*I, beta*S*I - gamma*I, gamma*I]
Note that the S, I, and R values correspond to S n , I n , and R n . These values are then
just inserted in the various formulas in the vector ODE. Here we collect the values
in a list since the ode_FE function will anyway wrap this list in an array. We could,
of course, returned an array instead:
def f(u, t):
S, I, R = u
return array([-beta*S*I, beta*S*I - gamma*I, gamma*I])
233
For all this to work, the complete numerical solution must be represented by
a two-dimensional array, created by u = zeros((N_t+1, m+1)). The first index
counts the time points and the second the components of the solution vector at one
time point. That is, u[n,i] corresponds to the mathematical quantity u
n
i . When we
use only one index, as in u[n], this is the same as u[n,:] and picks out all the
components in the solution at the time point with index n. Then the assignment
u[n+1] = ... becomes correct because it is actually an in-place assignment
u[n+1, :] = .... The nice feature of these facts is that the same piece of Python
code works for both a scalar ODE and a system of ODEs!
The ode_FE function for the vector ODE is placed in the file ode_system_FE.py
and was written as follows:
import numpy as np
import matplotlib.pyplot as plt
def ode_FE(f, U_0, dt, T):
N_t = int(round(T/dt))
# Ensure that any list/tuple returned from f_ is wrapped as array
f_ = lambda u, t: np.asarray(f(u, t))
u = np.zeros((N_t+1, len(U_0)))
t = np.linspace(0, N_t*dt, len(u))
u[0] = U_0
for n in range(N_t):
u[n+1] = u[n] + dt*f_(u[n], t[n])
return u, t
The line f_ = lambda ... needs an explanation. For a user, who just needs to
define the f in the ODE system, it is convenient to insert the various mathematical
expressions on the right-hand sides in a list and return that list. Obviously, we
could demand the user to convert the list to a numpy array, but it is so easy to do
a general such conversion in the ode_FE function as well. To make sure that the
result from f is indeed an array that can be used for array computing in the formula
u[n] + dt*f(u[n], t[n]), we introduce a new function f_ that calls the user’s
f and sends the results through the numpy function asarray, which ensures that its
argument is converted to a numpy array (if it is not already an array).
Note also the extra parenthesis required when calling zeros with two indices.
Let us show how the previous SIR model can be solved using the new general
ode_FE that can solve any vector ODE. The user’s f(u, t) function takes a vector
u, with three components corresponding to S, I , and R as argument, along with the
current time point t[n], and must return the values of the formulas of the right-hand
sides in the vector ODE. An appropriate implementation is
def f(u, t):
S, I, R = u
return [-beta*S*I, beta*S*I - gamma*I, gamma*I]
Note that the S, I, and R values correspond to S n , I n , and R n . These values are then
just inserted in the various formulas in the vector ODE. Here we collect the values
in a list since the ode_FE function will anyway wrap this list in an array. We could,
of course, returned an array instead:
def f(u, t):
S, I, R = u
return array([-beta*S*I, beta*S*I - gamma*I, gamma*I])