binom_mult.cpp

#include <RcppArmadillo.h>
using namespace Rcpp;
using namespace arma;
// [[Rcpp::depends(RcppArmadillo)]]

rowvec rdirichlet(const vec & alpha){
  int K = alpha.n_rows;
  rowvec y(K);
  for(int k=0;k<K ;k++){
    y[k] = R::rgamma(alpha[k],1); 
  }
  rowvec out = y/sum(y);
  return(out);
}

rowvec rcate(const rowvec & p){
  int K = p.n_cols;
  rowvec cump = cumsum(p);
  rowvec x(K);
  x.fill(0);
  double U = R::runif(0,1);
  if(U<=cump[0]){
    x[0] = 1;
  }else{
    for(int k=1; k<K; k++){
      if(cump[k-1]<U & U<=cump[k]){
        x[k] = 1;
      }
    }
  }
  return(x);
}

vec colSums(const mat & X){
  int nCols = X.n_cols;
  vec out(nCols);
  for(int i = 0; i < nCols; i++){
    out(i) = sum(X.col(i));
  }
  return(out);
}

vec colMeans(const mat & X){
  int nCols = X.n_cols;
  vec out(nCols);
  for(int i = 0; i < nCols; i++){
    out(i) = sum(X.col(i))/X.n_rows;
  }
  return(out);
}

vec colSums(const imat & X){
  int nCols = X.n_cols;
  vec out(nCols);
  for(int i = 0; i < nCols; i++){
    out(i) = sum(X.col(i));
  }
  return(out);
}

vec colMeans(const imat & X){
  int nCols = X.n_cols;
  vec out(nCols);
  for(int i = 0; i < nCols; i++){
    out(i) = sum(X.col(i))/X.n_rows;
  }
  return(out);
}

vec rowSums(const mat & X){
  int nRows = X.n_rows;
  vec out(nRows);
  for(int i = 0; i < nRows; i++){
    out(i) = sum(X.row(i));
  }
  return(out);
}

rowvec softmax(const rowvec & x){
  double tmp = max(x);
  rowvec res = exp(x-tmp)/sum(exp(x-tmp));
  return res;
}

double logsumexp(const rowvec & x){
  double tmp = max(x);
  return tmp + log(sum(exp(x-tmp)));
}

double multi_lpmf(rowvec W, double M, rowvec p){
  int K = W.n_cols;
  double lp = 0;
  for(int k=0; k<K; k++){
    if(!(W[k]==0 && p[k]==0)){
      lp = lp + W[k]*log(p[k])-lgamma(W[k]+1);
    }
  }
  lp = lp + lgamma(M+1);
  return lp;
}


// [[Rcpp::export]]
List binom_mult_Gibbs(vec y, vec s, mat W, int L, vec phiini, mat pini, vec rhoini, double alpha, double gamma, double beta, int iter){
  vec M;
  M = rowSums(W);
  int N = y.n_rows;
  int K = W.n_cols;
  rowvec unnorm(L);
  cube z = zeros(N,L,iter);
  mat ztmp = zeros(N,L);
  mat Wz = zeros(N,L);
  mat phi = zeros(iter,L);
  phi.row(0) = phiini.t();
  cube p = zeros(L,K,iter);
  p.slice(0) = pini;
  mat ptmp = pini;
  mat rho = zeros(iter,L);
  rho.row(0) = rhoini.t();
  mat loglik = zeros(iter-1);
  for(int i=1;i<iter;i++){
    for(int n=0;n<N;n++){
      for(int l=0; l<L; l++){
        unnorm[l] = beta*(std::log(phi.row(i-1)[l]) + R::dbinom(y[n],s[n],rho.row(i-1)[l],true) + multi_lpmf(W.row(n),M[n],ptmp.row(l)));
      }
      loglik(i-1) += logsumexp(unnorm/beta);
      ztmp.row(n) = rcate(softmax(unnorm));
    }
    phi.row(i)  = rdirichlet(colSums(ztmp)+1);
    for(int l=0; l<L; l++){
      rho.row(i)[l] = R::rbeta(beta*sum(ztmp.col(l) % y)+gamma,beta*sum(ztmp.col(l) % (s-y))+gamma);
      rho.row(i) = sort(rho.row(i));
      Wz = W.each_col() % ztmp.col(l);
      ptmp.row(l) = rdirichlet(beta*colSums(Wz)+alpha);
    }
    p.slice(i) = ptmp;
    z.slice(i) = ztmp;
  }
  return List::create(Named("p")=p,_["phi"]=phi,_["rho"]=rho,_["z"]=z,_["loglik"]=loglik);
}