Add loglik for occuTTD and fix some bugs

7 files changed, 97 insertions, 10 deletions

diff --git a/R/RcppExports.R b/R/RcppExports.R
index 504f6d6..ae97b81 100644
--- a/R/RcppExports.R
+++ b/R/RcppExports.R
@@ -29,6 +29,10 @@ get_loglik_colext <- function(y, M, Tsamp, J, si, psicube, phicube, pmat, nd) {
     .Call(`_ubms_get_loglik_colext`, y, M, Tsamp, J, si, psicube, phicube, pmat, nd)
 }
 
+get_loglik_occuTTD <- function(y, M, si, psimat, lammat, k, delta, ydist) {
+    .Call(`_ubms_get_loglik_occuTTD`, y, M, si, psimat, lammat, k, delta, ydist)
+}
+
 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/fit.R b/R/fit.R
index d7f60fc..dedb735 100644
--- a/R/fit.R
+++ b/R/fit.R
@@ -26,6 +26,8 @@ ubmsFit <- function(model, call, data, response, submodels, ...){
 
   #Fit model
   fit <- fit_model(model, response, submodels, ...)
+  # Exit early if just returning Stan inputs
+  if(check_return_inputs(...)) return(fit)
 
   #Remove placeholder submodels
   submodels <- remove_placeholders(submodels)
@@ -43,6 +45,17 @@ remove_placeholders <- function(submodels){
   submodels
 }
 
+# Function to check if just Stan inputs should be returned
+# Used by kfold method
+# Pass return_inputs=TRUE in ... when calling stan_*
+check_return_inputs <- function(...){
+  args <- list(...)
+  if("return_inputs" %in% names(args)){
+    if(args$return_inputs) return(TRUE)
+  }
+  FALSE
+}
+
 fit_class <- function(mod){
   cap <- paste0(toupper(substr(mod,1,1)), substr(mod,2,nchar(mod)))
   paste0("ubmsFit",cap)
@@ -53,6 +66,11 @@ fit_class <- function(mod){
 fit_model <- function(name, response, submodels, ...){
   model <- name_to_stanmodel(name, submodels)
   inp <- build_stan_inputs(name, response, submodels)
+  # Should just Stan inputs be returned?
+  if(check_return_inputs(...)){
+    inp$submodels <- submodels
+    return(inp)
+  }
   mod <- stanmodels[[model]]
   mod@model_name <- name
   fit <- sampling(mod, data=inp$stan_data, pars=inp$pars, ...)
diff --git a/R/inputs.R b/R/inputs.R
index 81735db..a3845d6 100644
--- a/R/inputs.R
+++ b/R/inputs.R
@@ -3,7 +3,7 @@
   model_code <- name_to_modelcode(name)
   y_data <- get_stan_data(response)
 
-  pars <- get_pars(submodels, name %in% c("occuTTD","distsamp"))
+  pars <- get_pars(submodels, name %in% c("distsamp"))
   submodels <- unname(submodels@submodels)
   types <- sapply(submodels, function(x) x@type)
   submodel_data <- lapply(submodels, get_stan_data)
diff --git a/R/kfold.R b/R/kfold.R
index f49c047..97ea675 100644
--- a/R/kfold.R
+++ b/R/kfold.R
@@ -40,12 +40,12 @@ ll_fold <- function(i, object, folds){
   train_data <- object@data[which(!folds==i),]
   cl$data <- train_data
   cl$refresh <- 0
-  refit <- suppressWarnings(eval(cl, parent.frame(3)))
+  refit <- suppressWarnings(eval(cl))
 
   test_data <- object@data[which(folds==i),]
   cl$data <- test_data
   cl$return_inputs <- TRUE
-  inps <- eval(cl, parent.frame(3))
+  inps <- eval(cl)
   refit@submodels <- inps$submodels
   ll <- get_loglik(refit, inps)
 
diff --git a/R/loglik.R b/R/loglik.R
index bccdf47..1309d64 100644
--- a/R/loglik.R
+++ b/R/loglik.R
@@ -85,6 +85,18 @@ setMethod("get_loglik", "ubmsFitColext", function(object, inps, ...){
                      phicube, pmat, 1 - inps$stan_data$Kmin)
 })
 
+#' @include occuTTD.R
+setMethod("get_loglik", "ubmsFitOccuTTD", function(object, inps, ...){
+  psi <- calculate_par(object, inps, submodel="state")
+  lam <- calculate_par(object, inps, submodel="det")
+  k <- rep(0, ncol(psi))
+  if(inps$stan_data$y_dist == 3){
+    k <- extract(object, "beta_shape", permute=FALSE)
+    k <- exp(as.vector(k))
+  }
+  get_loglik_occuTTD(inps$stan_data$aux2, inps$stan_data$M, inps$stan_data$si-1,
+                     psi, lam, k, inps$stan_data$aux1, inps$stan_data$y_dist)
+})
 
 #' Extract Pointwise Log-likelihood From Model
 #'
@@ -126,12 +138,6 @@ setMethod("extract_log_lik", "ubmsFit",
   ll
 })
 
-#' @include occuTTD.R
-setMethod("extract_log_lik", "ubmsFitOccuTTD",
-          function(object, parameter_name = "log_lik", merge_chains=TRUE){
-  loo::extract_log_lik(object@stanfit, parameter_name, merge_chains)
-})
-
 #' @include distsamp.R
 setMethod("extract_log_lik", "ubmsFitDistsamp",
           function(object, parameter_name = "log_lik", merge_chains=TRUE){
diff --git a/R/submodel.R b/R/submodel.R
index b7c2c13..79369bb 100644
--- a/R/submodel.R
+++ b/R/submodel.R
@@ -94,7 +94,7 @@ setMethod("model.matrix", "ubmsSubmodel",
   formula <- lme4::nobars(object@formula)
   out <- model.matrix(formula, mf)
   if(na.rm) out <- out[!object@missing,,drop=FALSE]
-  if(warn){
+  if(warn & !all(is.na(out))){
     max_cov <- max(out, na.rm=TRUE)
     if(max_cov > 4){
       warning(paste0("Covariates possibly not standardized (max value = ", max_cov,
diff --git a/src/loglik.cpp b/src/loglik.cpp
index 16b3b88..78a97b9 100644
--- a/src/loglik.cpp
+++ b/src/loglik.cpp
@@ -285,3 +285,62 @@ arma::mat get_loglik_colext(arma::vec y, int M, arma::ivec Tsamp, arma::imat J,
 
   return out;
 }
+
+
+arma::vec ttd_prob_exp(arma::vec y, arma::vec lam, arma::ivec delta){
+  int J = y.size();
+  vec e_lamt(J);
+
+  for (int j=0; j < J; j++){
+    e_lamt(j) = pow(lam(j), delta(j)) * exp(-lam(j)*y(j));
+  }
+  return e_lamt;
+}
+
+arma::vec ttd_prob_weib(arma::vec y, arma::vec lam, arma::ivec delta, double k){
+  int J = y.size();
+  vec e_lamt(J);
+  for (int j = 0; j < J; j++){
+    e_lamt(j) = pow(k*lam(j)*pow(lam(j)*y(j), (k-1)), delta(j)) *
+                exp(-1*pow(lam(j)*y(j), k));
+  }
+  return e_lamt;
+}
+
+
+// [[Rcpp::export]]
+arma::mat get_loglik_occuTTD(arma::vec y, int M, arma::imat si, arma::mat psimat,
+                             arma::mat lammat, arma::vec k,
+                             arma::ivec delta, int ydist){
+
+  int S = psimat.n_cols;
+  mat out(S,M);
+  vec lam_s(lammat.n_rows);
+  vec e_lamt;
+  int J;
+  vec y_m, lam_m;
+  ivec delta_m;
+  double lik;
+
+  for (int s = 0; s < S; s++){
+
+    lam_s = lammat.col(s);
+
+    for (int m = 0; m < M; m++){
+      y_m = y.subvec(si(m,0), si(m,1));
+      lam_m = lam_s.subvec(si(m,0), si(m,1));
+      delta_m = delta.subvec(si(m,0), si(m,1));
+
+      if(ydist == 1){
+        e_lamt = ttd_prob_exp(y_m, lam_m, delta_m);
+      } else if(ydist == 3){
+        e_lamt = ttd_prob_weib(y_m, lam_m, delta_m, k(s));
+      }
+
+      lik = psimat(m,s) * prod(e_lamt) + (1-psimat(m,s)) * (1-max(delta_m));
+      out(s,m) = log(lik);
+    }
+  }
+
+  return out;
+}