Don't do inverse metric decomposition every draw

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

@@ -19,7 +19,7 @@ class dense_e_metric : public base_hamiltonian<Model, dense_e_point, BaseRNG> {
       : base_hamiltonian<Model, dense_e_point, BaseRNG>(model) {}
 
   double T(dense_e_point& z) {
-    return 0.5 * z.p.transpose() * z.inv_e_metric_ * z.p;
+    return 0.5 * z.p.transpose() * z.get_inv_metric() * z.p;
   }
 
   double tau(dense_e_point& z) { return T(z); }
@@ -34,7 +34,7 @@ class dense_e_metric : public base_hamiltonian<Model, dense_e_point, BaseRNG> {
     return Eigen::VectorXd::Zero(this->model_.num_params_r());
   }
 
-  Eigen::VectorXd dtau_dp(dense_e_point& z) { return z.inv_e_metric_ * z.p; }
+  Eigen::VectorXd dtau_dp(dense_e_point& z) { return z.get_inv_metric() * z.p; }
 
   Eigen::VectorXd dphi_dq(dense_e_point& z, callbacks::logger& logger) {
     return z.g;
@@ -50,7 +50,10 @@ class dense_e_metric : public base_hamiltonian<Model, dense_e_point, BaseRNG> {
     for (idx_t i = 0; i < u.size(); ++i)
       u(i) = rand_dense_gaus();
 
-    z.p = z.inv_e_metric_.llt().matrixU().solve(u);
+    z.p = z.get_llt_inv_metric()
+              .transpose()
+              .triangularView<Eigen::Upper>()
+              .solve(u);
   }
 };
 

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

@@ -2,6 +2,7 @@
 #define STAN_MCMC_HMC_HAMILTONIANS_DENSE_E_POINT_HPP
 
 #include <stan/callbacks/writer.hpp>
+#include <stan/math/prim/err.hpp>
 #include <stan/mcmc/hmc/hamiltonians/ps_point.hpp>
 
 namespace stan {
@@ -11,29 +12,52 @@ namespace mcmc {
  * Euclidean manifold with dense metric
  */
 class dense_e_point : public ps_point {
- public:
+ private:
   /**
    * Inverse mass matrix.
    */
   Eigen::MatrixXd inv_e_metric_;
+  Eigen::MatrixXd inv_e_metric_llt_matrixL_;
 
+ public:
   /**
    * Construct a dense point in n-dimensional phase space
    * with identity matrix as inverse mass matrix.
    *
    * @param n number of dimensions
    */
-  explicit dense_e_point(int n) : ps_point(n), inv_e_metric_(n, n) {
+  explicit dense_e_point(int n)
+      : ps_point(n), inv_e_metric_(n, n), inv_e_metric_llt_matrixL_(n, n) {
     inv_e_metric_.setIdentity();
+    inv_e_metric_llt_matrixL_.setIdentity();
   }
 
   /**
-   * Set elements of mass matrix
+   * Set inverse metric
    *
    * @param inv_e_metric initial mass matrix
    */
-  void set_metric(const Eigen::MatrixXd& inv_e_metric) {
-    inv_e_metric_ = inv_e_metric;
+  template <typename EigMat, require_eigen_matrix_dynamic_t<EigMat>* = nullptr>
+  void set_inv_metric(EigMat&& inv_e_metric) {
+    inv_e_metric_ = std::forward<EigMat>(inv_e_metric);
+    inv_e_metric_llt_matrixL_ = inv_e_metric_.llt().matrixL();
+  }
+
+  /**
+   * Get inverse metric
+   *
+   * @return reference to the inverse metric
+   */
+  const Eigen::MatrixXd& get_inv_metric() const { return inv_e_metric_; }
+
+  /**
+   * Get the transpose of the lower Cholesky factor
+   *  of the inverse metric
+   *
+   * @return reference to transpose of Cholesky factor
+   */
+  const Eigen::MatrixXd& get_llt_inv_metric() const {
+    return inv_e_metric_llt_matrixL_;
   }
 
   /**

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

@@ -18,7 +18,7 @@ class diag_e_metric : public base_hamiltonian<Model, diag_e_point, BaseRNG> {
       : base_hamiltonian<Model, diag_e_point, BaseRNG>(model) {}
 
   double T(diag_e_point& z) {
-    return 0.5 * z.p.dot(z.inv_e_metric_.cwiseProduct(z.p));
+    return 0.5 * z.p.dot(z.get_inv_metric().cwiseProduct(z.p));
   }
 
   double tau(diag_e_point& z) { return T(z); }
@@ -34,7 +34,7 @@ class diag_e_metric : public base_hamiltonian<Model, diag_e_point, BaseRNG> {
   }
 
   Eigen::VectorXd dtau_dp(diag_e_point& z) {
-    return z.inv_e_metric_.cwiseProduct(z.p);
+    return z.get_inv_metric().cwiseProduct(z.p);
   }
 
   Eigen::VectorXd dphi_dq(diag_e_point& z, callbacks::logger& logger) {
@@ -45,8 +45,8 @@ class diag_e_metric : public base_hamiltonian<Model, diag_e_point, BaseRNG> {
     boost::variate_generator<BaseRNG&, boost::normal_distribution<> >
         rand_diag_gaus(rng, boost::normal_distribution<>());
 
-    for (int i = 0; i < z.p.size(); ++i)
-      z.p(i) = rand_diag_gaus() / sqrt(z.inv_e_metric_(i));
+    z.p = z.get_inv_metric().unaryExpr(
+        [&rand_diag_gaus](auto&& x) { return rand_diag_gaus() / sqrt(x); });
   }
 };
 

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

@@ -2,6 +2,7 @@
 #define STAN_MCMC_HMC_HAMILTONIANS_DIAG_E_POINT_HPP
 
 #include <stan/callbacks/writer.hpp>
+#include <stan/math/prim/err.hpp>
 #include <stan/mcmc/hmc/hamiltonians/ps_point.hpp>
 
 namespace stan {
@@ -11,12 +12,13 @@ namespace mcmc {
  * Euclidean manifold with diagonal metric
  */
 class diag_e_point : public ps_point {
- public:
+ private:
   /**
    * Vector of diagonal elements of inverse mass matrix.
    */
   Eigen::VectorXd inv_e_metric_;
 
+ public:
   /**
    * Construct a diag point in n-dimensional phase space
    * with vector of ones for diagonal elements of inverse mass matrix.
@@ -32,10 +34,18 @@ class diag_e_point : public ps_point {
    *
    * @param inv_e_metric initial mass matrix
    */
-  void set_metric(const Eigen::VectorXd& inv_e_metric) {
-    inv_e_metric_ = inv_e_metric;
+  template <typename EigVec, require_eigen_vector_t<EigVec>* = nullptr>
+  void set_inv_metric(EigVec&& inv_e_metric) {
+    inv_e_metric_ = std::forward<EigVec>(inv_e_metric);
   }
 
+  /**
+   * Get inverse metric
+   *
+   * @return reference to the inverse metric
+   */
+  const Eigen::VectorXd& get_inv_metric() const { return inv_e_metric_; }
+
   /**
    * Write elements of mass matrix to string and handoff to writer.
    *

src/stan/mcmc/hmc/nuts/adapt_dense_e_nuts.hpp

@@ -29,10 +29,14 @@ class adapt_dense_e_nuts : public dense_e_nuts<Model, BaseRNG>,
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->covar_adaptation_.learn_covariance(
-          this->z_.inv_e_metric_, this->z_.q);
+      Eigen::MatrixXd inv_metric;
+
+      bool update
+          = this->covar_adaptation_.learn_covariance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
 
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));

src/stan/mcmc/hmc/nuts/adapt_diag_e_nuts.hpp

@@ -29,10 +29,14 @@ class adapt_diag_e_nuts : public diag_e_nuts<Model, BaseRNG>,
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->var_adaptation_.learn_variance(this->z_.inv_e_metric_,
-                                                         this->z_.q);
+      Eigen::VectorXd inv_metric;
+
+      bool update
+          = this->var_adaptation_.learn_variance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
 
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));

src/stan/mcmc/hmc/nuts/base_nuts.hpp

@@ -57,12 +57,12 @@ class base_nuts : public base_hmc<Model, Hamiltonian, Integrator, BaseRNG> {
 
   ~base_nuts() {}
 
-  void set_metric(const Eigen::MatrixXd& inv_e_metric) {
-    this->z_.set_metric(inv_e_metric);
+  void set_inv_metric(const Eigen::MatrixXd& inv_e_metric) {
+    this->z_.set_inv_metric(inv_e_metric);
   }
 
-  void set_metric(const Eigen::VectorXd& inv_e_metric) {
-    this->z_.set_metric(inv_e_metric);
+  void set_inv_metric(const Eigen::VectorXd& inv_e_metric) {
+    this->z_.set_inv_metric(inv_e_metric);
   }
 
   void set_max_depth(int d) {

src/stan/mcmc/hmc/nuts_classic/adapt_dense_e_nuts_classic.hpp

@@ -29,10 +29,14 @@ class adapt_dense_e_nuts_classic : public dense_e_nuts_classic<Model, BaseRNG>,
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->covar_adaptation_.learn_covariance(
-          this->z_.inv_e_metric_, this->z_.q);
+      Eigen::MatrixXd inv_metric;
+
+      bool update
+          = this->covar_adaptation_.learn_covariance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
 
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));

src/stan/mcmc/hmc/nuts_classic/adapt_diag_e_nuts_classic.hpp

@@ -30,10 +30,14 @@ class adapt_diag_e_nuts_classic : public diag_e_nuts_classic<Model, BaseRNG>,
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->var_adaptation_.learn_variance(this->z_.inv_e_metric_,
-                                                         this->z_.q);
+      Eigen::VectorXd inv_metric;
+
+      bool update
+          = this->var_adaptation_.learn_variance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
 
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));

src/stan/mcmc/hmc/nuts_classic/dense_e_nuts_classic.hpp

@@ -19,11 +19,12 @@ class dense_e_nuts_classic
                                                                          rng) {}
 
   // Note that the points don't need to be swapped
-  // here since start.inv_e_metric_ = finish.inv_e_metric_
+  // here since start.get_inv_metric() = finish.get_inv_metric()
   bool compute_criterion(ps_point& start, dense_e_point& finish,
                          Eigen::VectorXd& rho) {
-    return finish.p.transpose() * finish.inv_e_metric_ * (rho - finish.p) > 0
-           && start.p.transpose() * finish.inv_e_metric_ * (rho - start.p) > 0;
+    return finish.p.transpose() * finish.get_inv_metric() * (rho - finish.p) > 0
+           && start.p.transpose() * finish.get_inv_metric() * (rho - start.p)
+                  > 0;
   }
 };
 

src/stan/mcmc/hmc/nuts_classic/diag_e_nuts_classic.hpp

@@ -19,11 +19,13 @@ class diag_e_nuts_classic
                                                                         rng) {}
 
   // Note that the points don't need to be swapped here
-  // since start.inv_e_metric_ = finish.inv_e_metric_
+  // since start.get_inv_metric() = finish.get_inv_metric()
   bool compute_criterion(ps_point& start, diag_e_point& finish,
                          Eigen::VectorXd& rho) {
-    return finish.inv_e_metric_.cwiseProduct(finish.p).dot(rho - finish.p) > 0
-           && finish.inv_e_metric_.cwiseProduct(start.p).dot(rho - start.p) > 0;
+    return finish.get_inv_metric().cwiseProduct(finish.p).dot(rho - finish.p)
+               > 0
+           && finish.get_inv_metric().cwiseProduct(start.p).dot(rho - start.p)
+                  > 0;
   }
 };
 

src/stan/mcmc/hmc/static/adapt_dense_e_static_hmc.hpp

@@ -32,13 +32,16 @@ class adapt_dense_e_static_hmc : public dense_e_static_hmc<Model, BaseRNG>,
                                                 s.accept_stat());
       this->update_L_();
 
-      bool update = this->covar_adaptation_.learn_covariance(
-          this->z_.inv_e_metric_, this->z_.q);
+      Eigen::MatrixXd inv_metric;
+
+      bool update
+          = this->covar_adaptation_.learn_covariance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
         this->update_L_();
-
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));
         this->stepsize_adaptation_.restart();
       }

src/stan/mcmc/hmc/static/adapt_diag_e_static_hmc.hpp

@@ -32,13 +32,16 @@ class adapt_diag_e_static_hmc : public diag_e_static_hmc<Model, BaseRNG>,
                                                 s.accept_stat());
       this->update_L_();
 
-      bool update = this->var_adaptation_.learn_variance(this->z_.inv_e_metric_,
-                                                         this->z_.q);
+      Eigen::VectorXd inv_metric;
+
+      bool update
+          = this->var_adaptation_.learn_variance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
         this->update_L_();
-
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));
         this->stepsize_adaptation_.restart();
       }

src/stan/mcmc/hmc/static/base_static_hmc.hpp

@@ -29,12 +29,12 @@ class base_static_hmc
 
   ~base_static_hmc() {}
 
-  void set_metric(const Eigen::MatrixXd& inv_e_metric) {
-    this->z_.set_metric(inv_e_metric);
+  void set_inv_metric(const Eigen::MatrixXd& inv_e_metric) {
+    this->z_.set_inv_metric(inv_e_metric);
   }
 
-  void set_metric(const Eigen::VectorXd& inv_e_metric) {
-    this->z_.set_metric(inv_e_metric);
+  void set_inv_metric(const Eigen::VectorXd& inv_e_metric) {
+    this->z_.set_inv_metric(inv_e_metric);
   }
 
   sample transition(sample& init_sample, callbacks::logger& logger) {

src/stan/mcmc/hmc/static_uniform/adapt_dense_e_static_uniform.hpp

@@ -31,11 +31,16 @@ class adapt_dense_e_static_uniform
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->covar_adaptation_.learn_covariance(
-          this->z_.inv_e_metric_, this->z_.q);
+      Eigen::MatrixXd inv_metric;
+
+      bool update
+          = this->covar_adaptation_.learn_covariance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
+
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));
         this->stepsize_adaptation_.restart();
       }

src/stan/mcmc/hmc/static_uniform/adapt_diag_e_static_uniform.hpp

@@ -31,10 +31,15 @@ class adapt_diag_e_static_uniform
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->var_adaptation_.learn_variance(this->z_.inv_e_metric_,
-                                                         this->z_.q);
+      Eigen::VectorXd inv_metric;
+
+      bool update
+          = this->var_adaptation_.learn_variance(inv_metric, this->z_.q);
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
+
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));
         this->stepsize_adaptation_.restart();
       }

src/stan/mcmc/hmc/xhmc/adapt_dense_e_xhmc.hpp

@@ -29,10 +29,14 @@ class adapt_dense_e_xhmc : public dense_e_xhmc<Model, BaseRNG>,
       this->stepsize_adaptation_.learn_stepsize(this->nom_epsilon_,
                                                 s.accept_stat());
 
-      bool update = this->covar_adaptation_.learn_covariance(
-          this->z_.inv_e_metric_, this->z_.q);
+      Eigen::MatrixXd inv_metric;
+
+      bool update
+          = this->covar_adaptation_.learn_covariance(inv_metric, this->z_.q);
 
       if (update) {
+        this->z_.set_inv_metric(std::move(inv_metric));
+
         this->init_stepsize(logger);
 
         this->stepsize_adaptation_.set_mu(log(10 * this->nom_epsilon_));