Python3: SEIR モデルを表示する

こちらにあるプログラムを改造して見通しをよくしました。
csdegraaf/CoronaVirusModel
改造したのは、corona_spread.py だけです。

corona_spread.py

#! /usr/bin/python
#
#   corona_spread.py
#
#                       Aug/25/2020
# ------------------------------------------------------------------
import  sys

import numpy as np
from scipy import integrate
import matplotlib.pyplot as plt
#
import parameters as parameters
from calculations_module import seir_function

# ------------------------------------------------------------------
def plot_proc(tspan,y):
    total_cases = y[:, 1] + y[:, 2] + y[:, 3]
    total_cases_active = y[:, 1] + y[:, 2]

    fig, ax = plt.subplots()
    ax.plot(tspan, total_cases, color="b", label="E+I+R: Total cases")
    ax.plot(tspan, total_cases_active, color="r", label="E+I: Active cases")
    ax.set(xlabel="time (days)", ylabel="Patients", title='Cumulative and active cases')
    plt.legend()
#
    plt.show()
#
# ------------------------------------------------------------------
def out_proc(tspan,y):
    nsteps = np.size(tspan)
    S_end = y[nsteps - 1, 0]
    E_end = y[nsteps - 1, 1]
    I_end = y[nsteps - 1, 2]
    R_end = y[nsteps - 1, 3]

    total = S_end + E_end + I_end + R_end

    print('time (days): % 2d' %tspan[nsteps-1])

    print('total population: % 2d' %total)

    print('initial infected: % 2d' %I_0)

    print('total cases (E+I+R) at t= % 2d : % 2d' %(tspan[nsteps-1], E_end + I_end + R_end))

    print('Recovered at t=  % 2d : % 2d \n' %(tspan[nsteps-1], R_end))
    print('Infected (infectious) at t= % 2d : % 2d \n' %(tspan[nsteps-1],I_end))
    print('Exposed (non-infectious) at t= % 2d : % 2d \n ' %(tspan[nsteps-1], E_end))
    print('Susceptable at t= % 2d : % 2d \n ' %(tspan[nsteps-1], S_end))
#
# ------------------------------------------------------------------
def seir_with_params(t, y):
    return seir_function(t, y, params)
#
# ------------------------------------------------------------------
def calculation_proc(S_0,E_0,I_0,R_0):
    t_0 = 0
    tspan = np.linspace(t_0, 181, 180)

    y_init = np.zeros(4)
    y_init[0] = S_0
    y_init[1] = E_0
    y_init[2] = I_0
    y_init[3] = R_0


    r = integrate.ode(seir_with_params).set_integrator("dopri5")
    r.set_initial_value(y_init, t_0)
    y = np.zeros((len(tspan), len(y_init)))
    y[0, :] = y_init
    for i in range(1, 180):
        y[i, :] = r.integrate(tspan[i])
        if not r.successful():
            raise RuntimeError("Could not integrate")
#
    return tspan,y

# ------------------------------------------------------------------
sys.stderr.write("*** start ***\n")

S_0 = 11.0e+6
I_0 = 40.0
E_0 = 20. * I_0
R_0 = 0

c = 0.0
N = S_0 + I_0 + E_0 + R_0

sigma = 1. / 5.2
gamma = 1. / 18.


r_zero_array = np.zeros([6, 2])
r_zero_array[0, :] = [0.0,  3.0]
r_zero_array[1, :] = [20.0,  2.6]
r_zero_array[2, :] = [70.0,  1.9]
r_zero_array[3, :] = [84.0,  1.0]
r_zero_array[4, :] = [90.0,  .50]
r_zero_array[5, :] = [1000, .50]

params = parameters.Params(c, N, sigma, gamma, r_zero_array)

tspan,yy = calculation_proc(S_0,E_0,I_0,R_0)

plot_proc(tspan,yy)

out_proc(tspan,yy)

sys.stderr.write("*** end ***\n")
# ------------------------------------------------------------------

calculations_module.py

import numpy as np


def seir_function(t, y, params):
    """
    dS / dt = -beta * S * I / N
    dE / dt = +beta * S * I / N - sigma * E
    dI / dt = +sigma * E - gamma * I + c * R * I / N
    dR / dt = gamma * I - c * R * I / N

    yprime = [dS / dt  dE / dt dI / dt   dRdt]

    input:
      t current time
      y vector of current soln values
      y(1) = S, y(2) = E, y(3) = I, y(4) = R

    parameters in "params"
      beta, N, sigma, gamma, c, R_zero_array(table of values)

    output: (col vector)
      yprime(1) = dS / dt
      yprime(2) = dE / dt
      yprime(3) = dI / dt
      yprime(4) = dR / dt

    """
    R_zero_array = params.r_zero

    min_t = np.min(R_zero_array[:, 0])
    max_t = np.max(R_zero_array[:, 0])
    t_val = max(min_t, min(t, max_t))

    R_zero = np.interp(t_val, R_zero_array[:, 0], R_zero_array[:, 1])

    gamma = params.gamma

    beta = R_zero * gamma

    N = params.N
    sigma = params.sigma
    c = params.c

    S = y[0]
    E = y[1]
    I = y[2]
    R = y[3]

    yprime = np.zeros(4)

    yprime[0] = -beta * S * I / N
    yprime[1] = +beta * S * I / N - sigma * E
    yprime[2] = +sigma * E - gamma * I + c * R * I / N
    yprime[3] = gamma * I - c * R * I / N
    return yprime

parameters.py

class Params:
    def __init__(self, c, n, sigma, gamma, r_zero):
        self.c = c
        self.N = n
        self.sigma = sigma
        self.gamma = gamma
        self.r_zero = r_zero

実行方法

./corona_spread.py

実行結果