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

実行結果