Merge pull request #3249 from stan-dev/feature/no-psis-or-lp-calc-pat…

…hfinder

src/stan/callbacks/unique_stream_writer.hpp

@@ -88,11 +88,9 @@ class unique_stream_writer final : public writer {
    * parameters in the rows and samples in the columns. The matrix is then
    * transposed for the output.
    */
-  void operator()(const Eigen::MatrixXd& values) {
+  void operator()(const Eigen::Ref<Eigen::Matrix<double, -1, -1>>& values) {
     if (output_ == nullptr)
       return;
-    Eigen::IOFormat CommaInitFmt(Eigen::StreamPrecision, Eigen::DontAlignCols,
-                                 ", ", "", "", "\n", "", "");
     *output_ << values.transpose().format(CommaInitFmt);
   }
 
@@ -117,6 +115,12 @@ class unique_stream_writer final : public writer {
   }
 
  private:
+  /**
+   * Comma formatter for writing Eigen matrices
+   */
+  Eigen::IOFormat CommaInitFmt{
+      Eigen::StreamPrecision, Eigen::DontAlignCols, ", ", "", "", "\n", "", ""};
+
   /**
    * Output stream
    */

src/stan/callbacks/writer.hpp

@@ -56,7 +56,8 @@ class writer {
    * parameters in the rows and samples in the columns. The matrix is then
    * transposed for the output.
    */
-  virtual void operator()(const Eigen::MatrixXd& values) {}
+  virtual void operator()(
+      const Eigen::Ref<Eigen::Matrix<double, -1, -1>>& values) {}
 };
 
 }  // namespace callbacks

src/stan/services/pathfinder/multi.hpp

@@ -77,6 +77,16 @@ 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] calculate_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 written to `parameter_writer`. If `false`, no psis resampling is
+ * performed and (`num_paths` * `num_draws`) samples are written to
+ * `parameter_writer`.
  * @return error_codes::OK if successful
  */
 template <class Model, typename InitContext, typename InitWriter,
@@ -92,7 +102,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 calculate_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 +128,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], calculate_lp);
           if (unlikely(std::get<0>(pathfinder_ret) != error_codes::OK)) {
             logger.error(std::string("Pathfinder iteration: ")
                          + std::to_string(iter) + " failed.");
@@ -158,53 +169,69 @@ 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);
-  Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic> samples(
+  // Rows are individual parameters and columns are samples per iteration
+  Eigen::Matrix<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];
+        = individual_samples[i].matrix();
     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 && calculate_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)";
+  std::string optim_time_str
+      = time_header + std::to_string(pathfinders_delta_time)
+        + std::string(" seconds")
+        + ((psis_resample && calculate_lp) ? " (Pathfinders)" : " (Total)");
   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 && calculate_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

@@ -208,6 +208,8 @@ generate_matrix(Generator&& variate_generator, const Eigen::Index num_params,
  * @param alpha The approximation of the diagonal hessian
  * @param iter_msg The beginning of messages that includes the iteration number
  * @param logger A callback writer for messages
+ * @param calculate_lp If true, calculate the log probability of the samples.
+ * Else set to `NaN` for each sample.
  * @return A struct with the ELBO estimate along with the samples and log
  * probability ratios.
  */
@@ -217,8 +219,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 calculate_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 +243,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 (calculate_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 +584,12 @@ 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] calculate_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 +605,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 calculate_lp = true) {
   const auto start_pathfinder_time = std::chrono::steady_clock::now();
   boost::ecuyer1988 rng
       = util::create_rng<boost::ecuyer1988>(random_seed, stride_id);
@@ -604,7 +619,7 @@ inline auto pathfinder_lbfgs_single(
     logger.error(e.what());
     return internal::ret_pathfinder<ReturnLpSamples>(
         error_codes::SOFTWARE, Eigen::Array<double, Eigen::Dynamic, 1>(0),
-        Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>(0, 0), 0);
+        Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(0, 0), 0);
   }
 
   const auto num_parameters = cont_vector.size();
@@ -821,7 +836,7 @@ inline auto pathfinder_lbfgs_single(
           + " Optimization failed to start, pathfinder cannot be run.");
       return internal::ret_pathfinder<ReturnLpSamples>(
           error_codes::SOFTWARE, Eigen::Array<double, Eigen::Dynamic, 1>(0),
-          Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>(0, 0),
+          Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(0, 0),
           std::atomic<size_t>{num_evals + lbfgs.grad_evals()});
     } else {
       logger.warn(prefix_err_msg +
@@ -835,15 +850,15 @@ inline auto pathfinder_lbfgs_single(
         "successfully");
     return internal::ret_pathfinder<ReturnLpSamples>(
         error_codes::SOFTWARE, Eigen::Array<double, Eigen::Dynamic, 1>(0),
-        Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>(0, 0), num_evals);
+        Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(0, 0), num_evals);
   } else {
     if (refresh != 0) {
       logger.info(path_num + "Best Iter: [" + std::to_string(best_iteration)
                   + "] ELBO (" + std::to_string(elbo_best.elbo) + ")"
                   + " evaluations: (" + std::to_string(num_evals) + ")");
     }
   }
-  Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic> constrained_draws_mat;
+  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> constrained_draws_mat;
   Eigen::Array<double, Eigen::Dynamic, 1> lp_ratio;
   auto&& elbo_draws = elbo_best.repeat_draws;
   auto&& elbo_lp_ratio = elbo_best.lp_ratio;
@@ -854,35 +869,37 @@ 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, calculate_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();
       constrained_draws_mat
-          = Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>(names.size(),
-                                                                 total_size);
+          = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(names.size(),
+                                                                  total_size);
       Eigen::VectorXd unconstrained_col;
       Eigen::VectorXd approx_samples_constrained_col;
       for (Eigen::Index i = 0; i < elbo_draws.cols(); ++i) {
-        constrained_draws_mat.col(i).head(2) = elbo_lp_mat.row(i);
+        constrained_draws_mat.col(i).head(2) = elbo_lp_mat.row(i).matrix();
         unconstrained_col = elbo_draws.col(i);
         constrained_draws_mat.col(i).tail(num_unconstrained_params)
             = constrain_fun(rng, unconstrained_col,
-                            approx_samples_constrained_col);
+                            approx_samples_constrained_col)
+                  .matrix();
       }
       for (Eigen::Index i = elbo_draws.cols(), j = 0; i < total_size;
            ++i, ++j) {
-        constrained_draws_mat.col(i).head(2) = lp_draws.row(j);
+        constrained_draws_mat.col(i).head(2) = lp_draws.row(j).matrix();
         unconstrained_col = new_draws.col(j);
         constrained_draws_mat.col(i).tail(num_unconstrained_params)
             = constrain_fun(rng, unconstrained_col,
-                            approx_samples_constrained_col);
+                            approx_samples_constrained_col)
+                  .matrix();
       }
     } catch (const std::exception& e) {
       std::string err_msg = e.what();
@@ -893,35 +910,37 @@ inline auto pathfinder_lbfgs_single(
           + "Returning the approximate samples used for ELBO calculation: "
           + err_msg);
       constrained_draws_mat
-          = Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>(
+          = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(
               names.size(), elbo_draws.cols());
       Eigen::VectorXd approx_samples_constrained_col;
       Eigen::VectorXd unconstrained_col;
       for (Eigen::Index i = 0; i < elbo_draws.cols(); ++i) {
-        constrained_draws_mat.col(i).head(2) = elbo_lp_mat.row(i);
+        constrained_draws_mat.col(i).head(2) = elbo_lp_mat.row(i).matrix();
         unconstrained_col = elbo_draws.col(i);
         constrained_draws_mat.col(i).tail(num_unconstrained_params)
             = constrain_fun(rng, unconstrained_col,
-                            approx_samples_constrained_col);
+                            approx_samples_constrained_col)
+                  .matrix();
       }
       lp_ratio = std::move(elbo_best.lp_ratio);
     }
   } else {
     constrained_draws_mat
-        = Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>(
+        = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(
             names.size(), elbo_draws.cols());
     Eigen::VectorXd approx_samples_constrained_col;
     Eigen::VectorXd unconstrained_col;
     for (Eigen::Index i = 0; i < elbo_draws.cols(); ++i) {
-      constrained_draws_mat.col(i).head(2) = elbo_lp_mat.row(i);
+      constrained_draws_mat.col(i).head(2) = elbo_lp_mat.row(i).matrix();
       unconstrained_col = elbo_draws.col(i);
       constrained_draws_mat.col(i).tail(num_unconstrained_params)
           = constrain_fun(rng, unconstrained_col,
-                          approx_samples_constrained_col);
+                          approx_samples_constrained_col)
+                .matrix();
     }
     lp_ratio = std::move(elbo_best.lp_ratio);
   }
-  parameter_writer(constrained_draws_mat.matrix());
+  parameter_writer(constrained_draws_mat);
   parameter_writer();
   const auto end_pathfinder_time = std::chrono::steady_clock::now();
   const double pathfinder_delta_time = stan::services::util::duration_diff(