From 3c932f0aef231fbfccdf7bf47cd4d31ea6162086 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Fri, 13 Dec 2024 14:35:38 -0500 Subject: [PATCH] add docs --- src/stan/callbacks/unique_stream_writer.hpp | 150 +++++++++++++++++++- src/stan/services/pathfinder/multi.hpp | 39 ++++- src/test/unit/services/pathfinder/util.hpp | 16 +++ 3 files changed, 199 insertions(+), 6 deletions(-) diff --git a/src/stan/callbacks/unique_stream_writer.hpp b/src/stan/callbacks/unique_stream_writer.hpp index 7ef68e6eb0..41648787c5 100644 --- a/src/stan/callbacks/unique_stream_writer.hpp +++ b/src/stan/callbacks/unique_stream_writer.hpp @@ -58,6 +58,20 @@ class unique_stream_writer final : public writer { return; write_vector(names); } + /** + * 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 + * @param[in, out] ss A stringstream to use as a buffer + */ + void operator()(const std::vector& names, std::stringstream& ss) { + if (output_ == nullptr) + return; + write_vector(names, ss); + } /** * Get the underlying stream @@ -77,6 +91,20 @@ class unique_stream_writer final : public writer { return; write_vector(values); } + /** + * 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] values Values in a std::vector + * @param[in, out] ss ignored + */ + void operator()(const std::vector& values, std::stringstream& ss) { + if (output_ == nullptr) + return; + write_vector(values); + } /** * Writes multiple rows and columns of values in csv format. @@ -88,12 +116,96 @@ 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::Ref>& values) { + void operator()(const Eigen::Matrix& values) { if (output_ == nullptr) return; *output_ << values.transpose().format(CommaInitFmt); } + /** + * Writes multiple rows and columns of values in csv format. + * + * Note: the precision of the output is determined by the settings + * of the stream on construction. + * + * @param[in] values A matrix of values. The input is expected to have + * parameters in the rows and samples in the columns. The matrix is then + * transposed for the output. + * @param[in, out] ss ignored + */ + void operator()(const Eigen::Matrix& values, std::stringstream& /* ss */) { + if (output_ == nullptr) + return; + *output_ << values.transpose().format(CommaInitFmt); + } + + /** + * Writes a row of values in csv format. + * + * Note: the precision of the output is determined by the settings + * of the stream on construction. + * + * @param[in] values A row vector of values. The input is expected to have + * samples in the columns. + * @param[in, out] ss ignored + */ + void operator()(const Eigen::Matrix& values) { + if (output_ == nullptr) + return; + *output_ << values.format(CommaInitFmt); + } + + /** + * Writes a row of values in csv format. + * + * Note: the precision of the output is determined by the settings + * of the stream on construction. + * + * @param[in] values A row vector of values. The input is expected to have + * samples in the columns. + * @param[in, out] ss ignored + */ + void operator()(const Eigen::Matrix& values, std::stringstream& /* ss */) { + if (output_ == nullptr) + return; + *output_ << values.format(CommaInitFmt); + } + + /** + * Writes a row of values in csv format. + * + * Note: the precision of the output is determined by the settings + * of the stream on construction. + * + * @param[in] values A column vector of values. The input is expected to have + * samples in the rows. The matrix is then + * transposed for the output. + * @param[in, out] ss ignored + */ + void operator()(const Eigen::Matrix& values) { + if (output_ == nullptr) + return; + *output_ << values.transpose().format(CommaInitFmt); + } + + /** + * Writes a row of values in csv format. + * + * Note: the precision of the output is determined by the settings + * of the stream on construction. + * + * @param[in] values A column vector of values. The input is expected to have + * samples in the rows. The matrix is then + * transposed for the output. + */ + void operator()(const Eigen::Matrix& values, std::stringstream& /* ss */) { + if (output_ == nullptr) + return; + *output_ << values.transpose().format(CommaInitFmt); + } + + + /** * Writes the comment_prefix to the stream followed by a newline. */ @@ -114,6 +226,18 @@ class unique_stream_writer final : public writer { *output_ << comment_prefix_ << message << std::endl; } + /** + * Writes the comment_prefix then the message followed by a newline. + * + * @param[in] message A string + @ @param[in, out] ss ignored + */ + void operator()(const std::string& message, std::stringstream& /* ss */) { + if (output_ == nullptr) + return; + *output_ << comment_prefix_ << message << std::endl; + } + private: /** * Comma formatter for writing Eigen matrices @@ -153,6 +277,30 @@ class unique_stream_writer final : public writer { } *output_ << v.back() << std::endl; } + /** + * 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 + * @param[in, out] ss A stringstream to use as a buffer + */ + template + void write_vector(const std::vector& v, std::stringstream& ss) { + if (output_ == nullptr) + return; + if (v.empty()) { + return; + } + auto last = v.end(); + --last; + for (auto it = v.begin(); it != last; ++it) { + ss << *it << ","; + } + ss << v.back() << std::endl; + *output_ << ss.str(); + } }; } // namespace callbacks diff --git a/src/stan/services/pathfinder/multi.hpp b/src/stan/services/pathfinder/multi.hpp index cd39deae42..ed48b3ac9e 100644 --- a/src/stan/services/pathfinder/multi.hpp +++ b/src/stan/services/pathfinder/multi.hpp @@ -24,11 +24,30 @@ namespace services { namespace pathfinder { namespace internal { +/** + * Generate a single draw from the pathfinder algorithm. + * @tparam DoAll If true, generate all draws. If false, generate only the + * draw at `path_sample_idx`. + * @tparam ConstrainFun A functor with a valid `operator(EigVector, EigVector, Model, RNG)` + * @tparam Model A model implementation + * @tparam ElboEst A struct with the ELBO estimate + * @tparam RNG A random number generator + * @tparam ParamWriter A writer callback for parameter values + * @param[in] constrain_fun A functor to transform parameters to the constrained + * space + * @param[in] model A model implementation + * @param[in] elbo_est A struct with the ELBO estimate + * @param[in] path_num The path number + * @param[in] path_sample_idx The index of the sample to generate + * @param[in] rng A random number generator + * @param[in] param_writer A writer callback for parameter values + * @param[in,out] stream A stream to write the output + */ template inline void gen_pathfinder_draw(ConstrainFun&& constrain_fun, Model&& model, ElboEst&& elbo_est, const Eigen::Index path_num, const Eigen::Index path_sample_idx, - RNG&& rng, ParamWriter&& param_writer) { + RNG&& rng, ParamWriter&& param_writer, std::stringstream& stream) { // FIX THESE auto&& lp_draws = elbo_est.lp_mat; auto&& new_draws = elbo_est.repeat_draws; @@ -42,14 +61,17 @@ namespace internal { Eigen::Matrix sample_row(uc_param_size + 2); sample_row.head(2) = lp_draws.row(path_sample_idx).matrix(); sample_row.tail(uc_param_size) = approx_samples_constrained_col; - param_writer(sample_row); + param_writer(sample_row, stream); if constexpr (DoAll) { for (int i = 1; i < num_samples; ++i) { unconstrained_col = new_draws.col(i); constrain_fun(approx_samples_constrained_col, unconstrained_col, model, rng); sample_row.head(2) = lp_draws.row(path_sample_idx).matrix(); sample_row.tail(uc_param_size) = approx_samples_constrained_col; - param_writer(sample_row); + param_writer(sample_row, stream); + if (stream.rdbuf()) { + stream.rdbuf()->pubseekpos(0); + } } } } @@ -176,14 +198,16 @@ inline int pathfinder_lbfgs_multi( if (psis_resample && calculate_lp) { elbo_estimates.push_back(std::move(std::get<1>(pathfinder_ret))); } else { + std::stringstream stream; stan::rng_t rng = util::create_rng(random_seed, stride_id + iter); std::lock_guard lock(write_mutex); internal::gen_pathfinder_draw(constrain_fun, model, std::get<1>(pathfinder_ret), 0, 0, - rng, parameter_writer); + rng, parameter_writer, stream); } } }); + } catch (const std::exception& e) { logger.error(e.what()); return error_codes::SOFTWARE; @@ -214,6 +238,7 @@ inline int pathfinder_lbfgs_multi( boost::iterator_range( weight_vals.data(), weight_vals.data() + weight_vals.size()))); + std::stringstream stream; for (size_t i = 0; i <= num_multi_draws - 1; ++i) { auto draw_idx = rand_psis_idx(); // Calculate which pathfinder the draw came from @@ -221,7 +246,11 @@ inline int pathfinder_lbfgs_multi( auto path_sample_idx = draw_idx % num_draws; auto&& elbo_est = elbo_estimates[path_num]; internal::gen_pathfinder_draw(constrain_fun, model, elbo_est, - path_num, path_sample_idx, rng, parameter_writer); + path_num, path_sample_idx, rng, parameter_writer, stream); + if (stream.rdbuf()) { + stream.rdbuf()->pubseekpos(0); + } + } const auto end_psis_time = std::chrono::steady_clock::now(); psis_delta_time diff --git a/src/test/unit/services/pathfinder/util.hpp b/src/test/unit/services/pathfinder/util.hpp index 980b048a4b..c01e865301 100644 --- a/src/test/unit/services/pathfinder/util.hpp +++ b/src/test/unit/services/pathfinder/util.hpp @@ -154,8 +154,10 @@ class in_memory_writer : public stan::callbacks::stream_writer { * @param[in] names Names in a std::vector */ void operator()(const std::vector& names) { names_ = names; } + void operator()(const std::vector& names, std::stringstream& /* ss */) { names_ = names; } void operator()(const std::string& times) { times_.push_back(times); } + void operator()(const std::string& times, std::stringstream& /* ss */) { times_.push_back(times); } void operator()() { times_.push_back("\n"); } /** * Writes a set of values. @@ -165,6 +167,9 @@ class in_memory_writer : public stan::callbacks::stream_writer { void operator()(const std::vector& state) { states_.push_back(state); } + void operator()(const std::vector& state, std::stringstream& /* ss */) { + states_.push_back(state); + } void operator()( const std::vector>& xx) { for (auto& x_i : xx) { @@ -174,15 +179,26 @@ class in_memory_writer : public stan::callbacks::stream_writer { void operator()(const std::tuple& xx) { optim_path_.push_back(xx); } + template * = nullptr> void operator()(const EigVec& vals) { eigen_states_.push_back(vals); } + template * = nullptr> + void operator()(const EigVec& vals, std::stringstream& /* ss */) { + eigen_states_.push_back(vals); + } + template * = nullptr> void operator()(const EigMat& vals) { values_ = vals; } + template * = nullptr> + void operator()(const EigMat& vals, std::stringstream& /* ss */) { + values_ = vals; + } }; Eigen::Matrix eight_schools_r_answer() {