operator_addition.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CORE_OPERATOR_ADDITION_HPP
#define STAN_MATH_REV_CORE_OPERATOR_ADDITION_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core/var.hpp>
#include <stan/math/prim/err/check_matching_dims.hpp>
#include <stan/math/rev/core/callback_vari.hpp>
#include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
#include <stan/math/prim/fun/as_array_or_scalar.hpp>
#include <stan/math/prim/fun/constants.hpp>
 
namespace stan {
namespace math {
 
inline var operator+(const var& a, const var& b) {
  return make_callback_vari(a.vi_->val_ + b.vi_->val_,
                            [avi = a.vi_, bvi = b.vi_](const auto& vi) mutable {
                              avi->adj_ += vi.adj_;
                              bvi->adj_ += vi.adj_;
                            });
}
 
template <typename Arith, require_arithmetic_t<Arith>* = nullptr>
inline var operator+(const var& a, Arith b) {
  if (unlikely(b == 0.0)) {
    return a;
  }
  return make_callback_vari(
      a.vi_->val_ + b,
      [avi = a.vi_](const auto& vi) mutable { avi->adj_ += vi.adj_; });
}
 
template <typename Arith, require_arithmetic_t<Arith>* = nullptr>
inline var operator+(Arith a, const var& b) {
  return b + a;  // by symmetry
}
 
template <typename VarMat1, typename VarMat2,
          require_all_rev_matrix_t<VarMat1, VarMat2>* = nullptr>
inline auto add(VarMat1&& a, VarMat2&& b) {
  check_matching_dims("add", "a", a, "b", b);
  using op_ret_type = decltype(a.val() + b.val());
  using ret_type = return_var_matrix_t<op_ret_type, VarMat1, VarMat2>;
  arena_t<VarMat1> arena_a(std::forward<VarMat1>(a));
  arena_t<VarMat2> arena_b(std::forward<VarMat2>(b));
  arena_t<ret_type> ret(arena_a.val() + arena_b.val());
  reverse_pass_callback([ret, arena_a, arena_b]() mutable {
    for (Eigen::Index j = 0; j < ret.cols(); ++j) {
      for (Eigen::Index i = 0; i < ret.rows(); ++i) {
        const auto ref_adj = ret.adj().coeffRef(i, j);
        arena_a.adj().coeffRef(i, j) += ref_adj;
        arena_b.adj().coeffRef(i, j) += ref_adj;
      }
    }
  });
  return ret;
}
 
template <typename Arith, typename VarMat,
          require_st_arithmetic<Arith>* = nullptr,
          require_rev_matrix_t<VarMat>* = nullptr>
inline auto add(VarMat&& a, const Arith& b) {
  if (is_eigen<Arith>::value) {
    check_matching_dims("add", "a", a, "b", b);
  }
  using op_ret_type
      = decltype((a.val().array() + as_array_or_scalar(b)).matrix());
  using ret_type = return_var_matrix_t<op_ret_type, VarMat>;
  arena_t<VarMat> arena_a(std::forward<VarMat>(a));
  arena_t<ret_type> ret(arena_a.val().array() + as_array_or_scalar(b));
  reverse_pass_callback(
      [ret, arena_a]() mutable { arena_a.adj() += ret.adj_op(); });
  return ret;
}
 
template <typename Arith, typename VarMat,
          require_st_arithmetic<Arith>* = nullptr,
          require_rev_matrix_t<VarMat>* = nullptr>
inline auto add(const Arith& a, VarMat&& b) {
  return add(std::forward<VarMat>(b), a);
}
 
template <typename Var, typename EigMat,
          require_var_vt<std::is_arithmetic, Var>* = nullptr,
          require_eigen_vt<std::is_arithmetic, EigMat>* = nullptr>
inline auto add(const Var& a, const EigMat& b) {
  using ret_type = return_var_matrix_t<EigMat>;
  arena_t<ret_type> ret(a.val() + b.array());
  reverse_pass_callback([ret, a]() mutable { a.adj() += ret.adj().sum(); });
  return ret;
}
 
template <typename EigMat, typename Var,
          require_eigen_vt<std::is_arithmetic, EigMat>* = nullptr,
          require_var_vt<std::is_arithmetic, Var>* = nullptr>
inline auto add(const EigMat& a, const Var& b) {
  return add(b, a);
}
 
template <typename Var, typename VarMat,
          require_var_vt<std::is_arithmetic, Var>* = nullptr,
          require_rev_matrix_t<VarMat>* = nullptr>
inline auto add(const Var& a, VarMat&& b) {
  using ret_type = return_var_matrix_t<VarMat>;
  arena_t<VarMat> arena_b(std::forward<VarMat>(b));
  arena_t<ret_type> ret(a.val() + arena_b.val().array());
  reverse_pass_callback([ret, a, arena_b]() mutable {
    for (Eigen::Index j = 0; j < ret.cols(); ++j) {
      for (Eigen::Index i = 0; i < ret.rows(); ++i) {
        const auto ret_adj = ret.adj().coeffRef(i, j);
        a.adj() += ret_adj;
        arena_b.adj().coeffRef(i, j) += ret_adj;
      }
    }
  });
  return ret;
}
 
template <typename Var, typename VarMat,
          require_var_vt<std::is_arithmetic, Var>* = nullptr,
          require_rev_matrix_t<VarMat>* = nullptr>
inline auto add(VarMat&& a, const Var& b) {
  return add(b, std::forward<VarMat>(a));
}
 
template <typename T1, typename T2,
          require_any_var_vt<std::is_arithmetic, T1, T2>* = nullptr,
          require_any_arithmetic_t<T1, T2>* = nullptr>
inline auto add(const T1& a, const T2& b) {
  return a + b;
}
 
template <typename T1, typename T2,
          require_all_var_vt<std::is_arithmetic, T1, T2>* = nullptr>
inline auto add(const T1& a, const T2& b) {
  return a + b;
}
 
template <typename VarMat1, typename VarMat2,
          require_any_var_matrix_t<VarMat1, VarMat2>* = nullptr>
inline auto operator+(VarMat1&& a, VarMat2&& b) {
  return add(std::forward<VarMat1>(a), std::forward<VarMat2>(b));
}
 
}  // namespace math
}  // namespace stan
#endif