sum_to_zero_constrain.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CONSTRAINT_SUM_TO_ZERO_CONSTRAIN_HPP
#define STAN_MATH_REV_CONSTRAINT_SUM_TO_ZERO_CONSTRAIN_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core/reverse_pass_callback.hpp>
#include <stan/math/rev/core/arena_matrix.hpp>
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/prim/fun/inv_sqrt.hpp>
#include <stan/math/prim/fun/sqrt.hpp>
#include <stan/math/prim/constraint/sum_to_zero_constrain.hpp>
#include <cmath>
 
namespace stan {
namespace math {
 
namespace internal {
 
template <typename T>
inline void sum_to_zero_vector_backprop(T&& y_adj,
                                        const Eigen::VectorXd& z_adj) {
  const auto N = y_adj.size();
 
  double sum_u_adj = 0;
  for (int i = 0; i < N; ++i) {
    double n = static_cast<double>(i + 1);
 
    // adjoint of the reverse cumulative sum computed in the forward mode
    sum_u_adj += z_adj.coeff(i);
 
    // adjoint of the offset subtraction
    double v_adj = -z_adj.coeff(i + 1) * n;
 
    double w_adj = v_adj + sum_u_adj;
 
    y_adj.coeffRef(i) += w_adj / sqrt(n * (n + 1));
  }
}
 
}  // namespace internal
 
template <typename T, require_rev_col_vector_t<T>* = nullptr>
inline auto sum_to_zero_constrain(T&& y) {
  using ret_type = plain_type_t<T>;
  if (unlikely(y.size() == 0)) {
    return arena_t<ret_type>(Eigen::VectorXd{{0}});
  }
  auto arena_y = to_arena(std::forward<T>(y));
  arena_t<ret_type> arena_z = sum_to_zero_constrain(arena_y.val());
 
  reverse_pass_callback([arena_y, arena_z]() mutable {
    internal::sum_to_zero_vector_backprop(arena_y.adj(), arena_z.adj());
  });
 
  return arena_z;
}
 
template <typename T, require_rev_matrix_t<T>* = nullptr,
          require_not_t<is_rev_vector<T>>* = nullptr>
inline auto sum_to_zero_constrain(T&& x) {
  using ret_type = plain_type_t<T>;
  if (unlikely(x.size() == 0)) {
    return arena_t<ret_type>(Eigen::MatrixXd{{0}});
  }
  auto arena_x = to_arena(std::forward<T>(x));
  arena_t<ret_type> arena_z = sum_to_zero_constrain(arena_x.val());
 
  reverse_pass_callback([arena_x, arena_z]() mutable {
    const auto Nf = arena_x.val().rows();
    const auto Mf = arena_x.val().cols();
 
    Eigen::VectorXd d_beta = Eigen::VectorXd::Zero(Nf);
 
    for (int j = 0; j < Mf; ++j) {
      double a_j = inv_sqrt((j + 1.0) * (j + 2.0));
      double b_j = (j + 1.0) * a_j;
 
      double d_ax = 0.0;
 
      for (int i = 0; i < Nf; ++i) {
        double a_i = inv_sqrt((i + 1.0) * (i + 2.0));
        double b_i = (i + 1.0) * a_i;
 
        double dY = arena_z.adj().coeff(i, j) - arena_z.adj().coeff(Nf, j)
                    + arena_z.adj().coeff(Nf, Mf) - arena_z.adj().coeff(i, Mf);
        double dI_from_beta = a_j * d_beta.coeff(i);
        d_beta.coeffRef(i) += -dY;
 
        double dI_from_alpha = b_j * dY;
        double dI = dI_from_alpha + dI_from_beta;
        arena_x.adj().coeffRef(i, j) += b_i * dI + a_i * d_ax;
        d_ax -= dI;
      }
    }
  });
 
  return arena_z;
}
 
template <typename T, typename Lp, require_t<is_rev_matrix<T>>* = nullptr>
inline auto sum_to_zero_constrain(T&& y, Lp& lp) {
  return sum_to_zero_constrain(std::forward<T>(y));
}
 
}  // namespace math
}  // namespace stan
#endif