laplace_marginal_density.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_MIX_FUNCTOR_LAPLACE_MARGINAL_DENSITY_HPP
#define STAN_MATH_MIX_FUNCTOR_LAPLACE_MARGINAL_DENSITY_HPP
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/mix/functor/laplace_likelihood.hpp>
#include <stan/math/mix/functor/laplace_marginal_density_estimator.hpp>
#include <stan/math/mix/functor/conditional_copy_and_promote.hpp>
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/fun.hpp>
#include <stan/math/rev/functor.hpp>
#include <stan/math/prim/fun/to_ref.hpp>
#include <stan/math/prim/functor/iter_tuple_nested.hpp>
#include <unsupported/Eigen/MatrixFunctions>
#include <cmath>
 
namespace stan {
namespace math {
 
template <
    typename LLFun, typename LLTupleArgs, typename CovarFun, typename CovarArgs,
    bool InitTheta,
    require_t<is_all_arithmetic_scalar<CovarArgs, LLTupleArgs>>* = nullptr>
inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args,
                                     CovarFun&& covariance_function,
                                     CovarArgs&& covar_args,
                                     const laplace_options<InitTheta>& options,
                                     std::ostream* msgs) {
  Eigen::MatrixXd covariance = stan::math::apply(
      [msgs, &covariance_function](auto&&... args) {
        return covariance_function(std::forward<decltype(args)>(args)..., msgs);
      },
      std::forward<CovarArgs>(covar_args));
  return internal::laplace_marginal_density_est(
             std::forward<LLFun>(ll_fun), std::forward<LLTupleArgs>(ll_args),
             std::move(covariance), options, msgs)
      .lmd;
}
 
namespace internal {
 
template <bool ZeroInput = false, typename Output, typename Input,
          require_tuple_t<Output>* = nullptr, require_tuple_t<Input>* = nullptr,
          require_t<std::bool_constant<
              std::tuple_size_v<std::decay_t<Output>> == 0>>* = nullptr,
          require_t<std::bool_constant<
              std::tuple_size_v<std::decay_t<Input>> == 0>>* = nullptr>
inline constexpr void copy_compute_s2(Output&& output, Input&& input) {}
 
template <bool ZeroInput = false, typename Output, typename Input,
          require_t<is_all_arithmetic_scalar<Output>>* = nullptr,
          require_t<is_any_var_scalar<Input>>* = nullptr>
inline constexpr void copy_compute_s2(Output&& output, Input&& input) {
  if constexpr (is_tuple_v<Output> && is_tuple_v<Input>) {
    static_assert(
        std::tuple_size<std::decay_t<Output>>::value
            == std::tuple_size<std::decay_t<Input>>::value,
        "INTERNAL ERROR:(laplace_marginal_lpdf) copy_compute_s2 called on "
        "tuples of different sizes. This is an internal error, please report "
        "it: "
        "https://github.com/stan-dev/math/issues");
  }
  return iter_tuple_nested(
      [](auto&& output_i, auto&& input_i) {
        using output_i_t = std::decay_t<decltype(output_i)>;
        if constexpr (is_std_vector_v<output_i_t>) {
          using dbl_map_t = Eigen::Map<Eigen::Matrix<double, -1, 1>>;
          using var_map_t = Eigen::Map<Eigen::Matrix<var, -1, 1>>;
          var_map_t input_map(input_i.data(), input_i.size());
          dbl_map_t(output_i.data(), output_i.size()).array()
              += 0.5 * input_map.adj().array();
          if constexpr (ZeroInput) {
            input_map.adj().setZero();
          }
        } else if constexpr (is_eigen_v<output_i_t>) {
          output_i.array() += 0.5 * input_i.adj().array();
          if constexpr (ZeroInput) {
            input_i.adj().setZero();
          }
        } else if constexpr (is_stan_scalar_v<output_i_t>) {
          output_i += (0.5 * input_i.adj());
          if constexpr (ZeroInput) {
            input_i.adj() = 0;
          }
        } else {
          static_assert(
              sizeof(std::decay_t<output_i_t>*) == 0,
              "INTERNAL ERROR:(laplace_marginal_lpdf) copy_compute_s2 was "
              "not able to deduce the actions needed for the given type. "
              "This is an internal error, please report it: "
              "https://github.com/stan-dev/math/issues");
        }
      },
      std::forward<Output>(output), std::forward<Input>(input));
}
 
template <typename Output, typename Input>
inline void reverse_pass_collect_adjoints(var ret, Output&& output,
                                          Input&& input) {
  if constexpr (is_tuple_v<Output>) {
    stan::math::for_each(
        [ret](auto&& inner_arg, auto&& inner_input) mutable {
          reverse_pass_collect_adjoints(
              ret, std::forward<decltype(inner_arg)>(inner_arg),
              std::forward<decltype(inner_input)>(inner_input));
        },
        std::forward<Output>(output), std::forward<Input>(input));
  } else if constexpr (is_std_vector_containing_tuple_v<Output>) {
    for (std::size_t i = 0; i < output.size(); ++i) {
      reverse_pass_collect_adjoints(ret, output[i], input[i]);
    }
  } else {
    reverse_pass_callback(
        [vi = ret.vi_, arg_arena = to_arena(std::forward<Output>(output)),
         input_arena = to_arena(std::forward<Input>(input))]() mutable {
          collect_adjoints(arg_arena, vi, input_arena);
        });
  }
}
 
}  // namespace internal
 
template <typename LLFun, typename LLTupleArgs, typename CovarFun,
          typename CovarArgs, bool InitTheta,
          require_t<is_any_var_scalar<LLTupleArgs, CovarArgs>>* = nullptr>
inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args,
                                     CovarFun&& covariance_function,
                                     CovarArgs&& covar_args,
                                     const laplace_options<InitTheta>& options,
                                     std::ostream* msgs) {
  auto covar_args_refs = to_ref(std::forward<CovarArgs>(covar_args));
  auto ll_args_refs = to_ref(std::forward<LLTupleArgs>(ll_args));
  // Solver 1, 2, 3
  constexpr bool ll_args_contain_var = is_any_var_scalar<LLTupleArgs>::value;
  auto partial_parm = internal::make_zeroed_arena(ll_args_refs);
  auto covar_args_adj = internal::make_zeroed_arena(covar_args_refs);
  double lmd = 0.0;
  {
    nested_rev_autodiff nested;
    // Make one hard copy here
    auto ll_args_copy = internal::deep_copy_vargs<var>(ll_args_refs);
    auto covar_args_copy = internal::deep_copy_vargs<var>(covar_args_refs);
    auto covariance = stan::math::apply(
        [&covariance_function, &msgs](auto&&... args) {
          if constexpr (is_any_var_scalar_v<decltype(args)...>) {
            return to_var_value(covariance_function(args..., msgs));
          } else {
            return covariance_function(args..., msgs);
          }
        },
        covar_args_copy);
    decltype(auto) covariance_val = value_of(covariance);
    decltype(auto) ll_args_vals = value_of(ll_args_copy);
    auto md_est = internal::laplace_marginal_density_est(
        ll_fun, ll_args_vals, covariance_val, options, msgs);
    auto ll_args_filter = internal::filter_var_scalar_types(ll_args_copy);
    // tuple of references to var types
    // Solver 1, 2
    const bool solver_1_or_2
        = md_est.solver_used == 1 || md_est.solver_used == 2;
    arena_t<Eigen::MatrixXd> R(md_est.theta.size() * solver_1_or_2,
                               md_est.theta.size() * solver_1_or_2);
    // Solver 3
    arena_t<Eigen::MatrixXd> LU_solve_covariance(
        covariance.rows() * (md_est.solver_used == 3),
        covariance.cols() * (md_est.solver_used == 3));
    // Solver 1, 2, 3
    arena_t<Eigen::VectorXd> s2(md_est.theta.size());
    using stan::math::internal::ZeroOut;
    if (md_est.solver_used == 1) {
      if (options.hessian_block_size == 1) {
        arena_t<Eigen::MatrixXd> tmp = md_est.W_r.toDense();
        md_est.L.template triangularView<Eigen::Lower>().solveInPlace(tmp);
        R.noalias() = tmp.transpose() * tmp;
        arena_t<Eigen::MatrixXd> C
            = md_est.L.template triangularView<Eigen::Lower>().solve(
                md_est.W_r * covariance_val);
        if constexpr (ll_args_contain_var) {
          arena_t<Eigen::MatrixXd> A = covariance_val - C.transpose() * C;
          auto s2_tmp = laplace_likelihood::compute_s2(
              ll_fun, md_est.theta, A, options.hessian_block_size, ll_args_copy,
              msgs);
          s2.deep_copy(s2_tmp);
          internal::copy_compute_s2<ZeroOut>(partial_parm, ll_args_filter);
        } else {
          s2.deep_copy(
              (0.5
               * (covariance_val.diagonal() - (C.transpose() * C).diagonal())
                     .cwiseProduct(laplace_likelihood::third_diff(
                         ll_fun, md_est.theta, ll_args_vals, msgs))));
        }
 
      } else {
        arena_t<Eigen::MatrixXd> tmp = md_est.W_r.toDense();
        md_est.L.template triangularView<Eigen::Lower>().solveInPlace(tmp);
        R.noalias() = tmp.transpose() * tmp;
        arena_t<Eigen::MatrixXd> C
            = md_est.L.template triangularView<Eigen::Lower>().solve(
                md_est.W_r * covariance_val);
        arena_t<Eigen::MatrixXd> A = covariance_val - C.transpose() * C;
        auto s2_tmp = laplace_likelihood::compute_s2(ll_fun, md_est.theta, A,
                                                     options.hessian_block_size,
                                                     ll_args_copy, msgs);
        s2.deep_copy(s2_tmp);
        internal::copy_compute_s2<ZeroOut>(partial_parm, ll_args_filter);
      }
    } else if (md_est.solver_used == 2) {
      R = md_est.W_r
          - md_est.W_r * md_est.K_root
                * md_est.L.transpose()
                      .template triangularView<Eigen::Upper>()
                      .solve(
                          md_est.L.template triangularView<Eigen::Lower>()
                              .solve(md_est.K_root.transpose() * md_est.W_r));
 
      arena_t<Eigen::MatrixXd> C
          = md_est.L.template triangularView<Eigen::Lower>().solve(
              md_est.K_root.transpose());
      auto s2_tmp = laplace_likelihood::compute_s2(
          ll_fun, md_est.theta, (C.transpose() * C).eval(),
          options.hessian_block_size, ll_args_copy, msgs);
      s2.deep_copy(s2_tmp);
      internal::copy_compute_s2<ZeroOut>(partial_parm, ll_args_filter);
    } else {  // options.solver with LU decomposition
      LU_solve_covariance = md_est.LU.solve(covariance_val);
      auto I_minus_BinvKW
          = Eigen::MatrixXd::Identity(md_est.W_r.rows(), md_est.W_r.cols())
            - LU_solve_covariance * md_est.W_r;
      R = md_est.W_r * I_minus_BinvKW;  // == W - W B^{-1} K W
      arena_t<Eigen::MatrixXd> A
          = covariance_val - covariance_val * md_est.W_r * LU_solve_covariance;
      auto s2_tmp = laplace_likelihood::compute_s2(ll_fun, md_est.theta, A,
                                                   options.hessian_block_size,
                                                   ll_args_copy, msgs);
      s2.deep_copy(s2_tmp);
      internal::copy_compute_s2<ZeroOut>(partial_parm, ll_args_filter);
    }
    if constexpr (is_any_var_scalar_v<scalar_type_t<CovarArgs>>) {
      arena_t<Eigen::MatrixXd> K_adj_arena
          = 0.5 * md_est.a * md_est.a.transpose() - 0.5 * R
            + s2 * md_est.theta_grad.transpose()
            - (R * (covariance.val() * s2)) * md_est.theta_grad.transpose();
      var Z = make_callback_var(
          0.0, [covariance, K_adj_arena](auto&& vi) mutable {
            covariance.adj().array() += vi.adj() * K_adj_arena.array();
          });
      grad(Z.vi_);
      auto covar_args_filter
          = internal::filter_var_scalar_types(covar_args_copy);
      internal::collect_adjoints(covar_args_adj, covar_args_filter);
    }
    if constexpr (ll_args_contain_var) {
      laplace_likelihood::ll_arg_grad(ll_fun, md_est.theta, ll_args_copy, msgs);
      internal::collect_adjoints<ZeroOut>(partial_parm, ll_args_filter);
      arena_t<Eigen::VectorXd> v;
      if (md_est.solver_used == 1 || md_est.solver_used == 2) {
        v = covariance_val * s2 - covariance_val * R * covariance_val * s2;
      } else {
        v = LU_solve_covariance * s2;
      }
      laplace_likelihood::diff_eta_implicit(ll_fun, v, md_est.theta,
                                            ll_args_copy, msgs);
      internal::collect_adjoints<ZeroOut>(partial_parm, ll_args_filter);
    }
    lmd = md_est.lmd;
  }
  var ret(lmd);
  if constexpr (is_any_var_scalar_v<CovarArgs>) {
    auto covar_args_filter = internal::filter_var_scalar_types(covar_args_refs);
    internal::reverse_pass_collect_adjoints(ret, covar_args_filter,
                                            std::move(covar_args_adj));
  }
  if constexpr (ll_args_contain_var) {
    auto ll_args_filter = internal::filter_var_scalar_types(ll_args_refs);
    internal::reverse_pass_collect_adjoints(ret, ll_args_filter,
                                            std::move(partial_parm));
  }
  return ret;
}
 
}  // namespace math
}  // namespace stan
 
#endif