dirichlet_multinomial_lpmf.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_PROB_DIRICHLET_MULTINOMIAL_LPMF_HPP
#define STAN_MATH_PRIM_PROB_DIRICHLET_MULTINOMIAL_LPMF_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/lbeta.hpp>
#include <stan/math/prim/fun/to_ref.hpp>
#include <stan/math/prim/fun/digamma.hpp>
#include <stan/math/prim/fun/log.hpp>
#include <stan/math/prim/fun/sum.hpp>
#include <stan/math/prim/fun/as_value_array_or_scalar.hpp>
#include <stan/math/prim/functor/partials_propagator.hpp>
#include <vector>
 
namespace stan {
namespace math {
 
template <bool propto, typename T_prior_size,
          require_eigen_col_vector_t<T_prior_size>* = nullptr>
return_type_t<T_prior_size> dirichlet_multinomial_lpmf(
    const std::vector<int>& ns, const T_prior_size& alpha) {
  static const char* function = "dirichlet_multinomial_lpmf";
  check_size_match(function, "Size of number of trials variable", ns.size(),
                   "rows of prior size parameter", alpha.rows());
  check_nonnegative(function, "Number of trials variable", ns);
  const auto& alpha_ref = to_ref(alpha);
  const auto& alpha_val = as_value_array_or_scalar(alpha_ref);
  check_positive_finite(function, "Prior size parameter", alpha_ref);
 
  if (!include_summand<propto, T_prior_size>::value) {
    return 0.0;
  }
 
  int n_sum = sum(ns);
  if (n_sum == 0) {
    return 0.0;
  }
 
  // Need to treat as double otherwise Eigen's array log truncates
  auto ns_array = as_array_or_scalar(ns).template cast<double>();
  partials_return_t<T_prior_size> a_sum = sum(alpha_val);
  partials_return_t<T_prior_size> lp(0.0);
  if (include_summand<propto>::value) {
    lp += log(n_sum) - (ns_array > 0).select(log(ns_array), 0.0).sum();
  }
 
  lp += lbeta(a_sum, n_sum)
        - (ns_array > 0).select(lbeta(alpha_val, ns_array), 0.0).sum();
 
  auto ops_partials = make_partials_propagator(alpha_ref);
  if (!is_constant_all<T_prior_size>::value) {
    partials<0>(ops_partials)
        = (ns_array > 0)
              .select(digamma(alpha_val + ns_array) - digamma(alpha_val), 0.0)
          + digamma(a_sum) - digamma(a_sum + n_sum);
  }
  return ops_partials.build(lp);
}
 
template <typename T_prior_size>
return_type_t<T_prior_size> dirichlet_multinomial_lpmf(
    const std::vector<int>& ns, const T_prior_size& alpha) {
  return dirichlet_multinomial_lpmf<false>(ns, alpha);
}
 
}  // namespace math
}  // namespace stan
#endif