Adds option to turn off psis resampling and lp calculation for pathfinder

src/stan/services/pathfinder/multi.hpp

@@ -77,6 +77,14 @@ namespace pathfinder {
  * @param[in,out] parameter_writer output for parameter values
  * @param[in,out] diagnostic_writer output for diagnostics values,
  * `error_codes::SOFTWARE` for failures
+ * @param[in] return_lp Whether single pathfinder should return lp calculations.
+ *  If `true`, calculates the joint log probability for each sample.
+ *  If `false`, (`num_draws` - `num_elbo_draws`) of the joint log probability
+ *  calculations will be `NA` and psis resampling will not be performed.
+ * @param[in] psis_resampling If `true`, psis resampling is performed over the
+ *  samples returned by all of the individual pathfinders and `num_multi_draws`
+ *  samples are returned. If `false`, no psis resampling is performed
+ *  and (`num_paths` * `num_draws`) samples are returned.
  * @return error_codes::OK if successful
  */
 template <class Model, typename InitContext, typename InitWriter,
@@ -92,7 +100,8 @@ inline int pathfinder_lbfgs_multi(
     callbacks::logger& logger, InitWriter&& init_writers,
     std::vector<SingleParamWriter>& single_path_parameter_writer,
     std::vector<SingleDiagnosticWriter>& single_path_diagnostic_writer,
-    ParamWriter& parameter_writer, DiagnosticWriter& diagnostic_writer) {
+    ParamWriter& parameter_writer, DiagnosticWriter& diagnostic_writer,
+    bool return_lp = true, bool psis_resample = true) {
   const auto start_pathfinders_time = std::chrono::steady_clock::now();
   std::vector<std::string> param_names;
   param_names.push_back("lp_approx__");
@@ -117,7 +126,7 @@ inline int pathfinder_lbfgs_multi(
                   num_elbo_draws, num_draws, save_iterations, refresh,
                   interrupt, logger, init_writers[iter],
                   single_path_parameter_writer[iter],
-                  single_path_diagnostic_writer[iter]);
+                  single_path_diagnostic_writer[iter], return_lp);
           if (unlikely(std::get<0>(pathfinder_ret) != error_codes::OK)) {
             logger.error(std::string("Pathfinder iteration: ")
                          + std::to_string(iter) + " failed.");
@@ -158,53 +167,68 @@ inline int pathfinder_lbfgs_multi(
   for (auto&& ilpr : individual_lp_ratios) {
     num_returned_samples += ilpr.size();
   }
-  Eigen::Array<double, Eigen::Dynamic, 1> lp_ratios(num_returned_samples);
+  // Rows are individual parameters and columns are samples per iteration
   Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic> samples(
       individual_samples[0].rows(), num_returned_samples);
   Eigen::Index filling_start_row = 0;
   for (size_t i = 0; i < successful_pathfinders; ++i) {
-    const Eigen::Index individ_num_samples = individual_lp_ratios[i].size();
-    lp_ratios.segment(filling_start_row, individ_num_samples)
-        = individual_lp_ratios[i];
+    const Eigen::Index individ_num_samples = individual_samples[i].cols();
     samples.middleCols(filling_start_row, individ_num_samples)
         = individual_samples[i];
     filling_start_row += individ_num_samples;
   }
-  const auto tail_len = std::min(0.2 * num_returned_samples,
-                                 3 * std::sqrt(num_returned_samples));
-  Eigen::Array<double, Eigen::Dynamic, 1> weight_vals
-      = stan::services::psis::psis_weights(lp_ratios, tail_len, logger);
-  boost::ecuyer1988 rng
-      = util::create_rng<boost::ecuyer1988>(random_seed, stride_id);
-  boost::variate_generator<
-      boost::ecuyer1988&,
-      boost::random::discrete_distribution<Eigen::Index, double>>
-      rand_psis_idx(rng,
-                    boost::random::discrete_distribution<Eigen::Index, double>(
-                        boost::iterator_range<double*>(
-                            weight_vals.data(),
-                            weight_vals.data() + weight_vals.size())));
-  for (size_t i = 0; i <= num_multi_draws - 1; ++i) {
-    parameter_writer(samples.col(rand_psis_idx()));
+  double psis_delta_time = 0;
+  if (psis_resample && return_lp) {
+    Eigen::Array<double, Eigen::Dynamic, 1> lp_ratios(num_returned_samples);
+    filling_start_row = 0;
+    for (size_t i = 0; i < successful_pathfinders; ++i) {
+      const Eigen::Index individ_num_samples = individual_lp_ratios[i].size();
+      lp_ratios.segment(filling_start_row, individ_num_samples)
+          = individual_lp_ratios[i];
+      filling_start_row += individ_num_samples;
+    }
+
+    const auto tail_len = std::min(0.2 * num_returned_samples,
+                                   3 * std::sqrt(num_returned_samples));
+    Eigen::Array<double, Eigen::Dynamic, 1> weight_vals
+        = stan::services::psis::psis_weights(lp_ratios, tail_len, logger);
+    boost::ecuyer1988 rng
+        = util::create_rng<boost::ecuyer1988>(random_seed, stride_id);
+    boost::variate_generator<
+        boost::ecuyer1988&,
+        boost::random::discrete_distribution<Eigen::Index, double>>
+        rand_psis_idx(
+            rng, boost::random::discrete_distribution<Eigen::Index, double>(
+                     boost::iterator_range<double*>(
+                         weight_vals.data(),
+                         weight_vals.data() + weight_vals.size())));
+    for (size_t i = 0; i <= num_multi_draws - 1; ++i) {
+      parameter_writer(samples.col(rand_psis_idx()));
+    }
+    const auto end_psis_time = std::chrono::steady_clock::now();
+    psis_delta_time
+        = stan::services::util::duration_diff(start_psis_time, end_psis_time);
+
+  } else {
+    parameter_writer(samples);
   }
-  const auto end_psis_time = std::chrono::steady_clock::now();
-  double psis_delta_time
-      = stan::services::util::duration_diff(start_psis_time, end_psis_time);
   parameter_writer();
   const auto time_header = std::string("Elapsed Time: ");
   std::string optim_time_str = time_header
                                + std::to_string(pathfinders_delta_time)
                                + " seconds (Pathfinders)";
   parameter_writer(optim_time_str);
-  std::string psis_time_str = std::string(time_header.size(), ' ')
-                              + std::to_string(psis_delta_time)
-                              + " seconds (PSIS)";
-  parameter_writer(psis_time_str);
-  std::string total_time_str
-      = std::string(time_header.size(), ' ')
-        + std::to_string(pathfinders_delta_time + psis_delta_time)
-        + " seconds (Total)";
-  parameter_writer(total_time_str);
+  if (psis_resample && return_lp) {
+    std::string psis_time_str = std::string(time_header.size(), ' ')
+                                + std::to_string(psis_delta_time)
+                                + " seconds (PSIS)";
+    parameter_writer(psis_time_str);
+    std::string total_time_str
+        = std::string(time_header.size(), ' ')
+          + std::to_string(pathfinders_delta_time + psis_delta_time)
+          + " seconds (Total)";
+    parameter_writer(total_time_str);
+  }
   parameter_writer();
   return 0;
 }

src/stan/services/pathfinder/single.hpp

@@ -217,8 +217,8 @@ inline elbo_est_t est_approx_draws(LPF&& lp_fun, ConstrainF&& constrain_fun,
                                    RNG&& rng,
                                    const taylor_approx_t& taylor_approx,
                                    size_t num_samples, const EigVec& alpha,
-                                   const std::string& iter_msg,
-                                   Logger&& logger) {
+                                   const std::string& iter_msg, Logger&& logger,
+                                   bool return_lp = true) {
   boost::variate_generator<boost::ecuyer1988&, boost::normal_distribution<>>
       rand_unit_gaus(rng, boost::normal_distribution<>());
   const auto num_params = taylor_approx.x_center.size();
@@ -241,18 +241,25 @@ inline elbo_est_t est_approx_draws(LPF&& lp_fun, ConstrainF&& constrain_fun,
         };
   Eigen::VectorXd approx_samples_col;
   std::stringstream pathfinder_ss;
-  for (Eigen::Index i = 0; i < num_samples; ++i) {
-    try {
-      approx_samples_col = approx_samples.col(i);
-      ++lp_fun_calls;
-      lp_mat.coeffRef(i, 1) = lp_fun(approx_samples_col, pathfinder_ss);
-    } catch (const std::exception& e) {
-      lp_mat.coeffRef(i, 1) = -std::numeric_limits<double>::infinity();
+  Eigen::Array<double, Eigen::Dynamic, 1> lp_ratio;
+  if (return_lp) {
+    for (Eigen::Index i = 0; i < num_samples; ++i) {
+      try {
+        approx_samples_col = approx_samples.col(i);
+        ++lp_fun_calls;
+        lp_mat.coeffRef(i, 1) = lp_fun(approx_samples_col, pathfinder_ss);
+      } catch (const std::exception& e) {
+        lp_mat.coeffRef(i, 1) = -std::numeric_limits<double>::infinity();
+      }
+      log_stream(logger, pathfinder_ss, iter_msg);
     }
-    log_stream(logger, pathfinder_ss, iter_msg);
+    lp_ratio = lp_mat.col(1) - lp_mat.col(0);
+  } else {
+    lp_ratio = Eigen::Array<double, Eigen::Dynamic, 1>::Constant(
+        lp_mat.rows(), std::numeric_limits<double>::quiet_NaN());
+    lp_mat.col(1) = Eigen::Matrix<double, Eigen::Dynamic, 1>::Constant(
+        lp_mat.rows(), std::numeric_limits<double>::quiet_NaN());
   }
-  Eigen::Array<double, Eigen::Dynamic, 1> lp_ratio
-      = lp_mat.col(1) - lp_mat.col(0);
   if (ReturnElbo) {
     double elbo = lp_ratio.mean();
     return elbo_est_t{elbo, lp_fun_calls, std::move(approx_samples),
@@ -575,6 +582,11 @@ auto pathfinder_impl(RNG&& rng, LPFun&& lp_fun, ConstrainFun&& constrain_fun,
  * @param[in,out] init_writer Writer callback for unconstrained inits
  * @param[in,out] parameter_writer Writer callback for parameter values
  * @param[in,out] diagnostic_writer output for diagnostics values
+ * @param[in] return_lp Whether single pathfinder should return lp calculations.
+ *  If `true`, calculates the joint log probability for each sample.
+ *  If `false`, (`num_draws` - `num_elbo_draws`) of the joint log probability
+ * calculations will be `NA` and psis resampling will not be performed. Setting
+ * this parameter to `false` will also set all of the lp ratios to `NaN`.
  * @return If `ReturnLpSamples` is `true`, returns a tuple of the error code,
  * approximate draws, and a vector of the lp ratio. If `false`, only returns an
  * error code `error_codes::OK` if successful, `error_codes::SOFTWARE`
@@ -590,7 +602,7 @@ inline auto pathfinder_lbfgs_single(
     int num_elbo_draws, int num_draws, bool save_iterations, int refresh,
     callbacks::interrupt& interrupt, callbacks::logger& logger,
     callbacks::writer& init_writer, ParamWriter& parameter_writer,
-    DiagnosticWriter& diagnostic_writer) {
+    DiagnosticWriter& diagnostic_writer, bool return_lp = true) {
   const auto start_pathfinder_time = std::chrono::steady_clock::now();
   boost::ecuyer1988 rng
       = util::create_rng<boost::ecuyer1988>(random_seed, stride_id);
@@ -854,13 +866,13 @@ inline auto pathfinder_lbfgs_single(
     try {
       internal::elbo_est_t est_draws = internal::est_approx_draws<false>(
           lp_fun, constrain_fun, rng, taylor_approx_best, remaining_draws,
-          taylor_approx_best.alpha, path_num, logger);
+          taylor_approx_best.alpha, path_num, logger, return_lp);
       num_evals += est_draws.fn_calls;
       auto&& new_lp_ratio = est_draws.lp_ratio;
       auto&& lp_draws = est_draws.lp_mat;
       auto&& new_draws = est_draws.repeat_draws;
-      lp_ratio = Eigen::Array<double, Eigen::Dynamic, 1>(
-          new_lp_ratio.size() + elbo_lp_ratio.size());
+      lp_ratio = Eigen::Array<double, Eigen::Dynamic, 1>(elbo_lp_ratio.size()
+                                                         + new_lp_ratio.size());
       lp_ratio.head(elbo_lp_ratio.size()) = elbo_lp_ratio.array();
       lp_ratio.tail(new_lp_ratio.size()) = new_lp_ratio.array();
       const auto total_size = elbo_draws.cols() + new_draws.cols();