R/rhatfun.R

In eifer4/stochasticSampling: Stochastic partitioning sampler.

split_rhat_rfun <- function(sims){ #from rstan package
  chains <- ncol(sims)
  n_samples <- nrow(sims)
  half_n <- floor(n_samples/2)
  idx_2nd <- n_samples - half_n + 1
  split_chain_mean <- numeric(chains * 2)
  split_chain_var <- numeric(chains * 2)
  for (i in 1:chains) {
    split_chain_mean[i] <- mean(sims[1:half_n, i], na.rm=TRUE)
    split_chain_var[i] <- var(sims[1:half_n, i], na.rm=TRUE)
    split_chain_mean[chains + i] <- mean(sims[idx_2nd:n_samples,
                                              i], na.rm=TRUE)
    split_chain_var[chains + i] <- var(sims[idx_2nd:n_samples,
                                            i], na.rm=TRUE)
  }
  var_between <- half_n * var(split_chain_mean,na.rm=TRUE)
  var_within <- mean(split_chain_var, na.rm=TRUE)
  sqrt((var_between/var_within + half_n - 1)/half_n)
}

corr_rhat_rfun <- function(sims, n_samples) {
  chains     <- dim(sims)[3]
  half_n     <- n_samples[1]
  n_samples  <- sum(n_samples)
  n_param    <- dim(sims)[1]
  
  split_chain_mean <- matrix(NA, ncol = (chains * 2), nrow = n_param)
  split_chain_var  <- matrix(NA, ncol = (chains * 2), nrow = n_param)
  
  for (i in 1:chains) {
    split_chain_mean[,i] <- sims[,1,i]
    split_chain_var[,i] <- sims[,1,i] * (1 - sims[,1,i])
    split_chain_mean[,chains + i] <- sims[,2,i]
    split_chain_var[,chains + i] <- sims[,2,i] * (1 - sims[,2,i])
  }
  var_between <- half_n * apply(split_chain_mean,1, var,na.rm=TRUE)
  var_within <- rowMeans(split_chain_var, na.rm=TRUE)
  sqrt((var_between/var_within + half_n - 1)/half_n)
}