rstanでの確率分布からのサンプリングの速さを比較する

rstanで自作関数、というかrstanに実装されていない確率分布からサンプリングをしたくてコードを書いていたが、その前にコードの書き方でサンプリングの効率というか速さが違うので速くなる書き方をしよう、という検証。
結論から言うと、実装されている関数ならば、vector 型でサンプリングするのがよく、自作関数を作るとひとつのrealもしくはintでサンプリングするので、遅い。

5つのパターンを比較する。どれも正規分布N(1, 1)からデータを作り、正規分布normal(m, s)で推定する。
code0は、vector型でy ~ normal(m, s)とサンプリングする。
code1は、ひとつひとつy[i] ~ normal(m, s)とサンプリングする。
code2は、target記法を用いてtarget += normal(y[i] | m, s)とサンプリングする。
code3は、自作関数tmp = delta/(2*pi()*lambda^2)^0.5 * 1/(1+((x-xi)/lambda)^2)^0.5 * exp(-0.5*(gamma+delta*asinh((x-xi)/lambda))^2); を定義し、code1のようにひとつひとつサンプリングする。
code4は、code3の自作関数をtarget記法でサンプリングする。
ちなみにrstanで自作関数を利用するとき、func_lpdfで関数を定義し、返り値はlog(return)とする。また、target記法で利用するときは、func_lpdfで記述する。

結果としてはcode0vector型でのサンプリングが一番早かった。rstanに定義されている関数を利用する場合は、ひとつひとつサンプリングするcode1のほうがtarget記法より速いようだった。
自作関数を利用する場合は、target記法のほうが速そうな気がする。
f:id:MikuHatsune:20200312225316p:plain

hoge <- rnorm(300, 1, 1)
code0 <- "
    data {
    int<lower=0> N;
    real y[N];
  }
  parameters {
    real m;
    real<lower=0> s;
  }
  model {
    y ~ normal(m, s);
  }
"
code1 <- "
    data {
    int<lower=0> N;
    real y[N];
  }
  parameters {
    real m;
    real<lower=0> s;
  }
  model {
    for(i in 1:N){
      y[i] ~ normal(m, s);
    }
  }
"
code2 <- "
    data {
    int<lower=0> N;
    real y[N];
  }
  parameters {
    real m;
    real<lower=0> s;
  }
  model {
    for(i in 1:N){
      target += normal_lpdf(y[i] | m, s);
    }
  }
"
code3 <- "
  functions{ //手作り正規分布を定義
    real my_lpdf(real x, real m, real s){
      real tmp;
      tmp = 1/(2*pi()*s^2)^0.5 * exp(-0.5*((x-m)/s)^2);
      return log(tmp);
    }
  }
  data {
    int<lower=0> N;
    real y[N];
  }
  parameters {
    real m;
    real<lower=0> s;
  }
  model {
    for(i in 1:N){
      y[i] ~ my(m, s);
    }
  }
"
code4 <- "
  functions{ //手作り正規分布を定義
    real my_lpdf(real x, real m, real s){
      real tmp;
      tmp = 1/(2*pi()*s^2)^0.5 * exp(-0.5*((x-m)/s)^2);
      return log(tmp);
    }
  }
  data {
    int<lower=0> N;
    real y[N];
  }
  parameters {
    real m;
    real<lower=0> s;
  }
  model {
    for(i in 1:N){
      target += my_lpdf(y[i] | m, s);
    }
  }
"
codes <- list(code0, code1, code2, code3, code4)
names(codes) <- c("vector", "elementwise", "element_target", "self_function", "self_target")

# 一括してコンパイルするが
# コンピュータのメモリが貧弱だと並列にすると死ぬかもしれない
library(rstan)
rstan_options(auto_write=TRUE)
# options(mc.cores=parallel::detectCores())


ms <- mapply(stan_model, model_code=codes)
standata <- list(N=length(hoge), y=hoge)
fits <- mapply(function(z) sampling(z, standata, iter=2000, warmup=1000, chain=100), ms)

elp <- mapply(function(w) colSums(mapply(function(z) attributes(z)$elapsed_time, w@sim$samples)), fits)


library(vioplot)
vioplot(elp, horizontal=TRUE, side="right", yaxt="n", xlab="elapsed time [sec]")
abline(h=seq(ncol(elp)), lty=3)
axis(1)
text(par()$usr[2], seq(ncol(elp))+0.3, colnames(elp), xpd=TRUE, pos=2)