ベイズマックスへようこそ。
統計に関するあれやこれやを自由気ままに投稿します。
全記事一覧はこちらから。

2020年09月20日

StanでWiener分布からの乱数を得る方法(+累積分布関数ほか)

 Stanには,Drift Diffusion Modelを実装するのにも使えるWiener分布の確率密度関数が実装されています。サンプリングステートメントとしてのy ~ wiener(alpha, tau, beta, delta)と対数確率密度を計算するwiener_lpdf(y | alpha, tau, beta, delta)を使えばデータをWiener分布に当てはめることは容易に行えますが,Wiener分布からの乱数を発生させる関数wiener_rng()は実装されていません。Stanで得たMCMCサンプルを使ってR側でRWiener::rwiener()関数を使えば乱数を得ることはできますが,事後予測分布を作りたいときにはStan側で実装できたほうが便利だと思います。そんな訳で,Stan用のwiener_rng()を自作してみました。2通りのやり方を思い付いたので2つ紹介します。1つ目はWiener過程そのものをStan内でシミュレートすることで乱数を得る方法で,2つ目はWiener分布の累積分布関数の逆関数を使う方法です。順番に紹介していきますが,先取りして結論を言ってしまうと,後者のやり方の方が高速かつ安定しそうです。※この記事で紹介しているやり方は必ずしも正しさが保証されたものではありませんので,実際にご自身の研究等で使用される際は自己責任でお願いします。

1. Wiener過程をシミュレートする方法(非推奨)

 RWiener::rwiener()の内部ではWiener過程をシミュレートする方法が用いられているようなので,それを真似してみました。Wiener過程はランダムウォークにおける各ステップの時間間隔Δtを0に近づけていったものなので,tの微小変化量として十分小さなdtの値を適当に与えてランダムウォークを作ります。このランダムウォークにトレンド(ドリフト率δ)を加えた確率過程が初期値αβからスタートし,その値がαを超えたときに上側反応,0を下回ったときに下側反応が生起すると考えます。Stanコードは以下の通り。

testmodel01.stan: wiener_rng()の実装その1(非推奨)
functions{
  //乱数生成: ウィーナー過程をシミュレートして求める方法
  vector wiener_rng(real alpha, real tau, real beta, real delta){
    vector[2] out = rep_vector(-999,2);
    real dt = 10^(-4); //1ステップの時間の微小変化量を指定;小さいほど正確だが時間がかかる
    int N = 1000000; //最大ステップ数;これを越えても閾値に到達しなければエラー
    vector[N] dis;
    for(i in 1:N) dis[i] = normal_rng(0,1);
    dis = delta*dt + sqrt(dt)*dis;
    dis = cumulative_sum(dis) + alpha*beta;
    for(i in 1:N){
      if(dis[i] > alpha){
        out[2] = 1;
        out[1] = i * dt + tau;
        break;
      } else if(dis[i] < 0){
        out[2] = 0;
        out[1] = i * dt + tau;
        break;
      }
    }
    
    return(out); //最大ステップ数以内に閾値を超えなければエラーとして(-999,-999)を返す
  }
}

data{ //今回はwiener_rng()の挙動の確認が目的なのでパラメータを外から与える。
  real<lower=0> alpha;
  real<lower=0> tau;
  real<lower=0,upper=1> beta;
  real delta;
}

parameters{
  real tmp; //エラー防止用の無意味なパラメータ
}

model{
  tmp ~ normal(0,1); //エラー防止用の無意味なサンプリング
}

generated quantities{
  vector[2] x_pred = wiener_rng(alpha,tau,beta,delta); //乱数生成
}

 なお,Wiener分布は反応 (上側か下側か) と時間(非負の実数)の同時分布なので,wiener_rng()も反応と時間の両方を返すように作ってあります。上記のコードではwiener_rng()からの乱数を受け取る生成量x_predが2次元のベクトルになっていることに注意してください。x_pred[1]が時間,xpred[2]が反応(上側なら1,下側なら0)を格納します。上記コードをRで実行するためのコードは以下の通り。

Rで実行する その1
library(tidyverse)
library(RWiener)
library(rstan)

# おまじない----
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

# 上記モデルのコンパイル----
sm <- stan_model(file="testmodel01.stan")

# パラメータの指定----
data <- list(
  alpha = 1,
  tau = 0.5,
  beta = 0.6,
  delta = 1
)

# サンプリングの実行 (単に乱数を発生させるだけなのでwarmupは小さくしている) ----
fit <- sampling(sm
                 ,data = data
                 ,seed = 610
                 ,chains = 4
                 ,iter = 2520
                 ,warmup = 20
)

# 結果の集計 ----
x_pred_stan <- rstan::extract(fit)$x_pred1 %>% as.data.frame() %>% 
  dplyr::rename(q = V1, resp = V2) %>%
  mutate(resp = as.factor(ifelse(resp==1,"upper","lower"))) %>%
  mutate(resp = factor(resp,levels=c("upper","lower")))

# 比較のためにRWiener::rwiener()からの乱数を得る ----
set.seed(8823)
x_pred_rwiener <- rwiener(10000,data$alpha,data$tau,data$beta,data$delta)

# 要約 ----
x_pred_stan %>% summary()
#       q             resp     
# Min.   :0.5092   upper:8051  
# 1st Qu.:0.5833   lower:1949  
# Median :0.6555               
# Mean   :0.7127               
# 3rd Qu.:0.7829               
# Max.   :2.3316   

x_pred_rwiener %>% summary()
#       q             resp     
# Min.   :0.5105   upper:8085  
# 1st Qu.:0.5805   lower:1915  
# Median :0.6509               
# Mean   :0.7090               
# 3rd Qu.:0.7760               
# Max.   :2.2582     

# プロット ----
wiener_plot(x_pred_stan)
wiener_plot(x_pred_rwiener)

 上記のコードに記している通り,要約統計量の値はRWiener::rwiener()の結果とほぼ一致しました。下図は分布をプロットものですが,こちらも見た目には問題なさそうです。

comp_wiener_rng01.png

 ところが,このやり方には問題があって,単に乱数を発生させるだけなのに計算にかなり時間がかかってしまいます。自分の環境では20分近くもかかってしまいました。dt = 10^(-4)という甘めの設定ですらこんなに時間がかかってしまうのは問題ですし,データの単位に応じて適切なdtを設定するのもかなり面倒そうです。Wiener過程の中身を理解するのには役立ちますが,実用的にはいまいちです。もしかしたらもっとうまい書き方があるのかもしれませんが,それよりは別の方法を試してみようと思います。

2. 累積分布関数の逆関数を利用する方法

 もうひとつの乱数生成のやり方として,累積分布関数の逆関数を利用する方法を紹介します。累積分布関数というのは,ある確率分布f(x)があったときに,任意の値qを入れると確率P(x <= q)を返す関数のことです。その逆関数は確率pを入れると値qを返す関数です。確率pの範囲が[0,1]であることと,pとqが一対一対応することを利用すれば,[0,1]の一様分布から発生させた乱数をこの逆関数に代入することで,元の確率分布からの乱数qを得ることができます。つまり,Wiener分布の累積分布関数の逆関数が分かれば一様分布を使ってWiener分布からの乱数を得ることができるのですが,僕が調べた限り,Wiener分布の累積分布関数の逆関数の数式を見つけることはできませんでした。RWiener::qwiener()の内部でも直接qを求めるのではなくpwiener()を利用して近似的にqを計算しているようなので,おそらく解析的に解くのは難しいのだと思います。RWiener::qwiener()の内部のアルゴリズムを真似すれば,逆関数が未知でも元の累積分布関数を使って近似的にqを計算することができるので,その方針で関数を作ってみました。Stanには累積分布関数wiener_cdf()も実装されていないので,これも自作しました。Wiener分布にはlarge-time representation(長い時間のときに収束が速い)とshort-time representation(短い時間のときに収束が速い)の二通りがありますが,ここではlarge-time representationを採用しました。下側の累積分布関数は次の式で表されます (Wabersich & Vandekerckhove, 2014より)。

wiener_cdf_formula.png

なお,この式のP0は下側選択確率です。これについては前回の記事で解説しているのでそちらをご覧ください。
Drift Diffusion Modelのパラメータから選択確率を計算する関数:RとStanによる実装

 累積分布関数の式にはシグマ記号が含まれており,kに1から+∞までの整数を代入したときの総和を取らなければなりません。しかし,許容する誤差をεと決めてやれば,+∞の代わりに高々有限のKLまでの総和を使って近似できることが分かっています。RWienerパッケージではε=10^(-10)に設定されています。今回実装する関数もこれに倣います。KLは以下の不等式を満たす整数であればOKです(再びWabersich & Vandekerckhove, 2014より)。

wiener_K_formula.png

 ここまでの内容をまとめると次のようになります:(1) Wiener分布の累積分布関数の逆関数を利用すればWiener分布からの乱数qを得ることが出来る。(2)逆関数を導出するのは難しいので累積分布関数を使って近似的にqを得ることにする。(3)累積分布関数もStanには実装されていないので自分で関数を実装する必要がある。(4)累積分布関数を計算するためには選択確率の関数が必要になる(前回の記事で作成済み)。ということで,全ての武器が揃ったのでStanに実装してみましょう。Stanコードは以下の通りです。

testmodel02.stan: wiener_rng()の実装その2
functions{
  //上側選択確率を返す関数
  real wiener_upper_cp(real alpha, real beta, real delta){
    real p_lower;
    if(delta == 0){
      p_lower = 1 - beta;
    } else{
      real tmp = -2*delta*alpha*(1-beta);
      p_lower = (1-exp(tmp)) / (exp(2*delta*alpha*beta) - exp(tmp));
    }
    return(1-p_lower);
  }
  
  //上側の累積確率を返す関数
  real wiener_upper_cdf(real x, real alpha, real tau, real beta, real delta){
    real p;
    if(x < tau){
      p = 0;
    } else{
      real beta2 = 1-beta;
      real delta2 = -delta;
      real epcilon = 10^(-10);
      real K1 = 1 / (x - tau) * (alpha / pi())^2;
      real K2 = -K1 * (
        log(
          0.5 * epcilon * pi() * (x - tau) * (delta2^2 + pi()^2 / alpha^2)
          ) 
          + delta2*alpha*beta2 + 0.5*delta2^2*(x-tau)
        );
      real k_max = fmax(fmax(sqrt(K1),sqrt(K2)),1.0);
      real P0 = 1 - wiener_upper_cp(alpha,beta2,delta2);
      real A = 2*pi() / alpha^2 * exp(-alpha*beta2*delta2 - (delta2^2 * (x-tau))/2 );
      real B = 0;
      int k = 1;
      while(k < (k_max + 1)){
        B += k * sin(pi()*k*beta2) / (delta2^2+(k*pi()/alpha)^2) * exp(-0.5 * (k*pi()/alpha)^2 * (x-tau));
        k += 1;
        if(k > 10000){
          P0 = -999; //1万回ループしても終わらなければエラーとして-999を返す。
          break;
        } 
      }
      p = P0 - A*B;
    }
    return(p);
  }
  
  //上側の累積確率を返す関数の逆関数
  real wiener_upper_cdf_inv(real p, real alpha, real tau, real beta, real delta){
    real q = 1;
    real pmin = 0;
    real pmax = wiener_upper_cp(alpha,beta,delta);
    real qmin = 0;
    real qmax = 10^(3); //時間の取り得る最大値を指定
    real epcilon = 10^(-10); //許容できる最大の誤差を指定
    
    if(p > pmax){
      q = -999; //与えられたpが選択確率を越えた場合にはエラーとして-999を返す
    } else{
      real pmid = wiener_upper_cdf(q,alpha,tau,beta,delta);
      int count = 0;
      while(fabs(p-pmid) > epcilon){
        if(p <= pmid){
          pmax = pmid;
          qmax = q;
          q = (qmax+qmin)/2;
        } else{
          pmin = pmid;
          qmin = q;
          q = (qmax+qmin)/2;
        }
        pmid = wiener_upper_cdf(q,alpha,tau,beta,delta);
        count += 1;
        if(count >= 10000){
          q = -9999; //10000回以内に収束しなかったらエラーとして-9999を返す
          break;
        }
      }
    }
    return(q);
  }
  
  //乱数生成: 逆累積確率関数を利用して求める方法
  vector wiener_rng(real alpha, real tau, real beta, real delta){
    vector[2] out;
    real rand = uniform_rng(0,1);
    real cp =  wiener_upper_cp(alpha,beta,delta);
    if(rand <= cp){
      out[2] = 1;
      out[1] = wiener_upper_cdf_inv(rand,alpha,tau,beta,delta);
    } else{
      out[2] = 0;
      out[1] = wiener_upper_cdf_inv(1-rand,alpha,tau,1-beta,-delta);
    }
    return(out);
  }
}

data{ //今回はwiener_rng()の挙動の確認が目的なのでパラメータを外から与える。
  real<lower=0> alpha;
  real<lower=0> tau;
  real<lower=0,upper=1> beta;
  real delta;
}

parameters{
  real tmp; ////エラー防止用の無意味なパラメータ
}

model{
  tmp ~ normal(0,1); //エラー防止用の無意味なサンプリング
}

generated quantities{
  vector[2] x_pred = wiener_rng(alpha,tau,beta,delta); //乱数生成
}

 実行するためのRコードは以下の通り。さっきのコードとファイル名以外全く同じです。

Rで実行する
library(tidyverse)
library(RWiener)
library(rstan)

# おまじない----
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

# 上記モデルのコンパイル----
sm <- stan_model(file="testmodel02.stan")

# パラメータの指定----
data <- list(
  alpha = 1,
  tau = 0.5,
  beta = 0.6,
  delta = 1
)

# サンプリングの実行 (単に乱数を発生させるだけなのでwarmupは小さくしている) ----
fit <- sampling(sm
                 ,data = data
                 ,seed = 610
                 ,chains = 4
                 ,iter = 2520
                 ,warmup = 20
)

# 結果の集計 ----
x_pred_stan <- rstan::extract(fit)$x_pred1 %>% as.data.frame() %>% 
  dplyr::rename(q = V1, resp = V2) %>%
  mutate(resp = as.factor(ifelse(resp==1,"upper","lower"))) %>%
  mutate(resp = factor(resp,levels=c("upper","lower")))

# 比較のためにRWiener::rwiener()からの乱数を得る ----
set.seed(8823)
x_pred_rwiener <- rwiener(10000,data$alpha,data$tau,data$beta,data$delta)

# 要約 ----
x_pred_stan %>% summary()
#       q             resp     
# Min.   :0.5089   upper:8076  
# 1st Qu.:0.5805   lower:1924  
# Median :0.6498               
# Mean   :0.7080               
# 3rd Qu.:0.7769               
# Max.   :2.4453

x_pred_rwiener %>% summary()
#       q             resp     
# Min.   :0.5105   upper:8085  
# 1st Qu.:0.5805   lower:1915  
# Median :0.6509               
# Mean   :0.7090               
# 3rd Qu.:0.7760               
# Max.   :2.2582     

# プロット ----
wiener_plot(x_pred_stan)
wiener_plot(x_pred_rwiener)

 要約統計量はちゃんと一致しているようです。下に示す分布のプロットを見ても特に問題なさそうです。

comp_wiener_rng02.png

 また,今回は計算にかかる時間もかなり短かったです。僕の環境では25秒で終わりました。実装するならこちらのやり方の方が良さそうですね。

3. おまけ:Rで作ったWiener分布関連の自作関数もろもろ

 Stanのwiener_rng()を自作することを目的に文献を探したりシミュレーションをしたりした結果,だいぶWiener分布のことが分かってきたように思います。せっかくなので,その過程で生まれたRの自作関数をここに貼っておこうと思います。基本的にRWienerの関数の下位互換(計算速度が遅いし正確性も不明)なので実用性は皆無だと思いますが,Wiener過程の中身を知る上で勉強の役に立つかもしれません。もし変なところがあったらご指摘頂けると嬉しいです。

おまけのR関数
# ウィーナー分布の上側・下側選択率を求める関数 ----
wiener_cp <- function(alpha,beta,delta,resp="upper",tau=NULL){
  if(resp!="upper" & resp!= "lower"){
    resp <- "upper"
    print("Warining: resp must be either 'upper' or 'lower.' The probability at the upper boundary was calculated.")
  }
  if(delta == 0){
    p <- 1-beta
  } else{
    tmp <- -2*delta*alpha*(1-beta)
    p <- (1-exp(tmp)) / (exp(2*delta*alpha*beta) - exp(tmp))
  }
  
  if(resp=="upper"){
    p <- 1-p
  }
  return(p)
}

# ウィーナー分布の累積分布関数(RWiener::pwiener()の下位互換) ----
pwiener2 <- function(x,alpha,tau,beta,delta,resp="upper",epcilon=10^(-10)){
  if(resp!="upper" & resp!= "lower"){
    resp <- "upper"
    print("Warining: resp must be either 'upper' or 'lower.' The probability at the upper boundary was calculated.")
  }
  if(resp=="upper"){
    beta <- 1-beta
    delta <- -delta
  }
  
  if(x < tau){
    return(0)
  } else{
    K1 <- 1 / (x - tau) * (alpha / pi)^2
    K2 <- -K1 * (log(0.5 * epcilon * pi * (x - tau) * (delta^2 + pi^2 / alpha^2)) + delta*alpha*beta + 0.5*delta^2*(x-tau))
    k_max <- ceiling(max(sqrt(K1),sqrt(max(K2,1))))
    
    P0 <- wiener_cp(alpha,beta,delta,resp="lower")
    A <- 2*pi / alpha^2 * exp(-alpha*beta*delta - (delta^2 * (x-tau))/2 )
    B <- 0
    for(k in 1:k_max){
      B <- B + k*sin(pi*k*beta) / (delta^2+(k*pi/alpha)^2) * exp(-0.5 * (k*pi/alpha)^2 * (x-tau))
    }
    return(P0 - A*B)
  }
}

# Wiener分布の累積分布関数の逆関数 (RWiener::qwiener()の下位互換) ----
qwiener2 <- function(p,alpha,tau,beta,delta,resp="upper",epcilon=10^(-10)){
  if(resp!="upper" & resp!= "lower"){
    resp <- "upper"
    print("Warining: resp must be either 'upper' or 'lower.' The probability at the upper boundary was calculated.")
  }
  if(resp=="upper"){
    beta <- 1-beta
    delta <- -delta
  }
  
  q <- 1
  pmin <- 0
  pmax <-  wiener_cp(alpha,beta,delta,resp="lower")
  qmin <- 0
  qmax <- 10^(10)
  
  if(p > pmax){
    q <- NA
  } else{
    pmid <- pwiener2(q,alpha,tau,beta,delta,resp="lower")
    
    while(abs(p-pmid) > epcilon){
      pmid <- pwiener2(q,alpha,tau,beta,delta,resp="lower")
      if(p <= pmid){
        pmax <- pmid
        qmax <- q
        q <- qmin + (qmax-qmin)/2
      } else{
        pmin <- pmid
        qmin <- q
        q <- qmin + (qmax-qmin)/2
      }
    }
  }
  
  return(q)
}

# Wiener分布からの乱数をシミュレーションで発生させる関数 (RWiener::rwiener()の下位互換;計算時間めっちゃ長い) ----
rwiener2 <- function(n,alpha,tau,beta,delta){
  dt <- 0.0001 #1ステップの時間[s]
  t_max <- 5 #非決定時間以降の最長時間[s]
  q <- rep(0,n)
  resp <- q
  
  N <- (t_max - tau)/dt
  
  for(i in 1:n){
    dis <- delta*dt + sqrt(dt)*rnorm(N, 0, 1);
    dis <- cumsum(dis) + alpha*beta;
    
    upper_first <- ifelse(sum(dis > alpha) > 0, which.max(dis > alpha) * dt, t_max*10)
    lower_first <- ifelse(sum(dis < 0) > 0, which.max(dis < 0) * dt, t_max*10)
    
    q[i] <- min(upper_first,lower_first,t_max)
    resp[i] <- which.min(c(lower_first,upper_first)) - 1 #upper = 1, lower = 0
    if(q[i] == t_max){
      q[i] <- NA
      resp[i] <- NA
    } else{
      q[i] <- q[i] + tau
    }
  }
  out <- data.frame(
    q = q,
    resp = resp
  ) %>%
    mutate(resp = as.factor(ifelse(resp==1,"upper","lower"))) %>%
    mutate(resp = factor(resp,levels=c("upper","lower")))
  return(out)
}

# Wiener分布からの乱数を累積分布関数の逆関数から得る関数 (なぜかRWiener::rwiener()より高速) ----
rwiener3 <- function(n,alpha,tau,beta,delta){
  S <- 1/alpha^2
  alpha2 <- alpha*sqrt(S)
  delta2 <- delta/sqrt(S)
  tau2 <- tau*S
  
  p_upper <- wiener_cp(alpha2,beta,delta2,resp="upper")
  rand <- runif(n,0,1)
  
  q <- rep(0,n)
  resp <- as.numeric(rand < p_upper)
  for(i in 1:n){
    q[i] <- ifelse(resp[i]==1,
                   qwiener2(rand[i],alpha2,tau2,beta,delta2,resp="upper"),
                   qwiener2(1-rand[i],alpha2,tau2,beta,delta2,resp="lower"))
  }
  
  out <- data.frame(
    q = q / S,
    resp = resp
  ) %>%
    mutate(resp = as.factor(ifelse(resp==1,"upper","lower"))) %>%
    mutate(resp = factor(resp,levels=c("upper","lower")))
  return(out)
}
posted by mutopsy at 21:49 | 統計