cov_matrix_constrain.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CONSTRAINT_COV_MATRIX_CONSTRAIN_HPP
#define STAN_MATH_REV_CONSTRAINT_COV_MATRIX_CONSTRAIN_HPP
 
#include <stan/math/rev/core.hpp>
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/prim/fun/exp.hpp>
#include <stan/math/prim/fun/log.hpp>
#include <stan/math/rev/fun/multiply_lower_tri_self_transpose.hpp>
#include <cmath>
 
namespace stan {
namespace math {
 
template <typename T, require_var_vector_t<T>* = nullptr>
var_value<Eigen::MatrixXd> cov_matrix_constrain(const T& x, Eigen::Index K) {
  using std::exp;
 
  check_size_match("cov_matrix_constrain", "x.size()", x.size(),
                   "K + (K choose 2)", (K * (K + 1)) / 2);
  arena_t<Eigen::MatrixXd> L_val = Eigen::MatrixXd::Zero(K, K);
  int i = 0;
  for (Eigen::Index m = 0; m < K; ++m) {
    L_val.row(m).head(m) = x.val().segment(i, m);
    i += m;
    L_val.coeffRef(m, m) = exp(x.val().coeff(i));
    i++;
  }
 
  var_value<Eigen::MatrixXd> L = L_val;
 
  reverse_pass_callback([x, L]() mutable {
    Eigen::Index K = L.rows();
    int i = x.size();
    for (int m = K - 1; m >= 0; --m) {
      i--;
      x.adj().coeffRef(i) += L.adj().coeff(m, m) * L.val().coeff(m, m);
      i -= m;
      x.adj().segment(i, m) += L.adj().row(m).head(m);
    }
  });
 
  return multiply_lower_tri_self_transpose(L);
}
 
template <typename T, require_var_vector_t<T>* = nullptr>
var_value<Eigen::MatrixXd> cov_matrix_constrain(const T& x, Eigen::Index K,
                                                scalar_type_t<T>& lp) {
  using std::exp;
  using std::log;
 
  check_size_match("cov_matrix_constrain", "x.size()", x.size(),
                   "K + (K choose 2)", (K * (K + 1)) / 2);
  arena_t<Eigen::MatrixXd> L_val = Eigen::MatrixXd::Zero(K, K);
  int pos = 0;
  for (Eigen::Index m = 0; m < K; ++m) {
    L_val.row(m).head(m) = x.val().segment(pos, m);
    pos += m;
    L_val.coeffRef(m, m) = exp(x.val().coeff(pos));
    pos++;
  }
 
  // Jacobian for complete transform, including exp() above
  double lp_val = (K * LOG_TWO);
  for (Eigen::Index k = 0; k < K; ++k) {
    // only +1 because index from 0
    lp_val += (K - k + 1) * log(L_val.coeff(k, k));
  }
 
  lp += lp_val;
 
  var_value<Eigen::MatrixXd> L = L_val;
 
  reverse_pass_callback([x, L, lp]() mutable {
    Eigen::Index K = L.rows();
    for (Eigen::Index k = 0; k < K; ++k) {
      L.adj().coeffRef(k, k) += (K - k + 1) * lp.adj() / L.val().coeff(k, k);
    }
    int pos = x.size();
    for (int m = K - 1; m >= 0; --m) {
      pos--;
      x.adj().coeffRef(pos) += L.adj().coeff(m, m) * L.val().coeff(m, m);
      pos -= m;
      x.adj().segment(pos, m) += L.adj().row(m).head(m);
    }
  });
 
  return multiply_lower_tri_self_transpose(L);
}
 
}  // namespace math
}  // namespace stan
 
#endif