Improve efficiency of model methods, tidy code

Author: andrjohns

R/fit.R

@@ -518,12 +518,8 @@ unconstrain_variables <- function(variables) {
          " not provided!", call. = FALSE)
   }
 
-  # Remove zero-length parameters from model_variables, otherwise process_init
-  # warns about missing inputs
-  model_variables$parameters <- model_variables$parameters[nonzero_length_params]
-
-  stan_pars <- process_init(list(variables), num_procs = 1, model_variables)
-  private$model_methods_env_$unconstrain_variables(private$model_methods_env_$model_ptr_, stan_pars)
+  variables_vector <- unlist(variables, recursive = TRUE, use.names = FALSE)
+  private$model_methods_env_$unconstrain_variables(private$model_methods_env_$model_ptr_, variables_vector)
 }
 CmdStanFit$set("public", name = "unconstrain_variables", value = unconstrain_variables)
 
@@ -571,11 +567,11 @@ unconstrain_draws <- function(files = NULL, draws = NULL,
           call. = FALSE)
     }
     if (!is.null(files)) {
-      read_csv <- read_cmdstan_csv(files = files, format = "draws_df")
+      read_csv <- read_cmdstan_csv(files = files, format = "draws_matrix")
       draws <- read_csv$post_warmup_draws
     }
     if (!is.null(draws)) {
-      draws <- maybe_convert_draws_format(draws, "draws_df")
+      draws <- maybe_convert_draws_format(draws, "draws_matrix")
     }
   } else {
     if (is.null(private$draws_)) {
@@ -584,7 +580,7 @@ unconstrain_draws <- function(files = NULL, draws = NULL,
       }
       private$read_csv_(format = "draws_df")
     }
-    draws <- maybe_convert_draws_format(private$draws_, "draws_df")
+    draws <- maybe_convert_draws_format(private$draws_, "draws_matrix")
   }
 
   model_par_names <- self$metadata()$stan_variables[self$metadata()$stan_variables != "lp__"]
@@ -599,19 +595,10 @@ unconstrain_draws <- function(files = NULL, draws = NULL,
   pars <- names(model_variables$parameters[nonzero_length_params])
 
   draws <- posterior::subset_draws(draws, variable = pars)
-  skeleton <- self$variable_skeleton(transformed_parameters = FALSE,
-                                     generated_quantities = FALSE)
-  par_columns <- !(names(draws) %in% c(".chain", ".iteration", ".draw"))
-  meta_columns <- !par_columns
-  unconstrained <- lapply(asplit(draws, 1), function(draw) {
-    par_list <- utils::relist(as.numeric(draw[par_columns]), skeleton)
-    self$unconstrain_variables(variables = par_list)
-  })
-
-  unconstrained <- do.call(rbind.data.frame, unconstrained)
+  unconstrained <- private$model_methods_env_$unconstrain_draws(private$model_methods_env_$model_ptr_, draws)
   uncon_names <- private$model_methods_env_$unconstrained_param_names(private$model_methods_env_$model_ptr_, FALSE, FALSE)
   names(unconstrained) <- repair_variable_names(uncon_names)
-  maybe_convert_draws_format(cbind.data.frame(unconstrained, draws[,meta_columns]), format)
+  maybe_convert_draws_format(unconstrained, format)
 }
 CmdStanFit$set("public", name = "unconstrain_draws", value = unconstrain_draws)
 

R/utils.R

@@ -757,13 +757,6 @@ expose_model_methods <- function(env, verbose = FALSE, hessian = FALSE) {
             readLines(system.file("include", "model_methods.cpp",
                                   package = "cmdstanr", mustWork = TRUE)))
 
-  if (hessian) {
-    code <- c("#include <stan/math/mix.hpp>",
-            code,
-            readLines(system.file("include", "hessian.cpp",
-                                  package = "cmdstanr", mustWork = TRUE)))
-  }
-
   code <- paste(code, collapse = "\n")
   rcpp_source_stan(code, env, verbose)
   invisible(NULL)

inst/include/hessian.cpp

@@ -1,4 +1,5 @@
 #include <Rcpp.h>
+#include <rcpp_eigen_interop.hpp>
 #include <stan/model/model_base.hpp>
 #include <stan/model/log_prob_grad.hpp>
 #include <stan/model/log_prob_propto.hpp>
@@ -26,10 +27,14 @@ using json_data_t = stan::json::json_data;
   return std::make_shared<json_data_t>(data_context);
 }
 
+stan::model::model_base&
+new_model(stan::io::var_context& data_context, unsigned int seed,
+          std::ostream* msg_stream);
+
 // [[Rcpp::export]]
 Rcpp::List model_ptr(std::string data_path, boost::uint32_t seed) {
-  Rcpp::XPtr<stan_model> ptr(
-    new stan_model(*var_context(data_path), seed, &Rcpp::Rcout)
+  Rcpp::XPtr<stan::model::model_base> ptr(
+    &new_model(*var_context(data_path), seed, &Rcpp::Rcout)
   );
   Rcpp::XPtr<boost::ecuyer1988> base_rng(new boost::ecuyer1988(seed));
   return Rcpp::List::create(
@@ -39,37 +44,56 @@ Rcpp::List model_ptr(std::string data_path, boost::uint32_t seed) {
 }
 
 // [[Rcpp::export]]
-double log_prob(SEXP ext_model_ptr, std::vector<double> upars,
-                bool jac_adjust) {
+double log_prob(SEXP ext_model_ptr, Eigen::VectorXd upars, bool jac_adjust) {
   Rcpp::XPtr<stan::model::model_base> ptr(ext_model_ptr);
-  std::vector<int> params_i;
   if (jac_adjust) {
-    return stan::model::log_prob_propto<true>(*ptr.get(), upars, params_i, &Rcpp::Rcout);
+    return stan::model::log_prob_propto<true>(*ptr.get(), upars, &Rcpp::Rcout);
   } else {
-    return stan::model::log_prob_propto<false>(*ptr.get(), upars, params_i, &Rcpp::Rcout);
+    return stan::model::log_prob_propto<false>(*ptr.get(), upars, &Rcpp::Rcout);
   }
 }
 
 // [[Rcpp::export]]
-Rcpp::NumericVector grad_log_prob(SEXP ext_model_ptr, std::vector<double> upars,
+Rcpp::NumericVector grad_log_prob(SEXP ext_model_ptr, Eigen::VectorXd upars,
                                   bool jac_adjust) {
   Rcpp::XPtr<stan::model::model_base> ptr(ext_model_ptr);
-  std::vector<double> gradients;
-  std::vector<int> params_i;
+  Eigen::VectorXd gradients;
 
   double lp;
   if (jac_adjust) {
-    lp = stan::model::log_prob_grad<true, true>(
-      *ptr.get(), upars, params_i, gradients);
+    lp = stan::model::log_prob_grad<true, true>(*ptr.get(), upars, gradients);
   } else {
-  lp = stan::model::log_prob_grad<true, false>(
-      *ptr.get(), upars, params_i, gradients);
+    lp = stan::model::log_prob_grad<true, false>(*ptr.get(), upars, gradients);
   }
-  Rcpp::NumericVector grad_rtn = Rcpp::wrap(gradients);
+  Rcpp::NumericVector grad_rtn(Rcpp::wrap(std::move(gradients)));
   grad_rtn.attr("log_prob") = lp;
   return grad_rtn;
 }
 
+// [[Rcpp::export]]
+Rcpp::List hessian(SEXP ext_model_ptr, Eigen::VectorXd upars, bool jacobian) {
+  Rcpp::XPtr<stan::model::model_base> ptr(ext_model_ptr);
+
+  auto hessian_functor = [&](auto&& x) {
+    if (jacobian) {
+      return ptr->log_prob<true, true>(x, 0);
+    } else {
+      return ptr->log_prob<true, false>(x, 0);
+    }
+  };
+
+  double log_prob;
+  Eigen::VectorXd grad;
+  Eigen::MatrixXd hessian;
+
+  stan::math::internal::finite_diff_hessian_auto(hessian_functor, upars, log_prob, grad, hessian);
+
+  return Rcpp::List::create(
+    Rcpp::Named("log_prob") = log_prob,
+    Rcpp::Named("grad_log_prob") = grad,
+    Rcpp::Named("hessian") = hessian);
+}
+
 // [[Rcpp::export]]
 size_t get_num_upars(SEXP ext_model_ptr) {
   Rcpp::XPtr<stan::model::model_base> ptr(ext_model_ptr);
@@ -95,12 +119,23 @@ Rcpp::List get_param_metadata(SEXP ext_model_ptr) {
 }
 
 // [[Rcpp::export]]
-std::vector<double> unconstrain_variables(SEXP ext_model_ptr, std::string init_path) {
+Eigen::VectorXd unconstrain_variables(SEXP ext_model_ptr, Eigen::VectorXd variables) {
   Rcpp::XPtr<stan::model::model_base> ptr(ext_model_ptr);
-  std::vector<int> params_i;
-  std::vector<double> vars;
-  ptr->transform_inits(*var_context(init_path), params_i, vars, &Rcpp::Rcout);
-  return vars;
+  Eigen::VectorXd unconstrained_variables;
+  ptr->unconstrain_array(variables, unconstrained_variables, &Rcpp::Rcout);
+  return unconstrained_variables;
+}
+
+// [[Rcpp::export]]
+Eigen::MatrixXd unconstrain_draws(SEXP ext_model_ptr, Eigen::MatrixXd variables) {
+  Rcpp::XPtr<stan::model::model_base> ptr(ext_model_ptr);
+  Eigen::MatrixXd unconstrained_draws(variables.cols(), variables.rows());
+  for (int i = 0; i < variables.rows(); i++) {
+    Eigen::VectorXd unconstrained_variables;
+    ptr->unconstrain_array(variables.transpose().col(i), unconstrained_variables, &Rcpp::Rcout);
+    unconstrained_draws.col(i) = unconstrained_variables;
+  }
+  return unconstrained_draws.transpose();
 }
 
 // [[Rcpp::export]]

inst/include/model_methods.cpp

@@ -34,17 +34,10 @@ test_that("Methods error if not compiled", {
   )
 })
 
-test_that("User warned about higher-order autodiff with hessian", {
-  skip_if(os_is_wsl())
-  expect_message(
-    fit$init_model_methods(hessian = TRUE, verbose = TRUE),
-    "The hessian method relies on higher-order autodiff which is still experimental. Please report any compilation errors that you encounter",
-    fixed = TRUE
-    )
-})
 
 test_that("Methods return correct values", {
   skip_if(os_is_wsl())
+  fit$init_model_methods(verbose = TRUE)
   lp <- fit$log_prob(unconstrained_variables=c(0.1))
   expect_equal(lp, -8.6327599208828509347)