gp_periodic_cov.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_FUN_GP_PERIODIC_COV_HPP
#define STAN_MATH_REV_FUN_GP_PERIODIC_COV_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/fun/value_of.hpp>
#include <stan/math/rev/fun/adjoint_of.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/square.hpp>
#include <stan/math/prim/fun/squared_distance.hpp>
#include <cmath>
#include <type_traits>
#include <vector>
 
namespace stan {
namespace math {
 
template <typename T_x, typename T_sigma, require_st_arithmetic<T_x>* = nullptr,
          require_stan_scalar_t<T_sigma>* = nullptr>
inline Eigen::Matrix<var, Eigen::Dynamic, Eigen::Dynamic> gp_periodic_cov(
    const std::vector<T_x>& x, const T_sigma sigma, const var l, const var p) {
  using std::exp;
  using std::sin;
  const char* fun = "gp_periodic_cov";
  check_positive(fun, "signal standard deviation", sigma);
  check_positive(fun, "length-scale", l);
  check_positive(fun, "period", p);
  size_t x_size = x.size();
  for (size_t i = 0; i < x_size; ++i) {
    check_not_nan(fun, "element of x", x[i]);
  }
 
  Eigen::Matrix<var, -1, -1> cov(x_size, x_size);
  if (x_size == 0) {
    return cov;
  }
  size_t l_tri_size = x_size * (x_size - 1) / 2;
  arena_matrix<Eigen::VectorXd> dists_lin(l_tri_size);
  arena_matrix<Eigen::VectorXd> sin_dists_sq_lin(l_tri_size);
  arena_matrix<Eigen::Matrix<var, -1, 1>> cov_l_tri_lin(l_tri_size);
  arena_matrix<Eigen::Matrix<var, -1, 1>> cov_diag(
      is_constant<T_sigma>::value ? 0 : x_size);
 
  double sigma_sq = square(value_of(sigma));
  double pi_div_p = pi() / value_of(p);
  double neg_two_inv_l_sq = -2.0 / square(value_of(l));
 
  size_t block_size = 10;
  size_t pos = 0;
  for (size_t jb = 0; jb < x_size; jb += block_size) {
    size_t j_end = std::min(x_size, jb + block_size);
    size_t j_size = j_end - jb;
    cov.diagonal().segment(jb, j_size)
        = Eigen::VectorXd::Constant(j_size, sigma_sq);
 
    if (!is_constant<T_sigma>::value) {
      cov_diag.segment(jb, j_size) = cov.diagonal().segment(jb, j_size);
    }
    for (size_t ib = jb; ib < x_size; ib += block_size) {
      size_t i_end = std::min(x_size, ib + block_size);
      for (size_t j = jb; j < j_end; ++j) {
        for (size_t i = std::max(ib, j + 1); i < i_end; ++i) {
          double dist = distance(x[i], x[j]);
          double sin_dist = sin(pi_div_p * dist);
          double sin_dist_sq = square(sin_dist);
          dists_lin.coeffRef(pos) = dist;
          sin_dists_sq_lin.coeffRef(pos) = sin_dist_sq;
          cov_l_tri_lin.coeffRef(pos) = cov.coeffRef(j, i) = cov.coeffRef(i, j)
              = sigma_sq * exp(sin_dist_sq * neg_two_inv_l_sq);
          ++pos;
        }
      }
    }
  }
 
  reverse_pass_callback([cov_l_tri_lin, cov_diag, dists_lin, sin_dists_sq_lin,
                         sigma, l, p, x_size]() {
    size_t l_tri_size = x_size * (x_size - 1) / 2;
    double l_val = value_of(l);
    double p_val = value_of(p);
    double two_pi_div_p = TWO_PI / p_val;
 
    double adjl = 0;
    double adjsigma = 0;
    double adjp = 0;
 
    for (size_t pos = 0; pos < l_tri_size; ++pos) {
      double prod_add
          = cov_l_tri_lin.coeff(pos).val() * cov_l_tri_lin.coeff(pos).adj();
      adjl += prod_add * sin_dists_sq_lin.coeff(pos);
      adjsigma += prod_add;
      double dist = dists_lin.coeff(pos);
      adjp += prod_add * sin(two_pi_div_p * dist) * dist;
    }
    if (!is_constant<T_sigma>::value) {
      adjsigma += (cov_diag.val().array() * cov_diag.adj().array()).sum();
      adjoint_of(sigma) += adjsigma * 2 / value_of(sigma);
    }
    double l_sq = square(l_val);
    l.adj() += adjl * 4 / (l_sq * l_val);
    p.adj() += adjp * TWO_PI / (l_sq * square(p_val));
  });
 
  return cov;
}  // namespace math
 
}  // namespace math
}  // namespace stan
#endif