Add pcount kfold support

3 files changed, 59 insertions, 13 deletions

diff --git a/R/RcppExports.R b/R/RcppExports.R
index fc60ed1..8a16411 100644
--- a/R/RcppExports.R
+++ b/R/RcppExports.R
@@ -13,6 +13,10 @@ get_loglik_occu <- function(y, M, si, psimat, pmat, Kmin) {
     .Call(`_ubms_get_loglik_occu`, y, M, si, psimat, pmat, Kmin)
 }
 
+get_loglik_pcount <- function(y, M, si, lammat, pmat, K, Kmin) {
+    .Call(`_ubms_get_loglik_pcount`, y, M, si, lammat, pmat, K, Kmin)
+}
+
 simz_pcount <- function(y, lam_post, p_post, K, Kmin, kvals) {
     .Call(`_ubms_simz_pcount`, y, lam_post, p_post, K, Kmin, kvals)
 }
diff --git a/R/kfold.R b/R/kfold.R
index 51ce328..27d25f1 100644
--- a/R/kfold.R
+++ b/R/kfold.R
@@ -57,3 +57,9 @@ setMethod("get_loglik", "ubmsFitOccu", function(object, inps, ...){
   p <- t(posterior_linpred(object, submodel="det", transform=TRUE))
   get_loglik_occu(inps$y, inps$M, inps$si-1, psi, p, inps$Kmin)
 })
+
+setMethod("get_loglik", "ubmsFitPcount", function(object, inps, ...){
+  lam <- t(posterior_linpred(object, submodel="state", transform=TRUE))
+  p <- t(posterior_linpred(object, submodel="det", transform=TRUE))
+  get_loglik_pcount(inps$y, inps$M, inps$si-1, lam, p, inps$K, inps$Kmin)
+})
diff --git a/src/kfold.cpp b/src/kfold.cpp
index f454480..c065819 100644
--- a/src/kfold.cpp
+++ b/src/kfold.cpp
@@ -9,38 +9,74 @@ arma::mat get_loglik_occu(arma::vec y, int M, arma::imat si, arma::mat psimat,
 
   int S = psimat.n_cols;
   mat out(S,M);
-
-  vec psi_s(psimat.n_rows);
   vec p_s(pmat.n_rows);
-
-  vec psite;
-  double psi;
   int J;
   double cp;
-
   vec psub, ysub;
 
-  for (unsigned s = 0; s < S; s++){
+  for (int s = 0; s < S; s++){
 
-    psi_s = psimat.col(s);
     p_s = pmat.col(s);
 
-    for (unsigned m = 0; m < M; m++){
+    for (int m = 0; m < M; m++){
 
       ysub = y.subvec(si(m,0), si(m,1));
-      psi = psi_s(m);
       psub = p_s.subvec(si(m,0), si(m,1));
       J = psub.size();
       cp = 1.0;
 
-      for (unsigned j = 0; j < J; j++){
+      for (int j = 0; j < J; j++){
         cp *= pow(psub(j), ysub(j)) * pow(1-psub(j), 1-ysub(j));
       }
 
       if(Kmin(m) == 1){
-        out(s,m) = log(cp * psi + DBL_MIN);
+        out(s,m) = log(cp * psimat(m,s) + DBL_MIN);
       } else if(Kmin(m) == 0){
-        out(s,m) = log(cp * psi + (1-psi) + DBL_MIN);
+        out(s,m) = log(cp * psimat(m,s) + (1-psimat(m,s)) + DBL_MIN);
+      }
+    }
+  }
+
+  return out;
+}
+
+
+// [[Rcpp::export]]
+arma::mat get_loglik_pcount(arma::vec y, int M, arma::imat si,
+                            arma::mat lammat, arma::mat pmat, int K,
+                            arma::ivec Kmin){
+
+  int S = lammat.n_cols;
+  mat out(S,M);
+  vec p_s(pmat.n_rows);
+  vec psub, ysub;
+  int J;
+  double fac, ff, N, ky, numN;
+
+  for (int s = 0; s < S; s++){
+
+    p_s = pmat.col(s);
+
+    for (int m = 0; m < M; m++){
+      ysub = y.subvec(si(m,0), si(m,1));
+      J = ysub.size();
+      psub = p_s.subvec(si(m,0), si(m,1));
+      fac = 1.0;
+      ff = lammat(m,s) * prod(1 - psub);
+      numN = K - Kmin(m);
+
+      for (int i = 1; i < (numN + 1); i++){
+        N = K - i + 1;
+        ky = 1.0;
+        for (int j = 0; j < J; j++){
+          ky *= N / (N - ysub(j));
+        }
+        fac = 1 + fac * ff * ky / N;
+      }
+
+      out(s,m) = log(fac) + Rf_dpois(Kmin(m), lammat(m,s), true);
+      for (int j = 0; j < J; j ++){
+        out(s,m) += Rf_dbinom(ysub(j), Kmin(m), psub(j), true);
       }
     }
   }