Parallel Run Adaptive Sampler

src/stan/callbacks/file_stream_writer.hpp

@@ -0,0 +1,117 @@
+#ifndef STAN_CALLBACKS_FILE_STREAM_WRITER_HPP
+#define STAN_CALLBACKS_FILE_STREAM_WRITER_HPP
+
+#include <stan/callbacks/writer.hpp>
+#include <ostream>
+#include <vector>
+#include <string>
+
+namespace stan {
+namespace callbacks {
+
+/**
+ * <code>file_stream_writer</code> is an implementation
+ * of <code>writer</code> that writes to a file.
+ */
+class file_stream_writer final : public writer {
+ public:
+  /**
+   * Constructs a file stream writer with an output stream
+   * and an optional prefix for comments.
+   *
+   * @param[in, out] A unique pointer to a type inheriting from `std::ostream`
+   * @param[in] comment_prefix string to stream before each comment line.
+   *  Default is "".
+   */
+  explicit file_stream_writer(std::unique_ptr<std::ostream>&& output,
+                              const std::string& comment_prefix = "")
+      : output_(std::move(output)), comment_prefix_(comment_prefix) {}
+
+  file_stream_writer();
+  file_stream_writer(file_stream_writer& other) = delete;
+  file_stream_writer(file_stream_writer&& other)
+      : output_(std::move(other.output_)),
+        comment_prefix_(std::move(other.comment_prefix_)) {}
+  /**
+   * Virtual destructor
+   */
+  virtual ~file_stream_writer() {}
+
+  /**
+   * Writes a set of names on a single line in csv format followed
+   * by a newline.
+   *
+   * Note: the names are not escaped.
+   *
+   * @param[in] names Names in a std::vector
+   */
+  void operator()(const std::vector<std::string>& names) {
+    write_vector(names);
+  }
+  /**
+   * Get the underlying stream
+   */
+  auto& get_stream() { return *output_; }
+
+  /**
+   * Writes a set of values in csv format followed by a newline.
+   *
+   * Note: the precision of the output is determined by the settings
+   *  of the stream on construction.
+   *
+   * @param[in] state Values in a std::vector
+   */
+  void operator()(const std::vector<double>& state) { write_vector(state); }
+
+  /**
+   * Writes the comment_prefix to the stream followed by a newline.
+   */
+  void operator()() { *output_ << comment_prefix_ << std::endl; }
+
+  /**
+   * Writes the comment_prefix then the message followed by a newline.
+   *
+   * @param[in] message A string
+   */
+  void operator()(const std::string& message) {
+    *output_ << comment_prefix_ << message << std::endl;
+  }
+
+ private:
+  /**
+   * Output stream
+   */
+  std::unique_ptr<std::ostream> output_;
+
+  /**
+   * Comment prefix to use when printing comments: strings and blank lines
+   */
+  std::string comment_prefix_;
+
+  /**
+   * Writes a set of values in csv format followed by a newline.
+   *
+   * Note: the precision of the output is determined by the settings
+   *  of the stream on construction.
+   *
+   * @param[in] v Values in a std::vector
+   */
+  template <class T>
+  void write_vector(const std::vector<T>& v) {
+    if (v.empty())
+      return;
+
+    typename std::vector<T>::const_iterator last = v.end();
+    --last;
+
+    for (typename std::vector<T>::const_iterator it = v.begin(); it != last;
+         ++it)
+      *output_ << *it << ",";
+    *output_ << v.back() << std::endl;
+  }
+};
+
+}  // namespace callbacks
+}  // namespace stan
+
+#endif

src/stan/callbacks/stream_logger.hpp

@@ -14,7 +14,7 @@ namespace callbacks {
  * <code>logger</code> that writes messages to separate
  * std::stringstream outputs.
  */
-class stream_logger : public logger {
+class stream_logger final : public logger {
  private:
   std::ostream& debug_;
   std::ostream& info_;

src/stan/callbacks/tee_writer.hpp

@@ -16,7 +16,7 @@ namespace callbacks {
  * For any call to this writer, it will tee the call to both writers
  * provided in the constructor.
  */
-class tee_writer : public writer {
+class tee_writer final : public writer {
  public:
   /**
    * Constructor accepting two writers.

src/stan/mcmc/hmc/hamiltonians/softabs_metric.hpp

@@ -174,26 +174,26 @@ class softabs_metric : public base_hamiltonian<Model, softabs_point, BaseRNG> {
 
   // Threshold below which a power series
   // approximation of the softabs function is used
-  static double lower_softabs_thresh;
+  static constexpr double lower_softabs_thresh = 1e-4;
 
   // Threshold above which an asymptotic
   // approximation of the softabs function is used
-  static double upper_softabs_thresh;
+  static constexpr double upper_softabs_thresh = 18;
 
   // Threshold below which an exact derivative is
   // used in the Jacobian calculation instead of
   // finite differencing
-  static double jacobian_thresh;
+  static constexpr double jacobian_thresh = 1e-10;
 };
 
 template <class Model, class BaseRNG>
-double softabs_metric<Model, BaseRNG>::lower_softabs_thresh = 1e-4;
+constexpr double softabs_metric<Model, BaseRNG>::lower_softabs_thresh;
 
 template <class Model, class BaseRNG>
-double softabs_metric<Model, BaseRNG>::upper_softabs_thresh = 18;
+constexpr double softabs_metric<Model, BaseRNG>::upper_softabs_thresh;
 
 template <class Model, class BaseRNG>
-double softabs_metric<Model, BaseRNG>::jacobian_thresh = 1e-10;
+constexpr double softabs_metric<Model, BaseRNG>::jacobian_thresh;
 }  // namespace mcmc
 }  // namespace stan
 #endif

src/stan/services/util/generate_transitions.hpp

@@ -44,14 +44,17 @@ void generate_transitions(stan::mcmc::base_mcmc& sampler, int num_iterations,
                           util::mcmc_writer& mcmc_writer,
                           stan::mcmc::sample& init_s, Model& model,
                           RNG& base_rng, callbacks::interrupt& callback,
-                          callbacks::logger& logger) {
+                          callbacks::logger& logger, size_t n_chain = 0) {
   for (int m = 0; m < num_iterations; ++m) {
     callback();
 
     if (refresh > 0
         && (start + m + 1 == finish || m == 0 || (m + 1) % refresh == 0)) {
       int it_print_width = std::ceil(std::log10(static_cast<double>(finish)));
       std::stringstream message;
+      if (n_chain > 0) {
+        message << "Chain [" << (n_chain + 1) << "]";
+      }
       message << "Iteration: ";
       message << std::setw(it_print_width) << m + 1 + start << " / " << finish;
       message << " [" << std::setw(3)

src/stan/services/util/run_adaptive_sampler.hpp

@@ -5,6 +5,7 @@
 #include <stan/callbacks/writer.hpp>
 #include <stan/services/util/generate_transitions.hpp>
 #include <stan/services/util/mcmc_writer.hpp>
+#include <tbb/parallel_for.h>
 #include <chrono>
 #include <vector>
 
@@ -34,7 +35,7 @@ namespace util {
  * @param[in,out] sample_writer writer for draws
  * @param[in,out] diagnostic_writer writer for diagnostic information
  */
-template <class Sampler, class Model, class RNG>
+template <typename Sampler, typename Model, typename RNG>
 void run_adaptive_sampler(Sampler& sampler, Model& model,
                           std::vector<double>& cont_vector, int num_warmup,
                           int num_samples, int num_thin, int refresh,
@@ -87,6 +88,83 @@ void run_adaptive_sampler(Sampler& sampler, Model& model,
                           / 1000.0;
   writer.write_timing(warm_delta_t, sample_delta_t);
 }
+
+template <class Sampler, class Model, class RNG, typename SampT,
+          typename DiagnoseT>
+void run_adaptive_sampler(std::vector<Sampler>& samplers, Model& model,
+                          std::vector<std::vector<double>>& cont_vectors,
+                          int num_warmup, int num_samples, int num_thin,
+                          int refresh, bool save_warmup, std::vector<RNG>& rngs,
+                          callbacks::interrupt& interrupt,
+                          callbacks::logger& logger,
+                          std::vector<SampT>& sample_writers,
+                          std::vector<DiagnoseT>& diagnostic_writers,
+                          size_t n_chain) {
+  if (n_chain == 0) {
+    run_adaptive_sampler(samplers[0], model, cont_vectors[0], num_warmup,
+                         num_samples, num_thin, refresh, save_warmup, rngs[0],
+                         interrupt, logger, sample_writers[0],
+                         diagnostic_writers[0]);
+  }
+  tbb::parallel_for(
+      tbb::blocked_range<size_t>(0, n_chain, 1),
+      [num_warmup, num_samples, num_thin, refresh, save_warmup, &samplers,
+       &model, &rngs, &interrupt, &logger, &sample_writers, &cont_vectors,
+       &diagnostic_writers](const tbb::blocked_range<size_t>& r) {
+        for (size_t i = r.begin(); i != r.end(); ++i) {
+          auto&& sampler = samplers[i];
+          sampler.engage_adaptation();
+          Eigen::Map<Eigen::VectorXd> cont_params(cont_vectors[i].data(),
+                                                  cont_vectors[i].size());
+          try {
+            sampler.z().q = cont_params;
+            sampler.init_stepsize(logger);
+          } catch (const std::exception& e) {
+            logger.info("Exception initializing step size.");
+            logger.info(e.what());
+            return;
+          }
+          stan::mcmc::sample samp(cont_params, 0, 0);
+          auto&& sample_writer = sample_writers[i];
+          auto writer = services::util::mcmc_writer(
+              sample_writer, diagnostic_writers[i], logger);
+
+          // Headers
+          writer.write_sample_names(samp, sampler, model);
+          writer.write_diagnostic_names(samp, sampler, model);
+
+          const auto start_warm = std::chrono::steady_clock::now();
+          util::generate_transitions(sampler, num_warmup, 0,
+                                     num_warmup + num_samples, num_thin,
+                                     refresh, save_warmup, true, writer, samp,
+                                     model, rngs[i], interrupt, logger, i);
+          const auto end_warm = std::chrono::steady_clock::now();
+          auto warm_delta
+              = std::chrono::duration_cast<std::chrono::milliseconds>(
+                    end_warm - start_warm)
+                    .count()
+                / 1000.0;
+          sampler.disengage_adaptation();
+          writer.write_adapt_finish(sampler);
+          sampler.write_sampler_state(sample_writer);
+
+          const auto start_sample = std::chrono::steady_clock::now();
+          util::generate_transitions(sampler, num_samples, num_warmup,
+                                     num_warmup + num_samples, num_thin,
+                                     refresh, true, false, writer, samp, model,
+                                     rngs[i], interrupt, logger, i);
+          const auto end_sample = std::chrono::steady_clock::now();
+          auto sample_delta
+              = std::chrono::duration_cast<std::chrono::milliseconds>(
+                    end_sample - start_sample)
+                    .count()
+                / 1000.0;
+          writer.write_timing(warm_delta, sample_delta);
+        }
+      },
+      tbb::simple_partitioner());
+}
+
 }  // namespace util
 }  // namespace services
 }  // namespace stan

src/test/unit/callbacks/file_stream_writer_test.cpp

@@ -0,0 +1,50 @@
+#include <gtest/gtest.h>
+#include <boost/lexical_cast.hpp>
+#include <stan/callbacks/file_stream_writer.hpp>
+
+class StanInterfaceCallbacksStreamWriter : public ::testing::Test {
+ public:
+  StanInterfaceCallbacksStreamWriter()
+      : writer(std::make_unique<std::stringstream>(std::stringstream{})) {}
+
+  void SetUp() {
+    static_cast<std::stringstream&>(writer.get_stream()).str(std::string());
+    static_cast<std::stringstream&>(writer.get_stream()).clear();
+  }
+  void TearDown() {}
+
+  stan::callbacks::file_stream_writer writer;
+};
+
+TEST_F(StanInterfaceCallbacksStreamWriter, double_vector) {
+  const int N = 5;
+  std::vector<double> x;
+  for (int n = 0; n < N; ++n)
+    x.push_back(n);
+
+  EXPECT_NO_THROW(writer(x));
+  EXPECT_EQ("0,1,2,3,4\n",
+            static_cast<std::stringstream&>(writer.get_stream()).str());
+}
+
+TEST_F(StanInterfaceCallbacksStreamWriter, string_vector) {
+  const int N = 5;
+  std::vector<std::string> x;
+  for (int n = 0; n < N; ++n)
+    x.push_back(boost::lexical_cast<std::string>(n));
+
+  EXPECT_NO_THROW(writer(x));
+  EXPECT_EQ("0,1,2,3,4\n",
+            static_cast<std::stringstream&>(writer.get_stream()).str());
+}
+
+TEST_F(StanInterfaceCallbacksStreamWriter, null) {
+  EXPECT_NO_THROW(writer());
+  EXPECT_EQ("\n", static_cast<std::stringstream&>(writer.get_stream()).str());
+}
+
+TEST_F(StanInterfaceCallbacksStreamWriter, string) {
+  EXPECT_NO_THROW(writer("message"));
+  EXPECT_EQ("message\n",
+            static_cast<std::stringstream&>(writer.get_stream()).str());
+}