vari.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CORE_VARI_HPP
#define STAN_MATH_REV_CORE_VARI_HPP
 
#include <stan/math/rev/core/var_value_fwd_declare.hpp>
#include <stan/math/rev/core/chainable_alloc.hpp>
#include <stan/math/rev/core/chainablestack.hpp>
#include <stan/math/rev/core/arena_matrix.hpp>
#include <stan/math/prim/meta.hpp>
#include <ostream>
#include <type_traits>
 
namespace stan {
namespace math {
 
// forward declaration of vari_value
template <typename T, typename = void>
class vari_value;
 
class vari_base {
 public:
  virtual void chain() = 0;
  virtual void set_zero_adjoint() = 0;
 
  static inline void* operator new(size_t nbytes) noexcept {
    return ChainableStack::instance_->memalloc_.alloc(nbytes);
  }
 
  static inline void operator delete(
      void* /* ignore arg */) noexcept { /* no op */
  }
};
 
template <typename T>
class vari_value<T, require_t<std::is_floating_point<T>>> : public vari_base {
 public:
  using value_type = std::decay_t<T>;
  const value_type val_;
  value_type adj_{0.0};
 
  template <typename S, require_convertible_t<S&, T>* = nullptr>
  vari_value(S x) noexcept : val_(x) {  // NOLINT
    ChainableStack::instance_->var_stack_.push_back(this);
  }
 
  template <typename S, require_convertible_t<S&, T>* = nullptr>
  vari_value(S x, bool stacked) noexcept : val_(x) {
    if (stacked) {
      ChainableStack::instance_->var_stack_.push_back(this);
    } else {
      ChainableStack::instance_->var_nochain_stack_.push_back(this);
    }
  }
 
  inline const auto& val() const { return val_; }
 
  inline auto& adj() const { return adj_; }
 
  inline auto& adj() { return adj_; }
 
  inline void chain() {}
 
  inline void init_dependent() noexcept { adj_ = 1.0; }
 
  inline void set_zero_adjoint() noexcept final { adj_ = 0.0; }
 
  friend std::ostream& operator<<(std::ostream& os, const vari_value<T>* v) {
    return os << v->val_ << ":" << v->adj_;
  }
 
 private:
  template <typename, typename>
  friend class var_value;
};
 
// For backwards compatability the default is double
using vari = vari_value<double>;
 
template <typename T, typename = void>
class vari_view;
 
template <typename Derived>
class vari_view_eigen {
 private:
  vari_view_eigen() = default;
  friend Derived;
 
  inline Derived& derived() { return static_cast<Derived&>(*this); }
  inline const Derived& derived() const {
    return static_cast<const Derived&>(*this);
  }
 
 public:
  inline auto block(Eigen::Index start_row, Eigen::Index start_col,
                    Eigen::Index num_rows, Eigen::Index num_cols) const {
    using inner_type = decltype(
        derived().val_.block(start_row, start_col, num_rows, num_cols));
    return vari_view<inner_type>(
        derived().val_.block(start_row, start_col, num_rows, num_cols),
        derived().adj_.block(start_row, start_col, num_rows, num_cols));
  }
  inline auto block(Eigen::Index start_row, Eigen::Index start_col,
                    Eigen::Index num_rows, Eigen::Index num_cols) {
    using inner_type = decltype(
        derived().val_.block(start_row, start_col, num_rows, num_cols));
    return vari_view<inner_type>(
        derived().val_.block(start_row, start_col, num_rows, num_cols),
        derived().adj_.block(start_row, start_col, num_rows, num_cols));
  }
 
  inline auto head(Eigen::Index n) const {
    using inner_type = decltype(derived().val_.head(n));
    return vari_view<inner_type>(derived().val_.head(n),
                                 derived().adj_.head(n));
  }
  inline auto head(Eigen::Index n) {
    using inner_type = decltype(derived().val_.head(n));
    return vari_view<inner_type>(derived().val_.head(n),
                                 derived().adj_.head(n));
  }
 
  inline auto tail(Eigen::Index n) const {
    using inner_type = decltype(derived().val_.tail(n));
    return vari_view<inner_type>(derived().val_.tail(n),
                                 derived().adj_.tail(n));
  }
  inline auto tail(Eigen::Index n) {
    using inner_type = decltype(derived().val_.tail(n));
    return vari_view<inner_type>(derived().val_.tail(n),
                                 derived().adj_.tail(n));
  }
 
  inline auto segment(Eigen::Index i, Eigen::Index n) const {
    using inner_type = decltype(derived().val_.segment(i, n));
    return vari_view<inner_type>(derived().val_.segment(i, n),
                                 derived().adj_.segment(i, n));
  }
  inline auto segment(Eigen::Index i, Eigen::Index n) {
    using inner_type = decltype(derived().val_.segment(i, n));
    return vari_view<inner_type>(derived().val_.segment(i, n),
                                 derived().adj_.segment(i, n));
  }
 
  inline auto transpose() const {
    using inner_type = decltype(derived().val_.transpose());
    return vari_view<inner_type>(derived().val_.transpose(),
                                 derived().adj_.transpose());
  }
  inline auto transpose() {
    using inner_type = decltype(derived().val_.transpose());
    return vari_view<inner_type>(derived().val_.transpose(),
                                 derived().adj_.transpose());
  }
 
  inline auto row(Eigen::Index i) const {
    using inner_type = decltype(derived().val_.row(i));
    return vari_view<inner_type>(derived().val_.row(i), derived().adj_.row(i));
  }
  inline auto row(Eigen::Index i) {
    using inner_type = decltype(derived().val_.row(i));
    return vari_view<inner_type>(derived().val_.row(i), derived().adj_.row(i));
  }
 
  inline auto col(Eigen::Index i) const {
    using inner_type = decltype(derived().val_.col(i));
    return vari_view<inner_type>(derived().val_.col(i), derived().adj_.col(i));
  }
  inline auto col(Eigen::Index i) {
    using inner_type = decltype(derived().val_.col(i));
    return vari_view<inner_type>(derived().val_.col(i), derived().adj_.col(i));
  }
 
  inline auto diagonal() const {
    using inner_type = decltype(derived().val_.diagonal());
    return vari_view<inner_type>(derived().val_.diagonal(),
                                 derived().adj_.diagonal());
  }
  inline auto diagonal() {
    using inner_type = decltype(derived().val_.diagonal());
    return vari_view<inner_type>(derived().val_.diagonal(),
                                 derived().adj_.diagonal());
  }
 
  inline auto coeff(Eigen::Index i, Eigen::Index j) const {
    return vari_value<double>(derived().val_.coeffRef(i, j),
                              derived().adj_.coeffRef(i, j));
  }
  inline auto coeff(Eigen::Index i, Eigen::Index j) {
    return vari_value<double>(derived().val_.coeffRef(i, j),
                              derived().adj_.coeffRef(i, j));
  }
 
  inline auto coeff(Eigen::Index i) const {
    return vari_value<double>(derived().val_.coeffRef(i),
                              derived().adj_.coeffRef(i));
  }
  inline auto coeff(Eigen::Index i) {
    return vari_value<double>(derived().val_.coeffRef(i),
                              derived().adj_.coeffRef(i));
  }
 
  inline auto operator()(Eigen::Index i) const { return this->coeff(i); }
  inline auto operator()(Eigen::Index i) { return this->coeff(i); }
 
  inline auto operator()(Eigen::Index i, Eigen::Index j) const {
    return this->coeff(i, j);
  }
  inline auto operator()(Eigen::Index i, Eigen::Index j) {
    return this->coeff(i, j);
  }
 
  inline auto rowwise_reverse() const {
    using inner_type = decltype(derived().val_.rowwise().reverse());
    return vari_view<inner_type>(derived().val_.rowwise().reverse(),
                                 derived().adj_.rowwise().reverse());
  }
  inline auto rowwise_reverse() {
    using inner_type = decltype(derived().val_.rowwise().reverse());
    return vari_view<inner_type>(derived().val_.rowwise().reverse(),
                                 derived().adj_.rowwise().reverse());
  }
 
  inline auto colwise_reverse() const {
    using inner_type = decltype(derived().val_.colwise().reverse());
    return vari_view<inner_type>(derived().val_.colwise().reverse(),
                                 derived().adj_.colwise().reverse());
  }
  inline auto colwise_reverse() {
    using inner_type = decltype(derived().val_.colwise().reverse());
    return vari_view<inner_type>(derived().val_.colwise().reverse(),
                                 derived().adj_.colwise().reverse());
  }
 
  inline auto reverse() const {
    using inner_type = decltype(derived().val_.reverse());
    return vari_view<inner_type>(derived().val_.reverse(),
                                 derived().adj_.reverse());
  }
  inline auto reverse() {
    using inner_type = decltype(derived().val_.reverse());
    return vari_view<inner_type>(derived().val_.reverse(),
                                 derived().adj_.reverse());
  }
 
  inline auto topRows(Eigen::Index n) const {
    using inner_type = decltype(derived().val_.topRows(n));
    return vari_view<inner_type>(derived().val_.topRows(n),
                                 derived().adj_.topRows(n));
  }
  inline auto topRows(Eigen::Index n) {
    using inner_type = decltype(derived().val_.topRows(n));
    return vari_view<inner_type>(derived().val_.topRows(n),
                                 derived().adj_.topRows(n));
  }
 
  inline auto bottomRows(Eigen::Index n) const {
    using inner_type = decltype(derived().val_.bottomRows(n));
    return vari_view<inner_type>(derived().val_.bottomRows(n),
                                 derived().adj_.bottomRows(n));
  }
  inline auto bottomRows(Eigen::Index n) {
    using inner_type = decltype(derived().val_.bottomRows(n));
    return vari_view<inner_type>(derived().val_.bottomRows(n),
                                 derived().adj_.bottomRows(n));
  }
 
  inline auto middleRows(Eigen::Index start_row, Eigen::Index n) const {
    using inner_type = decltype(derived().val_.middleRows(start_row, n));
    return vari_view<inner_type>(derived().val_.middleRows(start_row, n),
                                 derived().adj_.middleRows(start_row, n));
  }
  inline auto middleRows(Eigen::Index start_row, Eigen::Index n) {
    using inner_type = decltype(derived().val_.middleRows(start_row, n));
    return vari_view<inner_type>(derived().val_.middleRows(start_row, n),
                                 derived().adj_.middleRows(start_row, n));
  }
 
  inline auto leftCols(Eigen::Index n) const {
    using inner_type = decltype(derived().val_.leftCols(n));
    return vari_view<inner_type>(derived().val_.leftCols(n),
                                 derived().adj_.leftCols(n));
  }
  inline auto leftCols(Eigen::Index n) {
    using inner_type = decltype(derived().val_.leftCols(n));
    return vari_view<inner_type>(derived().val_.leftCols(n),
                                 derived().adj_.leftCols(n));
  }
 
  inline auto rightCols(Eigen::Index n) const {
    using inner_type = decltype(derived().val_.rightCols(n));
    return vari_view<inner_type>(derived().val_.rightCols(n),
                                 derived().adj_.rightCols(n));
  }
  inline auto rightCols(Eigen::Index n) {
    using inner_type = decltype(derived().val_.rightCols(n));
    return vari_view<inner_type>(derived().val_.rightCols(n),
                                 derived().adj_.rightCols(n));
  }
 
  inline auto middleCols(Eigen::Index start_col, Eigen::Index n) const {
    using inner_type = decltype(derived().val_.middleCols(start_col, n));
    return vari_view<inner_type>(derived().val_.middleCols(start_col, n),
                                 derived().adj_.middleCols(start_col, n));
  }
  inline auto middleCols(Eigen::Index start_col, Eigen::Index n) {
    using inner_type = decltype(derived().val_.middleCols(start_col, n));
    return vari_view<inner_type>(derived().val_.middleCols(start_col, n),
                                 derived().adj_.middleCols(start_col, n));
  }
 
  inline auto array() const {
    using inner_type = decltype(derived().val_.array());
    return vari_view<inner_type>(derived().val_.array(),
                                 derived().adj_.array());
  }
  inline auto array() {
    using inner_type = decltype(derived().val_.array());
    return vari_view<inner_type>(derived().val_.array(),
                                 derived().adj_.array());
  }
 
  inline auto matrix() const {
    using inner_type = decltype(derived().val_.matrix());
    return vari_view<inner_type>(derived().val_.matrix(),
                                 derived().adj_.matrix());
  }
  inline auto matrix() {
    using inner_type = decltype(derived().val_.matrix());
    return vari_view<inner_type>(derived().val_.matrix(),
                                 derived().adj_.matrix());
  }
 
  inline Eigen::Index rows() const { return derived().val_.rows(); }
  inline Eigen::Index cols() const { return derived().val_.cols(); }
  inline Eigen::Index size() const { return derived().val_.size(); }
};
 
template <typename T>
class vari_view<
    T, require_all_t<is_eigen<T>, bool_constant<!is_plain_type<T>::value>>>
    final : public vari_base,
            public vari_view_eigen<vari_view<T, require_not_plain_type_t<T>>> {
 public:
  using PlainObject = plain_type_t<T>;
  using value_type = std::decay_t<T>;  // The underlying type for this class
  static constexpr int RowsAtCompileTime = PlainObject::RowsAtCompileTime;
  static constexpr int ColsAtCompileTime = PlainObject::ColsAtCompileTime;
 
  T val_;
  T adj_;
  template <typename S, typename K,
            require_assignable_t<value_type, S>* = nullptr,
            require_assignable_t<value_type, K>* = nullptr>
  vari_view(const S& val, const K& adj) noexcept : val_(val), adj_(adj) {}
 
  inline const auto& val() const noexcept { return val_; }
  inline auto& val_op() noexcept { return val_; }
 
  inline auto& adj() noexcept { return adj_; }
  inline auto& adj() const noexcept { return adj_; }
  inline auto& adj_op() noexcept { return adj_; }
 
  void set_zero_adjoint() {}
  void chain() {}
};
 
template <typename T>
class vari_value<T, require_all_t<is_plain_type<T>, is_eigen_dense_base<T>>>
    : public vari_base,
      public vari_view_eigen<vari_value<
          T, require_all_t<is_plain_type<T>, is_eigen_dense_base<T>>>> {
 public:
  using PlainObject = plain_type_t<T>;
  using value_type = PlainObject;  // The underlying type for this class
  using eigen_scalar = value_type_t<PlainObject>;  // A floating point type
  using vari_type = vari_value<T>;
  static constexpr int RowsAtCompileTime = PlainObject::RowsAtCompileTime;
  static constexpr int ColsAtCompileTime = PlainObject::ColsAtCompileTime;
 
  arena_matrix<PlainObject> val_;
 
  arena_matrix<PlainObject> adj_;
 
  template <typename S, require_assignable_t<T, S>* = nullptr>
  explicit vari_value(const S& x)
      : val_(x),
        adj_((RowsAtCompileTime == 1 && S::ColsAtCompileTime == 1)
                     || (ColsAtCompileTime == 1 && S::RowsAtCompileTime == 1)
                 ? x.cols()
                 : x.rows(),
             (RowsAtCompileTime == 1 && S::ColsAtCompileTime == 1)
                     || (ColsAtCompileTime == 1 && S::RowsAtCompileTime == 1)
                 ? x.rows()
                 : x.cols()) {
    adj_.setZero();
    ChainableStack::instance_->var_stack_.push_back(this);
  }
 
  template <typename S, require_assignable_t<T, S>* = nullptr>
  vari_value(const S& x, bool stacked)
      : val_(x),
        adj_((RowsAtCompileTime == 1 && S::ColsAtCompileTime == 1)
                     || (ColsAtCompileTime == 1 && S::RowsAtCompileTime == 1)
                 ? x.cols()
                 : x.rows(),
             (RowsAtCompileTime == 1 && S::ColsAtCompileTime == 1)
                     || (ColsAtCompileTime == 1 && S::RowsAtCompileTime == 1)
                 ? x.rows()
                 : x.cols()) {
    adj_.setZero();
    if (stacked) {
      ChainableStack::instance_->var_stack_.push_back(this);
    } else {
      ChainableStack::instance_->var_nochain_stack_.push_back(this);
    }
  }
 
 protected:
  template <typename S, require_not_same_t<T, S>* = nullptr>
  explicit vari_value(const vari_value<S>* x) : val_(x->val_), adj_(x->adj_) {}
 
 public:
  inline const auto& val() const noexcept { return val_; }
  inline auto& val_op() noexcept { return val_; }
 
  inline auto& adj() noexcept { return adj_; }
  inline auto& adj() const noexcept { return adj_; }
  inline auto& adj_op() noexcept { return adj_; }
 
  virtual void chain() {}
  inline void init_dependent() { adj_.setOnes(); }
 
  inline void set_zero_adjoint() final { adj_.setZero(); }
 
  friend std::ostream& operator<<(std::ostream& os, const vari_value<T>* v) {
    return os << "val: \n" << v->val_ << " \nadj: \n" << v->adj_;
  }
 
 private:
  template <typename, typename>
  friend class var_value;
};
 
template <typename T>
class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
 public:
  using PlainObject = plain_type_t<T>;  // Base type of Eigen class
  using value_type = PlainObject;       // vari's adj_ and val_ member type
  using InnerIterator = typename arena_matrix<PlainObject>::InnerIterator;
  static constexpr int RowsAtCompileTime = T::RowsAtCompileTime;
  static constexpr int ColsAtCompileTime = T::ColsAtCompileTime;
 
  arena_matrix<PlainObject> val_;
  arena_matrix<PlainObject> adj_;
 
  template <typename S, require_convertible_t<S&, T>* = nullptr>
  explicit vari_value(S&& x)
      : val_(std::forward<S>(x)),
        adj_(val_.rows(), val_.cols(), val_.nonZeros(), val_.outerIndexPtr(),
             val_.innerIndexPtr(),
             arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
             val_.innerNonZeroPtr()) {
    ChainableStack::instance_->var_stack_.push_back(this);
  }
 
  vari_value(const arena_matrix<PlainObject>& val,
             const arena_matrix<PlainObject>& adj)
      : val_(val), adj_(adj) {
    ChainableStack::instance_->var_stack_.push_back(this);
  }
 
  template <typename S, require_convertible_t<S&, T>* = nullptr>
  vari_value(S&& x, bool stacked)
      : val_(std::forward<S>(x)),
        adj_(val_.rows(), val_.cols(), val_.nonZeros(), val_.outerIndexPtr(),
             val_.innerIndexPtr(),
             arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
             val_.innerNonZeroPtr()) {
    if (stacked) {
      ChainableStack::instance_->var_stack_.push_back(this);
    } else {
      ChainableStack::instance_->var_nochain_stack_.push_back(this);
    }
  }
 
  Eigen::Index rows() const { return val_.rows(); }
  Eigen::Index cols() const { return val_.cols(); }
  Eigen::Index size() const { return val_.size(); }
 
  inline const auto& val() const { return val_; }
  inline auto& val_op() { return val_; }
 
  inline auto& adj() { return adj_; }
  inline auto& adj() const { return adj_; }
  inline auto& adj_op() { return adj_; }
 
  void chain() {}
  inline void init_dependent() {
    std::fill(adj_.valuePtr(), adj_.valuePtr() + adj_.nonZeros(), 1.0);
  }
 
  inline void set_zero_adjoint() noexcept final {
    std::fill(adj_.valuePtr(), adj_.valuePtr() + adj_.nonZeros(), 0.0);
  }
 
  friend std::ostream& operator<<(std::ostream& os, const vari_value<T>* v) {
    return os << "val: \n" << v->val_ << " \nadj: \n" << v->adj_;
  }
 
 private:
  template <typename, typename>
  friend class var_value;
};
 
}  // namespace math
}  // namespace stan
#endif