Feature/3146 laplace approx

src/stan/services/optimize/laplace_sample.hpp

@@ -0,0 +1,183 @@
+#ifndef STAN_SERVICES_OPTIMIZE_LAPLACE_SAMPLE_HPP
+#define STAN_SERVICES_OPTIMIZE_LAPLACE_SAMPLE_HPP
+
+#include <stan/callbacks/interrupt.hpp>
+#include <stan/callbacks/logger.hpp>
+#include <stan/callbacks/writer.hpp>
+#include <stan/math/rev.hpp>
+#include <stan/services/error_codes.hpp>
+#include <stan/services/util/create_rng.hpp>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+namespace stan {
+namespace services {
+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) {
+  if (draws <= 0) {
+    throw std::domain_error("Number of draws must be > 0; found draws = "
+                            + std::to_string(draws));
+  }
+
+  std::vector<std::string> unc_param_names;
+  model.unconstrained_param_names(unc_param_names, false, false);
+  int num_unc_params = unc_param_names.size();
+
+  if (theta_hat.size() != num_unc_params) {
+    throw ::std::domain_error(
+        "Specified mode is wrong size; expected "
+        + std::to_string(num_unc_params)
+        + " unconstrained parameters, but specified mode has size = "
+        + std::to_string(theta_hat.size()));
+  }
+
+  std::vector<std::string> param_tp_gq_names;
+  model.constrained_param_names(param_tp_gq_names, true, true);
+  size_t draw_size = param_tp_gq_names.size();
+
+  // write names of params, tps, and gqs to sample writer
+  std::vector<std::string> names;
+  static const bool include_tp = true;
+  static const bool include_gq = true;
+  model.constrained_param_names(names, include_tp, include_gq);
+  names.push_back("log_p");
+  names.push_back("log_q");
+  sample_writer(names);
+
+  // create log density functor for vars and vals
+  std::stringstream log_density_msgs;
+  auto log_density_fun
+      = [&](const Eigen::Matrix<stan::math::var, -1, 1>& theta) {
+          return model.template log_prob<true, jacobian, stan::math::var>(
+              const_cast<Eigen::Matrix<stan::math::var, -1, 1>&>(theta),
+              &log_density_msgs);
+        };
+
+  // calculate inverse negative Hessian's Cholesky factor
+  if (refresh > 0) {
+    logger.info("Calculating Hessian\n");
+  }
+  double log_p;          // dummy
+  Eigen::VectorXd grad;  // dummy
+  Eigen::MatrixXd hessian;
+  interrupt();
+  math::internal::finite_diff_hessian_auto(log_density_fun, theta_hat, log_p,
+                                           grad, hessian);
+  if (refresh > 0) {
+    logger.info(log_density_msgs);
+    logger.info("\n");
+  }
+
+  // calculate Cholesky factor and inverse
+  interrupt();
+  if (refresh > 0) {
+    logger.info("Calculating inverse of Cholesky factor");
+    logger.info("\n");
+  }
+  Eigen::MatrixXd L_neg_hessian = (-hessian).llt().matrixL();
+  interrupt();
+  Eigen::MatrixXd inv_sqrt_neg_hessian = L_neg_hessian.inverse().transpose();
+  interrupt();
+  Eigen::MatrixXd half_hessian = 0.5 * hessian;
+
+  // generate draws and output to sample writer
+  if (refresh > 0) {
+    logger.info("Generating draws");
+    logger.info("\n");
+  }
+  boost::ecuyer1988 rng = util::create_rng(random_seed, 0);
+  Eigen::VectorXd draw_vec;  // declare draw_vec, msgs here to avoid re-alloc
+  for (int m = 0; m < draws; ++m) {
+    interrupt();  // allow interpution each iteration
+    if (refresh > 0 && m % refresh == 0) {
+      logger.info("iteration: ");
+      logger.info(std::to_string(m));
+      logger.info("\n");
+    }
+    Eigen::VectorXd z(num_unc_params);
+    for (int n = 0; n < num_unc_params; ++n) {
+      z(n) = math::std_normal_rng(rng);
+    }
+
+    Eigen::VectorXd unc_draw = theta_hat + inv_sqrt_neg_hessian * z;
+    double log_p = log_density_fun(unc_draw).val();
+    Eigen::VectorXd diff = unc_draw - theta_hat;
+    double log_q = diff.transpose() * half_hessian * diff;
+
+    std::stringstream msgs;
+    model.write_array(rng, unc_draw, draw_vec, include_tp, include_gq, &msgs);
+    if (refresh > 0) {
+      logger.info(msgs);
+      logger.info("\n");
+    }
+    std::vector<double> draw(&draw_vec(0), &draw_vec(0) + draw_size);
+    draw.push_back(log_p);
+    draw.push_back(log_q);
+    sample_writer(draw);
+  }
+}
+}  // namespace internal
+
+/**
+ * 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
+ * @parma[in] m model from which to sample
+ * @parma[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) {
+  try {
+    internal::laplace_sample<jacobian>(model, theta_hat, draws, random_seed,
+                                       refresh, interrupt, logger,
+                                       sample_writer);
+    return error_codes::OK;
+  } catch (const std::exception& e) {
+    if (refresh >= 0) {
+      logger.error(e.what());
+      logger.error("\n");
+    }
+  } catch (...) {
+    if (refresh >= 0) {
+      logger.error("unknown exception during execution\n");
+    }
+  }
+  return error_codes::DATAERR;
+}
+}  // namespace services
+}  // namespace stan
+
+#endif

src/test/test-models/good/services/multi_normal.stan

@@ -0,0 +1,10 @@
+transformed data {
+  vector[2] mu = [2, 3]';
+  cov_matrix[2] Sigma = [[1, 0.8], [0.8, 1]];
+}
+parameters {
+  vector[2] y;	   
+}
+model {
+  y ~ multi_normal(mu, Sigma);
+}

src/test/unit/services/optimize/laplace_sample_test.cpp

@@ -0,0 +1,128 @@
+#include <boost/algorithm/string.hpp>
+#include <gtest/gtest.h>
+#include <stan/callbacks/stream_logger.hpp>
+#include <stan/callbacks/stream_writer.hpp>
+#include <stan/io/dump.hpp>
+#include <stan/io/empty_var_context.hpp>
+#include <stan/io/stan_csv_reader.hpp>
+#include <stan/services/error_codes.hpp>
+#include <stan/services/optimize/laplace_sample.hpp>
+#include <test/test-models/good/services/multi_normal.hpp>
+#include <test/unit/services/instrumented_callbacks.hpp>
+#include <test/unit/util.hpp>
+#include <cmath>
+#include <iostream>
+#include <vector>
+
+class ServicesLaplaceSample : public ::testing::Test {
+ public:
+  ServicesLaplaceSample() : logger(msgs, msgs, msgs, msgs, msgs) {}
+
+  void SetUp() {
+    stan::io::empty_var_context var_context;
+    model = new stan_model(var_context);  // typedef from multi_normal.hpp
+  }
+
+  void TearDown() { delete model; }
+
+  stan_model* model;
+  std::stringstream msgs;
+  stan::callbacks::stream_logger logger;
+  stan::test::unit::instrumented_interrupt interrupt;
+};
+
+TEST_F(ServicesLaplaceSample, values) {
+  Eigen::VectorXd theta_hat(2);
+  theta_hat << 2, 3;
+  int draws = 50000;  // big to enable mean, var & covar test precision
+  unsigned int seed = 1234;
+  int refresh = 1;
+  std::stringstream sample_ss;
+  stan::callbacks::stream_writer sample_writer(sample_ss, "");
+  int return_code = stan::services::laplace_sample<true>(
+      *model, theta_hat, draws, seed, refresh, interrupt, logger,
+      sample_writer);
+  EXPECT_EQ(stan::services::error_codes::OK, return_code);
+  std::string samples_str = sample_ss.str();
+  EXPECT_EQ(2, count_matches("y", samples_str));
+  EXPECT_EQ(1, count_matches("y.1", samples_str));
+  EXPECT_EQ(1, count_matches("y.2", samples_str));
+  EXPECT_EQ(1, count_matches("log_p", samples_str));
+  EXPECT_EQ(1, count_matches("log_q", samples_str));
+
+  std::stringstream out;
+  stan::io::stan_csv draws_csv
+      = stan::io::stan_csv_reader::parse(sample_ss, &out);
+
+  EXPECT_EQ(4, draws_csv.header.size());
+  EXPECT_EQ("y[1]", draws_csv.header[0]);
+  EXPECT_EQ("y[2]", draws_csv.header[1]);
+  EXPECT_EQ("log_p", draws_csv.header[2]);
+  EXPECT_EQ("log_q", draws_csv.header[3]);
+
+  Eigen::MatrixXd sample = draws_csv.samples;
+  EXPECT_EQ(4, sample.cols());
+  EXPECT_EQ(draws, sample.rows());
+  Eigen::VectorXd y1 = sample.col(0);
+  Eigen::VectorXd y2 = sample.col(1);
+  Eigen::VectorXd log_p = sample.col(2);
+  Eigen::VectorXd log_q = sample.col(2);
+
+  // because target is normal, laplace approx is exact
+  for (int m = 0; m < draws; ++m) {
+    EXPECT_FLOAT_EQ(0, log_p(m) - log_q(m));
+  }
+
+  // expect mean draws to be location params +/0 MC error
+  EXPECT_NEAR(2, stan::math::mean(y1), 0.05);
+  EXPECT_NEAR(3, stan::math::mean(y2), 0.05);
+
+  double sum1 = 0;
+  for (int m = 0; m < draws; ++m) {
+    sum1 += std::pow(y1(m) - 2, 2);
+  }
+  double var1 = sum1 / draws;
+  EXPECT_NEAR(1, var1, 0.05);
+
+  double sum2 = 0;
+  for (int m = 0; m < draws; ++m) {
+    sum2 += std::pow(y2(m) - 3, 2);
+  }
+  double var2 = sum2 / draws;
+  EXPECT_NEAR(1, var2, 0.05);
+
+  double sum12 = 0;
+  for (int m = 0; m < draws; ++m) {
+    sum12 += (y1(m) - 2) * (y2(m) - 3);
+  }
+  double cov = sum12 / draws;
+  EXPECT_NEAR(0.8, cov, 0.05);
+}
+
+TEST_F(ServicesLaplaceSample, wrongSizeModeError) {
+  Eigen::VectorXd theta_hat(3);
+  theta_hat << 2, 3, 4;
+  int draws = 10;
+  unsigned int seed = 1234;
+  int refresh = 1;
+  std::stringstream sample_ss;
+  stan::callbacks::stream_writer sample_writer(sample_ss, "");
+  int RC = stan::services::laplace_sample<true>(*model, theta_hat, draws, seed,
+                                                refresh, interrupt, logger,
+                                                sample_writer);
+  EXPECT_EQ(stan::services::error_codes::DATAERR, RC);
+}
+
+TEST_F(ServicesLaplaceSample, nonPositiveDrawsError) {
+  Eigen::VectorXd theta_hat(2);
+  theta_hat << 2, 3;
+  int draws = 0;
+  unsigned int seed = 1234;
+  int refresh = 1;
+  std::stringstream sample_ss;
+  stan::callbacks::stream_writer sample_writer(sample_ss, "");
+  int RC = stan::services::laplace_sample<true>(*model, theta_hat, draws, seed,
+                                                refresh, interrupt, logger,
+                                                sample_writer);
+  EXPECT_EQ(stan::services::error_codes::DATAERR, RC);
+}