【PyMC3.11.5】pymc.GaussianRandomWalkで表現できない確率過程を記述する方法

PyMCっていろんな密度関数が実装されていて、疑似データを作るなどシミュレーションを行いたいときにとても便利です。

一方で、確率過程に関しては大まかに(Mv-)RW過程、AR(p)過程、GARCH(1,1)過程の3つしか実装されていない有様で、状態空間モデルを記述するときに不便だったりします。

例えば、次のようなデータ生成過程に基づいた疑似データの生成を行いたいたいとします。

(観測)　 $y_t = \alpha_t + \epsilon_t, ~~~ \epsilon_t \sim N(0,\sigma^2_{\epsilon});$

(状態)　 $\alpha_{t} = \mu + \phi (\alpha_{t-1} - c) + u_{t}, ~~~ u_{t} \sim N(0,\sigma^2_{u});$

$t = 1,\dots,T.$

上の特殊例として、 $c=0, \phi=1$ のときは、 $(\mu,\sigma^2_{u},\alpha_0)$ を適当に指定してからPyMCの既存の "GaussianRandomWalk" を使えば、状態方程式の記述に関しては完結します。

$(\mu,\sigma^2_{u},\alpha_0)=(1,3,0)$ と設定した上で、"sample_prior_predictive"関数を使って100回サンプリングした時系列を可視化してみると、

import numpy as np
import pandas as pd
import pymc3 as pm

T = 100
N_SIMS = 100

with pm.Model(coords={'time': range(T)}) as model:
  alpha = pm.GaussianRandomWalk('alpha', mu=1, sigma=3, dims=('time'), init=pm.Uniform.dist(0, 0))
  simulation = pm.sample_prior_predictive(samples=N_SIMS)

ax = pd.DataFrame(simulation['alpha'].transpose()).plot(figsize=(18,6), alpha=0.3, legend=False)

↑こうなる。ここまではよい、PyMC最高。

さて、例えば $|\phi|<1$ （に設定 | を推定）してーよなど、必ずしも GaussianRandomWalk とか GARCH(1,1) process で記述できないけれども、漸化式を記述したい場合もあるかと思います。
そのようなときには、Theanoのscanに再帰を入力してやればいいみたいです。

$(\mu,\sigma^2_{u},\alpha_0)=(5,0.3,0)$ 、特に $(\phi,c)=(0.9, \mu)$ と設定した上で100回サンプリングした時系列を可視化してみます。前回は $(\phi,c)=(1, 0)$ でした。

μ = 5
α_0 = np.float64(0)

φ = 0.9
c = μ

fn_t = lambda u_t, α_t: μ + φ * (α_t - c) + u_t

with pm.Model(coords={'time': range(T)}) as model:
  u = pm.Normal('u', mu=0, sigma=0.3, shape=T)
  α, _ = theano.scan(fn=fn_t, sequences=[u], outputs_info=[α_0], n_steps=T)
  α = pm.Deterministic('α', pm.math.concatenate([pm.Uniform('α_0', lower=α_0, upper=α_0, shape=1), α]))
  
  simulation = pm.sample_prior_predictive(samples=N_SIMS)

ax = pd.DataFrame(simulation['alpha'].transpose()).plot(figsize=(18,6), alpha=0.3, legend=False)

関数 fn_t として漸化式を、e として状態方程式の白色雑音の確率変数列を定義して、theano.scanに突っ込んどります。
結果、 $\alpha_{1:T}$ のシミュレーションが得られました。 $\alpha_{0:T}$ を可視化しています。

おわり