binomial_coefficient_log.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUN_BINOMIAL_COEFFICIENT_LOG_HPP
#define STAN_MATH_PRIM_FUN_BINOMIAL_COEFFICIENT_LOG_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/digamma.hpp>
#include <stan/math/prim/fun/is_any_nan.hpp>
#include <stan/math/prim/fun/log1p.hpp>
#include <stan/math/prim/fun/lbeta.hpp>
#include <stan/math/prim/fun/lgamma.hpp>
#include <stan/math/prim/fun/value_of.hpp>
#include <stan/math/prim/functor/partials_propagator.hpp>
#include <stan/math/prim/functor/apply_scalar_binary.hpp>
 
namespace stan {
namespace math {
 
template <typename T_n, typename T_k,
          require_all_stan_scalar_t<T_n, T_k>* = nullptr>
inline return_type_t<T_n, T_k> binomial_coefficient_log(const T_n n,
                                                        const T_k k) {
  using T_partials_return = partials_return_t<T_n, T_k>;
 
  if (is_any_nan(n, k)) {
    return NOT_A_NUMBER;
  }
 
  // Choosing the more stable of the symmetric branches
  if (n > -1 && k > value_of_rec(n) / 2.0 + 1e-8) {
    return binomial_coefficient_log(n, n - k);
  }
 
  const T_partials_return n_dbl = value_of(n);
  const T_partials_return k_dbl = value_of(k);
  const T_partials_return n_plus_1 = n_dbl + 1;
  const T_partials_return n_plus_1_mk = n_plus_1 - k_dbl;
 
  static constexpr const char* function = "binomial_coefficient_log";
  check_greater_or_equal(function, "first argument", n, -1);
  check_greater_or_equal(function, "second argument", k, -1);
  check_greater_or_equal(function, "(first argument - second argument + 1)",
                         n_plus_1_mk, 0.0);
 
  auto ops_partials = make_partials_propagator(n, k);
 
  T_partials_return value;
  if (k_dbl == 0) {
    value = 0;
  } else if (n_plus_1 < lgamma_stirling_diff_useful) {
    value = lgamma(n_plus_1) - lgamma(k_dbl + 1) - lgamma(n_plus_1_mk);
  } else {
    value = -lbeta(n_plus_1_mk, k_dbl + 1) - log1p(n_dbl);
  }
 
  if constexpr (is_any_autodiff_v<T_n, T_k>) {
    // Branching on all the edge cases.
    // In direct computation many of those would be NaN
    // But one-sided limits from within the domain exist, all of the below
    // follows from lim x->0 from above digamma(x) == -Inf
    //
    // Note that we have k < n / 2 (see the first branch in this function)
    // se we can ignore the n == k - 1 edge case.
    T_partials_return digamma_n_plus_1_mk = digamma(n_plus_1_mk);
 
    if constexpr (is_autodiff_v<T_n>) {
      if (n_dbl == -1.0) {
        if (k_dbl == 0) {
          partials<0>(ops_partials)[0] = 0;
        } else {
          partials<0>(ops_partials)[0] = NEGATIVE_INFTY;
        }
      } else {
        partials<0>(ops_partials)[0]
            = (digamma(n_plus_1) - digamma_n_plus_1_mk);
      }
    }
    if constexpr (is_autodiff_v<T_k>) {
      if (k_dbl == 0 && n_dbl == -1.0) {
        partials<1>(ops_partials)[0] = NEGATIVE_INFTY;
      } else if (k_dbl == -1) {
        partials<1>(ops_partials)[0] = INFTY;
      } else {
        partials<1>(ops_partials)[0]
            = (digamma_n_plus_1_mk - digamma(k_dbl + 1));
      }
    }
  }
 
  return ops_partials.build(value);
}
 
template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr>
inline auto binomial_coefficient_log(T1&& a, T2&& b) {
  return apply_scalar_binary(
      [](auto&& c, auto&& d) {
        return binomial_coefficient_log(std::forward<decltype(c)>(c),
                                        std::forward<decltype(d)>(d));
      },
      std::forward<T1>(a), std::forward<T2>(b));
}
 
}  // namespace math
}  // namespace stan
#endif