SIRベイズモデルで新型コロナの収束を予測

from numpy import inf
import numpy as np
from scipy.optimize import minimize
from scipy.integrate import odeint

def sir(y, t, beta, gamma):
    dydt = np.zeros(3)
    dydt[0] = - beta * y[0] * y[1] / (y[0] + y[1] + y[2])   # dS/dt
    dydt[1] = beta * y[0] * y[1] / (y[0] + y[1] + y[2]) - gamma * y[1]  # dI/dt
    dydt[2] = gamma * y[1]  # dR/dt
    return dydt

def estimate(target_pred_var, init_state, beta, gamma):
    sol = odeint(sir, init_state, np.arange(0, len(target_pred_var)), args=(beta, gamma))
    return sol

def loss(x, init_state, target_pred_var):
    beta, gamma = x[0], x[1]
    sol = estimate(target_pred_var, init_state, beta, gamma)
    return np.sum((sol[:, 1] - target_pred_var) ** 2)

def optimize(init_state, target_pred_var):
    result = minimize(loss, x0=[0.001, 0.001], args=(init_state, target_pred_var),
                      method="Nelder-Mead")
    return result

def prediction(result, n_pred, init_state):
    beta = result.x[0]
    gamma = result.x[1]
    sol = estimate(np.zeros(n_pred), init_state, beta, gamma)
    return sol

df_tmp = df.reset_index()
df_tmp['dt'] = df_tmp['dt'].astype('str')
df_tmp['dt2'] = df_tmp['dt'].apply(lambda x: dt.strptime(x, '%Y%m%d'))
target_pred_var = df_tmp['infected'].values

target_pred_var_r = df_tmp['recovered'].values
I = target_pred_var
R = target_pred_var_r
N = 14000000
S = N - I - R
init_state = [S[0], I[0], R[0]]

result = optimize(init_state, target_pred_var)
print(result)

 final_simplex: (array([[1.98561299e-01, 1.39422106e-01],
       [1.98561055e-01, 1.39421706e-01],
       [1.98563439e-01, 1.39424393e-01]]),
       array([106399.61884819, 106399.61886498, 106399.61935616]))
           fun: 106399.61884819488
       message: 'Optimization terminated successfully.'
           nit: 89
          nfev: 178
        status: 0
       success: True
             x: array([0.19856129, 0.13942211])

n_pred = 50
t_max = len(target_pred_var) + n_pred
dt = np.arange(0, t_max)
beta = result.x[0]
gamma = result.x[1]

dates = data.index.tolist()
for i in range(n_pred):
    d = d + datetime.timedelta(days=1)
    dates.append(d)

sol = odeint(sir, init_state, dt,
             args=(beta, gamma))

plt.figure(figsize=(10,6))
plt.plot(dates, sol[:,0] / N, 'b', label='Susceptible (estimation)')
plt.plot(dates, sol[:,1] / N, 'r', label='Infected (estimation)')
plt.plot(dates, sol[:,2] / N, 'g', label='Recovered (estimation)')

ax1.plot_date(dates[:len(I)], I / N, 'or', label='Infected (observation)')
ax1.plot_date(dates[:len(R)], R / N, 'og', label='Recovered (observation)')

plt.legend()
plt.xlabel('Date')
plt.ylabel('Population')
plt.title('Prediction of Covid-19 epidemic in ' + target_pref_str)
plt.show()

model_code = """
functions {
    real[] sir(real t,
               real[] y,
               real[] theta,
               real[] x_r,
               int[] x_i) {
        real dydt[3];
        dydt[1] = - theta[1] * y[1] * y[2] / (y[1] + y[2] + y[3]);
        dydt[2] = theta[1] * y[1] * y[2] / (y[1] + y[2] + y[3]) - theta[2] * y[2];
        dydt[3] = theta[2] * y[2];
        return dydt;
    }
}
data {
    int T;
    real y0[3];
    real t0;
    real ts[T];
    int I[T];
    int R[T];
}
transformed data {
    real x_r[0];
    int x_i[0];
}
parameters {
    real beta;
    real gamma;
}
transformed parameters {
    real y_hat[T, 3];
    y_hat = integrate_ode_rk45(sir, y0, t0, ts, {beta, gamma}, x_r, x_i);
}
model {
    beta ~ normal(0, 10);
    gamma ~ normal(0, 10);
    for (t in 1:T) {
        I[t] ~ poisson(max({y_hat[t, 2], 0.0}));
        R[t] ~ poisson(max({y_hat[t, 3], 0.0}));
    }
}
generated quantities {
    real y_pred[T, 3];
    y_pred = integrate_ode_rk45(sir, y0, t0, ts, {beta, gamma}, x_r, x_i);
}
"""

fig_save = plt.subplots(nrows=2, ncols=2, figsize=(20, 15))
axs = fig_save[1].flatten()

for i, target_pref_no in enumerate(['Tokyo', 'Osaka', 'Aichi', 'Fukuoka']):
    # データ準備
    t_obs = np.arange(1, len(I)+1).tolist()
    t_obs_pred = np.arange(1, len(I)+n_pred+1).tolist()

    stan_data = {
        'T': len(t_obs),
        'y0': [S[0], I[0], R[0]],
        't0': 0,
        'ts': t_obs,
        'I': I.astype(int).tolist(),
        'R': R.astype(int).tolist(),
    }

    # コンパイルとサンプリング
    sm = pystan.StanModel(model_code=model_code)
    fit = sm.sampling(data=stan_data, iter=2000, chains=4, warmup=500, seed=42)

    # 予測結果の取得
    y_hat = fit.extract()['y_hat']
    beta_samples = fit.extract()['beta']
    gamma_samples = fit.extract()['gamma']

# 予測用のデータ
t_obs = np.arange(1, len(I)+1).tolist()
n_pred = 50
t_max = len(I) + n_pred
t_all = np.arange(1, t_max+1).tolist()

# 予測の再計算
stan_data_pred = stan_data.copy()
stan_data_pred['T'] = len(t_all)
stan_data_pred['ts'] = t_all

# 可視化
y_hat = fit.extract()['y_hat']

median = np.median(y_hat, axis=0)
lower = np.percentile(y_hat, 2.5, axis=0)
upper = np.percentile(y_hat, 97.5, axis=0)

ax = axs[i]

# 実データ
ax.plot_date(dates[:len(I)], I, 'or', label='Infected (observation)')
ax.plot_date(dates[:len(R)], R, 'og', label='Recovered (observation)')

# 推定結果（中央値と95%信用区間）
ax.plot(dates, median[:, 1], color='r', alpha=0.7, linestyle='dashed', label='Infected (estimation)')
ax.fill_between(dates, lower[:, 1], upper[:, 1], color='r', alpha=0.1)
ax.plot(dates, median[:, 2], color='g', alpha=0.7, linestyle='dashed', label='Recovered (estimation)')
ax.fill_between(dates, lower[:, 2], upper[:, 2], color='g', alpha=0.1)

ax.legend()
ax.set_xlabel('Date')
ax.set_ylabel('Population')
ax.set_title('Prediction of Covid-19 epidemic in ' + target_pref_str)

plt.tight_layout()
plt.show()