Feature/reduce mem pathfinder

src/stan/callbacks/multi_writer.hpp

@@ -0,0 +1,60 @@
+#ifndef STAN_CALLBACKS_MULTI_WRITER_HPP
+#define STAN_CALLBACKS_MULTI_WRITER_HPP
+
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/callbacks/writer.hpp>
+#include <stan/math/prim/functor/for_each.hpp>
+#include <memory>
+#include <ostream>
+#include <string>
+#include <vector>
+
+namespace stan {
+namespace callbacks {
+
+/**
+ * `multi_writer` is an layer on top of a writer class that
+ *  allows for multiple output streams to be written to.
+ * @tparam Writers A parameter pack of types that inherit from `writer`
+ */
+template <typename... Writers>
+class multi_writer {
+ public:
+  /**
+   * Constructs a multi stream writer from a parameter pack of writers.
+   *
+   * @param[in, out] args A parameter pack of writers
+   */
+  template <typename... Args>
+  explicit multi_writer(Args&&... args)
+      : output_(std::forward<Args>(args)...) {}
+
+  multi_writer();
+  /**
+   * @tparam T Any type accepted by a `writer` overload
+   * @param[in] x A value to write to the output streams
+   */
+  template <typename T>
+  void operator()(T&& x) {
+    stan::math::for_each([&](auto&& output) { output(x); }, output_);
+  }
+  void operator()() {
+    stan::math::for_each([](auto&& output) { output(); }, output_);
+  }
+
+  /**
+   * Get the underlying stream
+   */
+  inline auto& get_stream() noexcept { return output_; }
+
+ private:
+  /**
+   * Output stream
+   */
+  std::tuple<std::reference_wrapper<Writers>...> output_;
+};
+
+}  // namespace callbacks
+}  // namespace stan
+
+#endif

src/stan/callbacks/unique_stream_writer.hpp

@@ -141,6 +141,7 @@ class unique_stream_writer final : public writer {
    */
   template <class T>
   void write_vector(const std::vector<T>& v) {
+    std::stringstream ss;
     if (output_ == nullptr)
       return;
     if (v.empty()) {
@@ -149,9 +150,10 @@ class unique_stream_writer final : public writer {
     auto last = v.end();
     --last;
     for (auto it = v.begin(); it != last; ++it) {
-      *output_ << *it << ",";
+      ss << *it << ",";
     }
-    *output_ << v.back() << std::endl;
+    ss << v.back() << std::endl;
+    *output_ << ss.str();
   }
 };
 

src/stan/services/pathfinder/multi.hpp

@@ -4,6 +4,7 @@
 #include <stan/callbacks/interrupt.hpp>
 #include <stan/callbacks/logger.hpp>
 #include <stan/callbacks/writer.hpp>
+#include <stan/callbacks/multi_writer.hpp>
 #include <stan/io/var_context.hpp>
 #include <stan/optimization/bfgs.hpp>
 #include <stan/optimization/lbfgs_update.hpp>
@@ -14,6 +15,7 @@
 #include <stan/services/util/duration_diff.hpp>
 #include <stan/services/util/initialize.hpp>
 #include <tbb/parallel_for.h>
+#include <tbb/concurrent_vector.h>
 #include <boost/random/discrete_distribution.hpp>
 #include <string>
 #include <vector>
@@ -22,6 +24,26 @@ namespace stan {
 namespace services {
 namespace pathfinder {
 
+namespace internal {
+
+template <typename Writer>
+struct concurrent_writer {
+  std::reference_wrapper<Writer> writer;
+  std::reference_wrapper<std::mutex> mut_;
+  explicit concurrent_writer(std::mutex& mut, Writer& writer)
+      : writer(writer), mut_(mut) {}
+  template <typename T>
+  void operator()(T&& t) {
+    std::lock_guard<std::mutex> lock(mut_.get());
+    writer.get()(t);
+  }
+  void operator()() {
+    std::lock_guard<std::mutex> lock(mut_.get());
+    writer.get()();
+  }
+};
+}  // namespace internal
+
 /**
  * Runs multiple pathfinders with final approximate samples drawn using PSIS.
  *
@@ -108,114 +130,125 @@ inline int pathfinder_lbfgs_multi(
   std::vector<std::string> param_names;
   param_names.push_back("lp_approx__");
   param_names.push_back("lp__");
+  param_names.push_back("pathfinder__");
   model.constrained_param_names(param_names, true, true);
-
   parameter_writer(param_names);
-  std::vector<Eigen::Array<double, Eigen::Dynamic, 1>> individual_lp_ratios;
-  individual_lp_ratios.resize(num_paths);
-  std::vector<Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic>>
-      individual_samples;
-  individual_samples.resize(num_paths);
   std::atomic<size_t> lp_calls{0};
+  tbb::concurrent_vector<internal::elbo_est_t> elbo_estimates;
+  // This should never allocate after this
+  elbo_estimates.reserve(num_paths + 10);
+  auto constrain_fun = [](auto&& constrained_draws, auto&& unconstrained_draws,
+                          auto&& model, auto&& rng) {
+    model.write_array(rng, unconstrained_draws, constrained_draws);
+    return constrained_draws;
+  };
   try {
+    // Idea: Instead of a mutex to lock the writes use a queue and busy thread
+    std::mutex write_mutex;
     tbb::parallel_for(
         tbb::blocked_range<int>(0, num_paths), [&](tbb::blocked_range<int> r) {
           for (int iter = r.begin(); iter < r.end(); ++iter) {
-            auto pathfinder_ret
-                = stan::services::pathfinder::pathfinder_lbfgs_single<true>(
-                    model, *(init[iter]), random_seed, stride_id + iter,
-                    init_radius, history_size, init_alpha, tol_obj, tol_rel_obj,
-                    tol_grad, tol_rel_grad, tol_param, num_iterations,
-                    num_elbo_draws, num_draws, save_iterations, refresh,
-                    interrupt, logger, init_writers[iter],
-                    single_path_parameter_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.");
-              return;
+            if (psis_resample && calculate_lp) {
+              auto pathfinder_ret
+                  = stan::services::pathfinder::pathfinder_lbfgs_single<true>(
+                      model, *(init[iter]), random_seed, stride_id + iter,
+                      init_radius, history_size, init_alpha, tol_obj,
+                      tol_rel_obj, tol_grad, tol_rel_grad, tol_param,
+                      num_iterations, num_elbo_draws, num_draws,
+                      save_iterations, refresh, interrupt, logger,
+                      init_writers[iter], single_path_parameter_writer[iter],
+                      single_path_diagnostic_writer[iter], calculate_lp,
+                      psis_resample);
+              if (unlikely(std::get<0>(pathfinder_ret) != error_codes::OK)) {
+                logger.error(std::string("Pathfinder iteration: ")
+                             + std::to_string(iter) + " failed.");
+                return;
+              }
+              lp_calls += std::get<2>(pathfinder_ret);
+              elbo_estimates.push_back(std::move(std::get<1>(pathfinder_ret)));
+            } else {
+              // For no psis, have single write to both single and multi writers
+              using multi_writer_t = stan::callbacks::multi_writer<
+                  SingleParamWriter, internal::concurrent_writer<ParamWriter>>;
+              internal::concurrent_writer safe_write{write_mutex,
+                                                     parameter_writer};
+              multi_writer_t multi_param_writer(
+                  single_path_parameter_writer[iter], safe_write);
+              auto pathfinder_ret
+                  = stan::services::pathfinder::pathfinder_lbfgs_single<true>(
+                      model, *(init[iter]), random_seed, stride_id + iter,
+                      init_radius, history_size, init_alpha, tol_obj,
+                      tol_rel_obj, tol_grad, tol_rel_grad, tol_param,
+                      num_iterations, num_elbo_draws, num_draws,
+                      save_iterations, refresh, interrupt, logger,
+                      init_writers[iter], multi_param_writer,
+                      single_path_diagnostic_writer[iter], calculate_lp,
+                      psis_resample);
+              if (unlikely(std::get<0>(pathfinder_ret) != error_codes::OK)) {
+                logger.error(std::string("Pathfinder iteration: ")
+                             + std::to_string(iter) + " failed.");
+                return;
+              }
+              lp_calls += std::get<2>(pathfinder_ret);
             }
-            individual_lp_ratios[iter] = std::move(std::get<1>(pathfinder_ret));
-            individual_samples[iter] = std::move(std::get<2>(pathfinder_ret));
-            lp_calls += std::get<3>(pathfinder_ret);
           }
         });
   } catch (const std::exception& e) {
     logger.error(e.what());
     return error_codes::SOFTWARE;
   }
-
-  // if any pathfinders failed, we want to remove their empty results
-  individual_lp_ratios.erase(
-      std::remove_if(individual_lp_ratios.begin(), individual_lp_ratios.end(),
-                     [](const auto& v) { return v.size() == 0; }),
-      individual_lp_ratios.end());
-  individual_samples.erase(
-      std::remove_if(individual_samples.begin(), individual_samples.end(),
-                     [](const auto& v) { return v.size() == 0; }),
-      individual_samples.end());
-
   const auto end_pathfinders_time = std::chrono::steady_clock::now();
-
-  const double pathfinders_delta_time = stan::services::util::duration_diff(
+  auto pathfinders_delta_time = stan::services::util::duration_diff(
       start_pathfinders_time, end_pathfinders_time);
-  const auto start_psis_time = std::chrono::steady_clock::now();
-  const size_t successful_pathfinders = individual_samples.size();
-  if (successful_pathfinders == 0) {
-    logger.info("No pathfinders ran successfully");
-    return error_codes::SOFTWARE;
-  }
-  if (refresh != 0) {
-    logger.info("Total log probability function evaluations:"
-                + std::to_string(lp_calls));
-  }
-  size_t num_returned_samples = 0;
-  // Because of failure in single pathfinder we can have multiple returned sizes
-  for (auto&& ilpr : individual_lp_ratios) {
-    num_returned_samples += ilpr.size();
-  }
-  // 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_samples[i].cols();
-    samples.middleCols(filling_start_row, individ_num_samples)
-        = individual_samples[i].matrix();
-    filling_start_row += individ_num_samples;
-  }
   double psis_delta_time = 0;
+  stan::rng_t rng = util::create_rng(random_seed, stride_id);
   if (psis_resample && calculate_lp) {
+    const auto start_psis_time = std::chrono::steady_clock::now();
+    const auto num_successful_paths = elbo_estimates.size();
+    const Eigen::Index num_returned_samples = num_draws * num_successful_paths;
     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;
+    Eigen::Index filling_start_row = 0;
+    for (const auto& elbo_est : elbo_estimates) {
+      const Eigen::Index individ_num_lp = elbo_est.lp_ratio.size();
+      lp_ratios.segment(filling_start_row, individ_num_lp) = elbo_est.lp_ratio;
+      filling_start_row += individ_num_lp;
     }
-
     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);
-    stan::rng_t rng = util::create_rng(random_seed, stride_id);
     boost::variate_generator<stan::rng_t&, 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())));
+    Eigen::VectorXd unconstrained_col;
+    Eigen::VectorXd approx_samples_constrained_col;
+    const auto uc_param_size = param_names.size() - 3;
+    Eigen::Matrix<double, 1, Eigen::Dynamic> sample_row(param_names.size());
     for (size_t i = 0; i <= num_multi_draws - 1; ++i) {
-      parameter_writer(samples.col(rand_psis_idx()));
+      auto draw_idx = rand_psis_idx();
+      // Calculate which pathfinder the draw came from
+      Eigen::Index path_num = std::floor(draw_idx / num_draws);
+      auto path_sample_idx = draw_idx % num_draws;
+      auto&& elbo_est = elbo_estimates[path_num];
+      auto&& lp_draws = elbo_est.lp_mat;
+      auto&& new_draws = elbo_est.repeat_draws;
+      const Eigen::Index param_size = new_draws.rows();
+      const auto num_samples = new_draws.cols();
+      unconstrained_col = new_draws.col(path_sample_idx);
+      constrain_fun(approx_samples_constrained_col, unconstrained_col, model,
+                    rng);
+      sample_row.head(2) = lp_draws.row(path_sample_idx).matrix();
+      sample_row(2) = path_num;
+      sample_row.tail(uc_param_size) = approx_samples_constrained_col;
+      parameter_writer(sample_row);
     }
     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);
   }
   parameter_writer();
   const auto time_header = std::string("Elapsed Time: ");
@@ -224,17 +257,15 @@ inline int pathfinder_lbfgs_multi(
         + std::string(" seconds")
         + ((psis_resample && calculate_lp) ? " (Pathfinders)" : " (Total)");
   parameter_writer(optim_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);
-  }
+  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 error_codes::OK;
 }