Skip to content

Commit

Permalink
Merge pull request #3261 from stan-dev/feature/3215-3260-laplace-addi…
Browse files Browse the repository at this point in the history
…tions

Add option to output Hessian, disable LP calculations in Laplace
  • Loading branch information
WardBrian authored Jan 30, 2024
2 parents bb4ab32 + 32032d8 commit a49908b
Show file tree
Hide file tree
Showing 2 changed files with 128 additions and 10 deletions.
83 changes: 73 additions & 10 deletions src/stan/services/optimize/laplace_sample.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include <stan/callbacks/interrupt.hpp>
#include <stan/callbacks/logger.hpp>
#include <stan/callbacks/writer.hpp>
#include <stan/callbacks/structured_writer.hpp>
#include <stan/math/rev.hpp>
#include <stan/services/error_codes.hpp>
#include <stan/services/util/create_rng.hpp>
Expand All @@ -17,9 +18,10 @@ namespace internal {

template <bool jacobian, typename Model>
void laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
int draws, unsigned int random_seed, int refresh,
callbacks::interrupt& interrupt, callbacks::logger& logger,
callbacks::writer& sample_writer) {
int draws, bool calculate_lp, unsigned int random_seed,
int refresh, callbacks::interrupt& interrupt,
callbacks::logger& logger, callbacks::writer& sample_writer,
callbacks::structured_writer& hessian_writer) {
if (draws <= 0) {
throw std::domain_error("Number of draws must be > 0; found draws = "
+ std::to_string(draws));
Expand Down Expand Up @@ -72,6 +74,13 @@ void laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
if (refresh > 0 && log_density_msgs.peek() != std::char_traits<char>::eof())
logger.info(log_density_msgs);

interrupt();
hessian_writer.begin_record();
hessian_writer.write("lp_mode", log_p);
hessian_writer.write("gradient", grad);
hessian_writer.write("Hessian", hessian);
hessian_writer.end_record();

// calculate Cholesky factor and inverse
interrupt();
if (refresh > 0) {
Expand Down Expand Up @@ -109,7 +118,13 @@ void laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
logger.info(write_array_msgs);
// output draw, log_p, log_q
std::vector<double> draw(&draw_vec(0), &draw_vec(0) + draw_size);
double log_p = log_density_fun(unc_draw).val();

double log_p;
if (calculate_lp) {
log_p = log_density_fun(unc_draw).val();
} else {
log_p = std::numeric_limits<double>::quiet_NaN();
}
draw.insert(draw.begin(), log_p);
Eigen::VectorXd diff = unc_draw - theta_hat;
double log_q = diff.transpose() * half_hessian * diff;
Expand Down Expand Up @@ -139,6 +154,8 @@ void laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
* @param[in] theta_hat unconstrained mode at which to center the
* Laplace approximation
* @param[in] draws number of draws to generate
* @param[in] calculate_lp whether to calculate the log probability of the
* approximate draws
* @param[in] random_seed seed for generating random numbers in the
* Stan program and in sampling
* @param[in] refresh period between iterations at which updates are
Expand All @@ -148,23 +165,69 @@ void laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
* sampler and from Stan programs
* @param[in,out] sample_writer callback for writing parameter names
* and then draws
* @param[in,out] hessian_writer callback for writing the log probability,
* gradient, and Hessian at the mode for diagnostic purposes
* @return a return code, with 0 indicating success
*/
template <bool jacobian, typename Model>
int laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
int draws, unsigned int random_seed, int refresh,
callbacks::interrupt& interrupt, callbacks::logger& logger,
callbacks::writer& sample_writer) {
int draws, bool calculate_lp, unsigned int random_seed,
int refresh, callbacks::interrupt& interrupt,
callbacks::logger& logger, callbacks::writer& sample_writer,
callbacks::structured_writer& hessian_writer) {
try {
internal::laplace_sample<jacobian>(model, theta_hat, draws, random_seed,
refresh, interrupt, logger,
sample_writer);
internal::laplace_sample<jacobian>(model, theta_hat, draws, calculate_lp,
random_seed, refresh, interrupt, logger,
sample_writer, hessian_writer);
} catch (const std::exception& e) {
logger.error(e.what());
return error_codes::CONFIG;
}
return error_codes::OK;
}

/**
* Take the specified number of draws from the Laplace approximation
* for the model at the specified unconstrained mode, writing the
* draws, unnormalized log density, and unnormalized density of the
* approximation to the sample writer and writing messages to the
* logger, returning a return code of zero if successful.
*
* Interrupts are called between compute-intensive operations. To
* turn off all console messages sent to the logger, set refresh to 0.
* If an exception is thrown by the model, the return value is
* non-zero, and if refresh > 0, its message is given to the logger as
* an error.
*
* @tparam jacobian `true` to include Jacobian adjustment for
* constrained parameters
* @tparam Model a Stan model
* @param[in] model model from which to sample
* @param[in] theta_hat unconstrained mode at which to center the
* Laplace approximation
* @param[in] draws number of draws to generate
* @param[in] random_seed seed for generating random numbers in the
* Stan program and in sampling
* @param[in] refresh period between iterations at which updates are
* given, with a value of 0 turning off all messages
* @param[in] interrupt callback for interrupting sampling
* @param[in,out] logger callback for writing console messages from
* sampler and from Stan programs
* @param[in,out] sample_writer callback for writing parameter names
* and then draws
* @return a return code, with 0 indicating success
*/
template <bool jacobian, typename Model>
int laplace_sample(const Model& model, const Eigen::VectorXd& theta_hat,
int draws, unsigned int random_seed, int refresh,
callbacks::interrupt& interrupt, callbacks::logger& logger,
callbacks::writer& sample_writer) {
callbacks::structured_writer dummy_hessian_writer;
return laplace_sample<jacobian>(model, theta_hat, draws, true, random_seed,
refresh, interrupt, logger, sample_writer,
dummy_hessian_writer);
}

} // namespace services
} // namespace stan

Expand Down
55 changes: 55 additions & 0 deletions src/test/unit/services/optimize/laplace_sample_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#include <gtest/gtest.h>
#include <stan/callbacks/stream_logger.hpp>
#include <stan/callbacks/stream_writer.hpp>
#include <stan/callbacks/json_writer.hpp>
#include <stan/io/dump.hpp>
#include <stan/io/empty_var_context.hpp>
#include <stan/io/stan_csv_reader.hpp>
Expand Down Expand Up @@ -100,6 +101,60 @@ TEST_F(ServicesLaplaceSample, values) {
EXPECT_NEAR(0.8, cov, 0.05);
}

struct deleter_noop {
template <typename T>
constexpr void operator()(T* arg) const {}
};

TEST_F(ServicesLaplaceSample, hessianOutput) {
Eigen::VectorXd theta_hat(2);
theta_hat << 2, 3;
int draws = 10;
bool calculate_lp = true;
unsigned int seed = 1234;
int refresh = 100;
std::stringstream sample_ss;
stan::callbacks::stream_writer sample_writer(sample_ss, "");

std::stringstream hessian_ss;
stan::callbacks::json_writer<std::stringstream, deleter_noop> hessian_writer{
std::unique_ptr<std::stringstream, deleter_noop>(&hessian_ss)};

int return_code = stan::services::laplace_sample<true>(
*model, theta_hat, draws, calculate_lp, seed, refresh, interrupt, logger,
sample_writer, hessian_writer);
EXPECT_EQ(stan::services::error_codes::OK, return_code);

std::string hessian_str = hessian_ss.str();

ASSERT_TRUE(stan::test::is_valid_JSON(hessian_str));
EXPECT_EQ(count_matches("lp_mode", hessian_str), 1);
EXPECT_EQ(count_matches("gradient", hessian_str), 1);
EXPECT_EQ(count_matches("Hessian", hessian_str), 1);
}

TEST_F(ServicesLaplaceSample, noLP) {
Eigen::VectorXd theta_hat(2);
theta_hat << 2, 3;
unsigned int seed = 1234;
int refresh = 100;
std::stringstream sample_ss;
stan::callbacks::stream_writer sample_writer(sample_ss, "");
stan::callbacks::structured_writer dummy_hessian_writer;

int draws = 11;
bool calculate_lp = false;

int return_code = stan::services::laplace_sample<true>(
*model, theta_hat, draws, calculate_lp, seed, refresh, interrupt, logger,
sample_writer, dummy_hessian_writer);
EXPECT_EQ(stan::services::error_codes::OK, return_code);

std::string samples_str = sample_ss.str();
EXPECT_EQ(1, count_matches("log_p__", samples_str));
EXPECT_EQ(draws, count_matches("nan", samples_str));
}

TEST_F(ServicesLaplaceSample, wrongSizeModeError) {
Eigen::VectorXd theta_hat(3);
theta_hat << 2, 3, 4;
Expand Down

0 comments on commit a49908b

Please sign in to comment.