SIRモデル方程式を数値的に解いて、gifで表示させてみる

伝染病の拡大過程を記述するSIRモデル

の初期値問題を Python で解き, 得られた解をgifにします.

PCとインターネットがあればどなたでも, 以下のノートブックを逐次実行することでgifファイルを生成することができます.
colab.research.google.com

Google Colaboratory を使ったノートブック形式での数値計算

qiita.com

に則り, ノートブックで数値計算を行っていきましょう.
以下のリンク先で, [ファイル] -> [ノートブックを新規作成] により環境の設定が完了します
colab.research.google.com

数値計算

ライブラリのインポート

まずは数値計算や描画に必要なライブラリをインポートします:

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML
from scipy.integrate import solve_ivp

今回は微分方程式の数値計算を実装することはせず, scipy.integrate.solve_ivp を用いて
常微分方程式の初期値問題を解きます.

docs.scipy.org

パラメータの設定

感染症の感染率, 隔離率を表すパラメータはそれぞれ とし,
総人口は 人に対して感染者は最初に 人いるものとします.

# parameter
beta = 2.0e-4
gamma = 1.0e-1
T = 100 # last time
dt = 1.0e-2 # step width for time variable

# initial condition
M = 1000
epsilon = 10

モデルの設定

scipy.integrate.solve_ivp に渡す, 常微分方程式系を標準化した際の右辺を関数を設定します.
3成分の連立方程式系のため, 引数は時間変数 と長さ3のリストである で,
返り値は長さ3のリストである必要があります.

def sir(t, var):
    s, i, r = var
    dsdt = - beta * s * i
    didt =   beta * s * i - gamma * i
    drdt =                  gamma * i
    return [dsdt, didt, drdt]

解く

常微分方程式を解くタイムスパンとして区間 , 刻み幅 となるように np.array型の配列を与えます.

t = np.arange(0, T, dt)
S, I, R = solve_ivp(sir, [0, T], [M-epsilon, epsilon, 0], t_eval=t).y

描画する

animationを作成します.
IPython.display.HTML により, 再生ボタンを押してgifを再生することができます.

fig = plt.figure()
fig.set_dpi(100)
ax = fig.add_subplot(1,1,1)
def animate(k):
    ax.clear()
    plt.xlim([-.1, 100.1])
    plt.ylim([-.0, 1000.])
    u = I + R + S
    ax.fill_between( t[:100*k], u[:100*k], color="lightblue", alpha=0.9)
    u = I + R
    ax.fill_between( t[:100*k], u[:100*k], color="black", alpha=0.9)
    u = I
    ax.fill_between( t[:100*k], u[:100*k], color="yellow", alpha=0.9)
    time_text = ax.text(0.05, 0.9, f"{k}/100", transform=ax.transAxes)
anim = animation.FuncAnimation(fig, animate, frames=len(t)//100, interval=200, repeat=True)
anim.save("sir.gif", writer="pillow")
HTML(anim.to_jshtml())