Python教材用のSIRモデル

友だちから「じぶんでSIRモデルが書けるようになりたい，でもそれだけだと面白くないから，できれば何かプラスアルファの機能がほしい．」というような依頼があったので，教材を作っている．

それで，先にゴールとして「このくらい書けるようになる」プログラムの例をかんたんな表現で書いてみた．

S．I，Rがベクトルになっているが，これは年齢層や地域などで別々に計算したいときに使う．いまは，3つとも同条件なので，このままだと意味がないが……．βを変更するか，グループ同士の交流を隣接行列のようなものを使って制御する必要がある．まあ，それはプラスアルファということで．現状のColabへのリンクも最後に張っておきます．

import numpy as np
import matplotlib.pyplot as plt

#基本設定
GROUPS=["A","B","C"]
ONES=np.ones_like(GROUPS, dtype=np.float64)
ZEROS=np.zeros_like(GROUPS, dtype=np.float64)

#SIRモデル定数
R0=1.0
γ=ONES*1/15
β=ONES*R0*γ
print("γ",γ)
print("β",β)

#SIRモデル変数初期値
S=[ONES]
I=[ONES*0.001]
R=[ZEROS]

def calc(s, i ,r):
    #差分の計算
    ds=-β*s*sum(i) #グループに関係なく接触
    #ds=-β*s*i #グループ同士が接触しない
    dr=γ*i
    di=-ds-dr

    #差分を適用した値を返す
    return s+ds, i+di, r+dr

def run(S, I, R, TMAX=100):
    s=S[-1]
    i=I[-1]
    r=R[-1]
    for t in range(1,TMAX):
        s,i,r=calc(s, i, r)
        S.append(s)
        I.append(i)
        R.append(r)
    return S, I, R

S, I, R=run(S, I, R)

plt.plot(S)
plt.plot(I)
plt.plot(R)
plt.legend()
plt.show()

colab.research.google.com