NumPyro による微分方程式のパラメータ推定 – SIR model

!pip install --upgrade jax==0.2.17 jaxlib==0.1.71+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
!pip install numpyro==0.7.2
!pip install arviz==0.11.2
!pip install japanize_matplotlib

import jax
import jax.numpy as jnp
import jax.experimental.ode as ode

import numpyro
import numpyro.distributions as dist

import arviz as az
import numpy as np

import matplotlib.pyplot as plt
import japanize_matplotlib

plt.rcParams['font.size'] = 12

numpyro.set_platform('cpu')
numpyro.set_host_device_count(4)

def dz_dt(z, t, beta, gamma):
        
    S = z[0]
    I = z[1]
    
    dS_dt = -beta * S * I
    dI_dt = beta * S * I - gamma * I
    
    return jnp.stack([dS_dt, dI_dt])

beta = 1.0
gamma = 0.25

t = jnp.arange(20).astype(float)

z = ode.odeint(dz_dt, [0.99, 0.01], t, beta, gamma, rtol=1e-8)

S = z[0]
I = z[1]

fig = plt.figure(figsize=(8, 4))

plt.plot(t, S, color='C0', label='感受性保持者（割合）')
plt.plot(t, I, color='C1', label='感染者（割合）')


plt.xlabel('Time')
plt.ylabel('Rate')
plt.legend();

num_people = 1000
num_data = 8

t_observed = t[:num_data]
I_observed = np.random.poisson(I[:num_data] * num_people)

fig = plt.figure(figsize=(8, 4))

plt.plot(t_observed, I_observed, 'o', color='C2', label='観測データ')

plt.plot(t, S * num_people, '--', color='C0', label='感受性保持者数')
plt.plot(t, I * num_people, '--', color='C1', label='感染者数')

plt.xlabel('Time')
plt.ylabel('Number')
plt.legend();

def model(num_people, t, y=None):
    
    # 感受性保持者（S）と感染者（I）の初期値に関する事前分布
    S_init = numpyro.sample('S_init', dist.Beta(1, 1))
    I_init = numpyro.deterministic('I_init', 1 - S_init)
    z_init = jnp.array([S_init, I_init])
    
    # 感染率（beta）と回復率（gamma）に関する事前分布
    beta = numpyro.sample('beta', dist.HalfNormal(5))
    gamma = numpyro.sample('gamma', dist.HalfNormal(5))
    
    # 微分方程式のソルバー
    z = ode.odeint(dz_dt, z_init, t, beta, gamma, rtol=1e-6, atol=1e-5, mxstep=1000)
    
    # 観測データからの尤度の計算
    numpyro.sample('y', dist.Poisson(z[:, 1] * num_people), obs=y)

%%time

nuts = numpyro.infer.NUTS(model)
mcmc = numpyro.infer.MCMC(nuts, num_warmup=1000, num_samples=3000, num_chains=4, progress_bar=False)

mcmc.run(jax.random.PRNGKey(0), num_people, t_observed, y=I_observed)
trace = mcmc.get_samples()

CPU times: user 13.1 s, sys: 27.1 ms, total: 13.2 s
Wall time: 9.39 s

idata = az.from_numpyro(mcmc)

az.plot_trace(idata);

az.summary(idata)

fig, axes = plt.subplots(2, 2, figsize=(10, 8))

az.plot_posterior(idata, ax=axes);

t_pred = jnp.arange(30).astype(float)

predictive = numpyro.infer.Predictive(model, trace)
ppc_samples = predictive(jax.random.PRNGKey(2), num_people, t_pred)

y_pred = ppc_samples['y']

mu_pred = jnp.mean(y_pred, 0)
pi_pred = jnp.percentile(y_pred, (5, 95), 0)

plt.figure(figsize=(9, 5))

plt.plot(t_observed, I_observed, 'o', color='C2', label='感染者数（観測値）')

plt.plot(t_pred, mu_pred, '-.', color='C1', label="感染者数（予測値）")
plt.fill_between(t_pred, pi_pred[0, :], pi_pred[1, :], color='C1', alpha=0.2)

plt.title('事後予測分布（90%-Bayesian Prediction Interval）')
plt.xlabel('Time')
plt.ylabel('Number')
plt.legend();