laplace_likelihood.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_MIX_FUNCTOR_LAPLACE_LIKELIHOOD_HPP
#define STAN_MATH_MIX_FUNCTOR_LAPLACE_LIKELIHOOD_HPP
 
#include <stan/math/mix/functor/hessian_block_diag.hpp>
#include <stan/math/mix/functor/conditional_copy_and_promote.hpp>
#include <stan/math/prim/functor.hpp>
#include <stan/math/prim/fun.hpp>
 
namespace stan {
namespace math {
 
namespace laplace_likelihood {
 
namespace internal {
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_t<Theta>* = nullptr>
inline auto log_likelihood(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
  return std::forward<F>(f)(std::forward<Theta>(theta),
                            std::forward<Args>(args)..., msgs);
}
 
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
inline auto theta_grad(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  nested_rev_autodiff nested;
  arena_t<Matrix<var, Dynamic, 1>> theta_var = theta;
  var f_var = f(theta_var, args..., msgs);
  grad(f_var.vi_);
  return theta_var.adj().eval();
}
 
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
inline void ll_arg_grad(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  nested_rev_autodiff nested;
  var f_var = f(theta, args..., msgs);
  grad(f_var.vi_);
}
 
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
inline auto diagonal_hessian(F&& f, Theta&& theta, Stream* msgs,
                             Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  const Eigen::Index theta_size = theta.size();
  auto v = Eigen::VectorXd::Ones(theta_size);
  Eigen::VectorXd hessian_v = Eigen::VectorXd::Zero(theta_size);
  hessian_times_vector(f, hessian_v, std::forward<Theta>(theta), std::move(v),
                       value_of(args)..., msgs);
  return (-hessian_v).eval();
}
 
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
inline auto block_hessian(F&& f, Theta&& theta,
                          const Eigen::Index hessian_block_size, Stream* msgs,
                          Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  const Eigen::Index theta_size = theta.size();
  if (hessian_block_size == 1) {
    auto v = Eigen::VectorXd::Ones(theta_size);
    Eigen::VectorXd hessian_v = Eigen::VectorXd::Zero(theta_size);
    hessian_times_vector(f, hessian_v, std::forward<Theta>(theta), std::move(v),
                         value_of(args)..., msgs);
    Eigen::SparseMatrix<double> hessian_theta(theta_size, theta_size);
    hessian_theta.reserve(Eigen::VectorXi::Constant(theta_size, 1));
    for (Eigen::Index i = 0; i < theta_size; i++) {
      hessian_theta.insert(i, i) = hessian_v(i);
    }
    return (-hessian_theta).eval();
  } else {
    return (-hessian_block_diag(f, std::forward<Theta>(theta),
                                hessian_block_size, value_of(args)..., msgs))
        .eval();
  }
}
 
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
inline auto diff(F&& f, Theta&& theta, const Eigen::Index hessian_block_size,
                 Stream* msgs, Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  const Eigen::Index theta_size = theta.size();
  auto theta_gradient = [&theta, &f, &msgs](auto&&... args) {
    nested_rev_autodiff nested;
    Matrix<var, Dynamic, 1> theta_var = theta;
    var f_var = f(theta_var, args..., msgs);
    grad(f_var.vi_);
    return theta_var.adj().eval();
  }(args...);
  if (hessian_block_size == 1) {
    auto v = Eigen::VectorXd::Ones(theta_size);
    Eigen::VectorXd hessian_v = Eigen::VectorXd::Zero(theta_size);
    hessian_times_vector(f, hessian_v, std::forward<Theta>(theta), std::move(v),
                         value_of(args)..., msgs);
    Eigen::SparseMatrix<double> hessian_theta(theta_size, theta_size);
    hessian_theta.reserve(Eigen::VectorXi::Constant(theta_size, 1));
    for (Eigen::Index i = 0; i < theta_size; i++) {
      hessian_theta.insert(i, i) = hessian_v(i);
    }
    return std::make_pair(std::move(theta_gradient), (-hessian_theta).eval());
  } else {
    return std::make_pair(
        std::move(theta_gradient),
        (-hessian_block_diag(f, std::forward<Theta>(theta), hessian_block_size,
                             value_of(args)..., msgs))
            .eval());
  }
}
 
template <typename F, typename Theta, typename Stream, typename... Args,
          require_eigen_vector_t<Theta>* = nullptr>
inline Eigen::VectorXd third_diff(F&& f, Theta&& theta, Stream&& msgs,
                                  Args&&... args) {
  nested_rev_autodiff nested;
  const Eigen::Index theta_size = theta.size();
  arena_t<Eigen::Matrix<var, Eigen::Dynamic, 1>> theta_var
      = std::forward<Theta>(theta);
  arena_t<Eigen::Matrix<fvar<fvar<var>>, Eigen::Dynamic, 1>> theta_ffvar(
      theta_size);
  for (Eigen::Index i = 0; i < theta_size; ++i) {
    theta_ffvar(i) = fvar<fvar<var>>(fvar<var>(theta_var(i), 1.0), 1.0);
  }
  fvar<fvar<var>> ftheta_ffvar = f(theta_ffvar, args..., msgs);
  grad(ftheta_ffvar.d_.d_.vi_);
  return theta_var.adj().eval();
}
 
template <typename F, typename Theta, typename AMat, typename Stream,
          typename... Args, require_eigen_vector_t<Theta>* = nullptr>
inline auto compute_s2(F&& f, Theta&& theta, AMat&& A,
                       const int hessian_block_size, Stream* msgs,
                       Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  using Eigen::MatrixXd;
  using Eigen::VectorXd;
 
  nested_rev_autodiff nested;
  const Eigen::Index theta_size = theta.size();
  arena_t<Matrix<var, Dynamic, 1>> theta_var = std::forward<Theta>(theta);
  int n_blocks = theta_size / hessian_block_size;
  arena_t<VectorXd> v(theta_size);
  arena_t<VectorXd> w(theta_size);
  Matrix<fvar<fvar<var>>, Dynamic, 1> theta_ffvar(theta_size);
  auto shallow_copy_args
      = stan::math::internal::shallow_copy_vargs<fvar<fvar<var>>>(
          std::forward_as_tuple(args...));
  for (Eigen::Index i = 0; i < hessian_block_size; ++i) {
    nested_rev_autodiff nested;
    v.setZero();
    for (int j = i; j < theta_size; j += hessian_block_size) {
      v(j) = 1;
    }
    w.setZero();
    for (int j = 0; j < n_blocks; ++j) {
      for (int k = 0; k < hessian_block_size; ++k) {
        w(k + j * hessian_block_size)
            = A(k + j * hessian_block_size, i + j * hessian_block_size);
      }
    }
    for (int j = 0; j < theta_size; ++j) {
      theta_ffvar(j) = fvar<fvar<var>>(fvar<var>(theta_var(j), v(j)), w(j));
    }
    fvar<fvar<var>> target_ffvar = stan::math::apply(
        [](auto&& f, auto&& theta_ffvar, auto&& msgs, auto&&... inner_args) {
          return f(theta_ffvar, inner_args..., msgs);
        },
        shallow_copy_args, f, theta_ffvar, msgs);
    grad(target_ffvar.d_.d_.vi_);
  }
  return (0.5 * theta_var.adj()).eval();
}
 
template <typename F, typename V_t, typename Theta, typename Stream,
          typename... Args, require_eigen_vector_t<Theta>* = nullptr>
inline auto diff_eta_implicit(F&& f, V_t&& v, Theta&& theta, Stream* msgs,
                              Args&&... args) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  using Eigen::VectorXd;
  constexpr bool contains_var = is_any_var_scalar<Args...>::value;
  if constexpr (!contains_var) {
    return;
  }
  nested_rev_autodiff nested;
  const Eigen::Index theta_size = theta.size();
  arena_t<Matrix<var, Dynamic, 1>> theta_var = std::forward<Theta>(theta);
  Matrix<fvar<var>, Dynamic, 1> theta_fvar(theta_size);
  for (Eigen::Index i = 0; i < theta_size; i++) {
    theta_fvar(i) = fvar<var>(theta_var(i), v(i));
  }
  auto shallow_copy_args = stan::math::internal::shallow_copy_vargs<fvar<var>>(
      std::forward_as_tuple(args...));
  fvar<var> f_sum = stan::math::apply(
      [](auto&& f, auto&& theta_fvar, auto&& msgs, auto&&... inner_args) {
        return f(theta_fvar, inner_args..., msgs);
      },
      shallow_copy_args, f, theta_fvar, msgs);
  grad(f_sum.d_.vi_);
}
 
}  // namespace internal
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto theta_grad(F&& f, Theta&& theta, TupleArgs&& ll_tup,
                       Stream* msgs = nullptr) {
  return apply(
      [](auto&& f, auto&& theta, auto&& msgs, auto&&... args) {
        return internal::theta_grad(std::forward<decltype(f)>(f),
                                    std::forward<decltype(theta)>(theta), msgs,
                                    std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_tup), std::forward<F>(f),
      std::forward<Theta>(theta), msgs);
}
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto ll_arg_grad(F&& f, Theta&& theta, TupleArgs&& ll_tup,
                        Stream* msgs = nullptr) {
  return apply(
      [](auto&& f, auto&& theta, auto&& msgs, auto&&... args) {
        return internal::ll_arg_grad(std::forward<decltype(f)>(f),
                                     std::forward<decltype(theta)>(theta), msgs,
                                     std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_tup), std::forward<F>(f),
      std::forward<Theta>(theta), msgs);
}
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto diagonal_hessian(F&& f, Theta&& theta, TupleArgs&& ll_tuple,
                             Stream* msgs) {
  return apply(
      [](auto&& f, auto&& theta, auto* msgs, auto&&... args) {
        return internal::diagonal_hessian(
            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
            msgs, std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_tuple), std::forward<F>(f),
      std::forward<Theta>(theta), msgs);
}
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto block_hessian(F&& f, Theta&& theta,
                          const Eigen::Index hessian_block_size,
                          TupleArgs&& ll_tuple, Stream* msgs) {
  return apply(
      [](auto&& f, auto&& theta, auto hessian_block_size, auto* msgs,
         auto&&... args) {
        return internal::block_hessian(
            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
            hessian_block_size, msgs, std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_tuple), std::forward<F>(f),
      std::forward<Theta>(theta), hessian_block_size, msgs);
}
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto log_likelihood(F&& f, Theta&& theta, TupleArgs&& ll_tup,
                           Stream* msgs) {
  return apply(
      [](auto&& f, auto&& theta, auto&& msgs, auto&&... args) {
        return internal::log_likelihood(
            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
            msgs, std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_tup), std::forward<F>(f),
      std::forward<Theta>(theta), msgs);
}
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto diff(F&& f, Theta&& theta, const Eigen::Index hessian_block_size,
                 TupleArgs&& ll_tuple, Stream* msgs) {
  return apply(
      [](auto&& f, auto&& theta, auto hessian_block_size, auto* msgs,
         auto&&... args) {
        return internal::diff(
            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
            hessian_block_size, msgs, std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_tuple), std::forward<F>(f),
      std::forward<Theta>(theta), hessian_block_size, msgs);
}
 
template <typename F, typename Theta, typename TupleArgs, typename Stream,
          require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline Eigen::VectorXd third_diff(F&& f, Theta&& theta, TupleArgs&& ll_args,
                                  Stream* msgs) {
  return apply(
      [](auto&& f, auto&& theta, auto&& msgs, auto&&... args) {
        return internal::third_diff(std::forward<decltype(f)>(f),
                                    std::forward<decltype(theta)>(theta), msgs,
                                    std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_args), std::forward<F>(f),
      std::forward<Theta>(theta), msgs);
}
 
template <typename F, typename Theta, typename AMat, typename TupleArgs,
          typename Stream, require_eigen_vector_t<Theta>* = nullptr,
          require_tuple_t<TupleArgs>* = nullptr>
inline auto compute_s2(F&& f, Theta&& theta, AMat&& A, int hessian_block_size,
                       TupleArgs&& ll_args, Stream* msgs) {
  return apply(
      [](auto&& f, auto&& theta, auto&& A, auto hessian_block_size, auto* msgs,
         auto&&... args) {
        return internal::compute_s2(
            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
            std::forward<decltype(A)>(A), hessian_block_size, msgs,
            std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_args), std::forward<F>(f),
      std::forward<Theta>(theta), std::forward<AMat>(A), hessian_block_size,
      msgs);
}
 
template <typename F, typename V_t, typename Theta, typename TupleArgs,
          typename Stream, require_tuple_t<TupleArgs>* = nullptr,
          require_eigen_vector_t<Theta>* = nullptr>
inline auto diff_eta_implicit(F&& f, V_t&& v, Theta&& theta,
                              TupleArgs&& ll_args, Stream* msgs) {
  return apply(
      [](auto&& f, auto&& v, auto&& theta, auto&& msgs, auto&&... args) {
        return internal::diff_eta_implicit(
            std::forward<decltype(f)>(f), std::forward<decltype(v)>(v),
            std::forward<decltype(theta)>(theta), msgs,
            std::forward<decltype(args)>(args)...);
      },
      std::forward<TupleArgs>(ll_args), std::forward<F>(f),
      std::forward<V_t>(v), std::forward<Theta>(theta), msgs);
}
 
}  // namespace laplace_likelihood
 
}  // namespace math
}  // namespace stan
 
#endif