NumPyro による微分方程式のパラメータ推定 – バネマス系

!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'] = 14

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

def f(t, F=1):
    
    return jnp.where(t < 2, F, 0)

def dz_dt(z, t, m, c, k):
    
    y = z[0]
    v = z[1]
    
    dy_dt = v
    dv_dt = (- k * y - c * v + f(t)) / m
        
    return jnp.stack([dy_dt, dv_dt])

m = 1.0 # 質量
c = 1.0 # 減衰係数
k = 1.0 # バネ定数

t_true = jnp.arange(0, 10, 0.1).astype(float)
z_init = jnp.array([0, 0]).astype(float)

z = ode.odeint(dz_dt, z_init, t_true, m, c, k)

y_true = z[:, 0]
v_true = z[:, 1]

plt.plot(t_true, y_true)
plt.title('Mass-Spring-Dumper System')
plt.xlabel('時間 [s]')
plt.ylabel('変位 [m]');

np.random.seed(0)

t_observed = t_true[::5]
y_observed = np.random.normal(y_true[::5], 0.05)

plt.plot(t_true, y_true)
plt.plot(t_observed, y_observed, 'o')

plt.title('Mass-Spring-Dumper System')
plt.xlabel('時間 [s]')
plt.ylabel('変位 [m]');

def model(t, y_observed=None):
    
    # 変位（y）と速度（v）の初期値に関する事前分布
    y_init = numpyro.sample('y_init', dist.Normal(0, 10))
    v_init = numpyro.sample('v_init', dist.Normal(0, 10))
    
    z_init = jnp.stack([y_init, v_init])
    
    # 減衰係数（c）とバネ定数（k）に関する事前分布
    c = numpyro.sample('c', dist.HalfNormal(10))
    k = numpyro.sample('k', dist.HalfNormal(10))
    
    # 微分方程式のソルバー
    z = ode.odeint(dz_dt, z_init, t, m, c, k)
    
    # 観測データからの尤度の計算
    sd_y = numpyro.sample('sigma', dist.HalfNormal(10))
    numpyro.sample('y', dist.Normal(z[:, 0], sd_y), obs=y_observed)

%%time

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

mcmc.run(jax.random.PRNGKey(0), t_observed, y_observed=y_observed)
mcmc_samples = mcmc.get_samples()

idata = az.from_numpyro(mcmc)

CPU times: user 19.2 s, sys: 44 ms, total: 19.3 s
Wall time: 13.3 s

az.plot_trace(idata);

az.summary(idata)

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

az.plot_posterior(idata, var_names=['y_init', 'v_init', 'c', 'k'], ax=axes)

fig.subplots_adjust(hspace=0.4)

t_pred = jnp.arange(0, 10, 0.1).astype(float)

predictive = numpyro.infer.Predictive(model, mcmc_samples)
ppc_samples = predictive(jax.random.PRNGKey(2), 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=(8, 6))

plt.plot(t_observed, y_observed, 'o', color='C1', label='観測値')
plt.plot(t_true, y_true, '--', color='C2', label='真値')

plt.plot(t_pred, mu_pred, '-.', color='C0', label='予測値 (平均)')
plt.fill_between(t_pred, pi_pred[0, :], pi_pred[1, :], color='C0', alpha=0.2)

plt.title('事後予測分布')
plt.xlabel('時間 [s]')
plt.ylabel('変位 [m]')

plt.legend();