JAGSで対数尤度を得る方法(Stanも)

WAICなどを計算するには，対数尤度のMCMCサンプルが必要。
個々のMCMCサンプルからがんばって計算する方法もあるが，より簡単にはMCMCサンプラーのなかで対数尤度を得ると楽。
Stanではgenerated quantityが容易されているが，JAGSでは少し工夫が必要。

参考
https://sourceforge.net/p/mcmc-jags/discussion/610037/thread/a3bef062/
https://sourceforge.net/p/mcmc-jags/discussion/610037/thread/87902622/

その他参考
ito-hi.blog.so-net.ne.jp
hamada.hatenablog.jp

JAGSのモデルの中でのようなものでも対数尤度が得られる
一変量の分布に関しては，密度関数などがRのように用意されている。

model{
for(i in 1:I) logLik[i] <- log(dnorm(x[i], mu, 1/sigma))
}

そのほか, logdensity.**系の関数が用意されているが，これはマニュアルに書かれていないので必要に応じて調べる。
多分，すべての分布に対してこれらの関数が利用できる？
（未確認）

#
# 相関係数については，以下の頁のものを参考にした。
# https://kyougokumakoto.blogspot.jp/2016/09/jags_91.html
# 
# logdensity.mnormで多変量正規分布の対数密度を得られる。
# https://sourceforge.net/p/mcmc-jags/discussion/610037/thread/a3bef062/
#
library(MASS)
set.seed(1234) 
mu <- c(0,0) 
Sig <- matrix(c(1,0.75,0.75,1),2,2) 
x <- mvrnorm(100,mu,Sig) 
cor.test(x[,1],x[,2])

data = list(x = x, N = nrow(x)) 

# JAGS code 
cat("model{
# likelihood
for(i in 1:N){
x[i,1:2] ~ dmnorm(mu[], prec[ , ]) 

logLik[i] <- logdensity.mnorm(x[i,1:2], mu, prec) 
} 

# corresponding the precision matrix
prec[1:2, 1:2] <- inverse(cov[,]) 

# constructing the covariance matrix
cov[1,1] <- sigma[1]^2
cov[1,2] <- sigma[1]*sigma[2]*rho
cov[2,1] <- sigma[1]*sigma[2]*rho
cov[2,2] <- sigma[2]^2 

#priors
rho ~ dunif(-1,1) 
mu[1] ~ dnorm(0, 1.0e-4) 
mu[2] ~ dnorm(0, 1.0e-4)
sigma[1] ~ dt(0, 1, 1)T(0,)
sigma[2] ~ dt(0, 1, 1)T(0,) 

}", file="cor.jag")

fit_jag <- jags(data=data,
            parameters.to.save=c("mu", "sigma","rho","logLik"), 
            model.file="cor.jag",
            n.chains = 3,
            n.adapt = 100,
            n.iter = 5000,
            n.burnin = 2500,
            n.thin = 2)

library("loo")

# チェインがたくさんある場合には，長い行列をとしてみる？
log_lik1 <- sapply(1:nrow(x), function(i) do.call("rbind",fit_jag$samples)[,paste0("logLik[",i,"]")])
waic(log_lik1)

log_lik2 <- sapply(1:nrow(x), function(i) fit_jag$samples[[1]][,paste0("logLik[",i,"]")])
waic(log_lik2)


library(rstan)

#
# Stanのコード参考：https://qiita.com/hoxo_m/items/0f1b05681f5d6c4b560a
# ベクトルの配列について注意。y[i]で得られるのはi番目の長さが2のベクトル。
# y[i][1]でi番目の配列にあるベクトルの1番目の要素を取り出すということになる。 
#

cat("data{
 int<lower=0> N;
    vector[2] y[N];
    }
    parameters{
    vector[2] mu;
    real<lower=0> var_x;
    real<lower=0> var_y;
    real cov;
    }
    
    transformed parameters{
    matrix[2,2] sigma;
    sigma[1,1] = var_x;
    sigma[2,1] = cov;
    sigma[1,2] = cov;
    sigma[2,2] = var_y;
    }
    
    model{
    for(i in 1:N){
    target += multi_normal_lpdf(y[i]| mu, sigma);
    //target += normal_lpdf(y[i][2]|0, 1);
    //target += normal_lpdf(y[i][1]|0, 1);
    //y[i] ~ multi_normal(mu, sigma);
    }
    }
    
    
    generated quantities {
    real cor;
    real log_lik[N];
    cor = cov/sqrt(var_x * var_y);
    for(i in 1:N) log_lik[i] = multi_normal_lpdf(y[i] | mu, sigma);
    }

", file="cor.stan")
datastan <- list(N=nrow(x), y = x)
fit <- stan(file = "cor.stan", data=datastan, iter = 5000, chains = 2)

log_lik_stan <- extract_log_lik(fit)
str(log_lik_stan)
waic(log_lik_stan)
waic(log_lik1)
waic(log_lik2)

# 値は大体一致している。