1 #ifndef STAN_MATH_PRIM_FUN_BINOMIAL_COEFFICIENT_LOG_HPP
2 #define STAN_MATH_PRIM_FUN_BINOMIAL_COEFFICIENT_LOG_HPP
3
4 #include <stan/math/prim/meta.hpp>
5 #include <stan/math/prim/err.hpp>
6 #include <stan/math/prim/fun/constants.hpp>
7 #include <stan/math/prim/fun/digamma.hpp>
8 #include <stan/math/prim/fun/is_any_nan.hpp>
9 #include <stan/math/prim/fun/log1p.hpp>
10 #include <stan/math/prim/fun/lbeta.hpp>
11 #include <stan/math/prim/fun/lgamma.hpp>
12 #include <stan/math/prim/fun/value_of.hpp>
13 #include <stan/math/prim/functor/operands_and_partials.hpp>
14 #include <stan/math/prim/functor/apply_scalar_binary.hpp>
15
16 namespace stan {
17 namespace math {
18
19 /**
20 * Return the log of the binomial coefficient for the specified
21 * arguments.
22 *
23 * The binomial coefficient, \f${n \choose k}\f$, read "n choose k", is
24 * defined for \f$0 \leq k \leq n\f$ by
25 *
26 * \f${n \choose k} = \frac{n!}{k! (n-k)!}\f$.
27 *
28 * This function uses Gamma functions to define the log
29 * and generalize the arguments to continuous n and k.
30 *
31 * \f$ \log {n \choose k}
32 * = \log \ \Gamma(n+1) - \log \Gamma(k+1) - \log \Gamma(n-k+1)\f$.
33 *
34 *
35 \f[
36 \mbox{binomial\_coefficient\_log}(x, y) =
37 \begin{cases}
38 \textrm{error} & \mbox{if } y > x + 1 \textrm{ or } y < -1 \textrm{ or } x
39 < -1\\
40 \ln\Gamma(x+1) & \mbox{if } -1 < y < x + 1 \\
41 \quad -\ln\Gamma(y+1)& \\
42 \quad -\ln\Gamma(x-y+1)& \\[6pt]
43 \textrm{NaN} & \mbox{if } x = \textrm{NaN or } y = \textrm{NaN}
44 \end{cases}
45 \f]
46
47 \f[
48 \frac{\partial\, \mbox{binomial\_coefficient\_log}(x, y)}{\partial x} =
49 \begin{cases}
50 \textrm{error} & \mbox{if } y > x + 1 \textrm{ or } y < -1 \textrm{ or } x
51 < -1\\
52 \Psi(x+1) & \mbox{if } 0\leq y \leq x \\
53 \quad -\Psi(x-y+1)& \\[6pt]
54 \textrm{NaN} & \mbox{if } x = \textrm{NaN or } y = \textrm{NaN}
55 \end{cases}
56 \f]
57
58 \f[
59 \frac{\partial\, \mbox{binomial\_coefficient\_log}(x, y)}{\partial y} =
60 \begin{cases}
61 \textrm{error} & \mbox{if } y > x + 1 \textrm{ or } y < -1 \textrm{ or } x
62 < -1\\
63 -\Psi(y+1) & \mbox{if } 0\leq y \leq x \\
64 \quad +\Psi(x-y+1)& \\[6pt]
65 \textrm{NaN} & \mbox{if } x = \textrm{NaN or } y = \textrm{NaN}
66 \end{cases}
67 \f]
68 *
69 * This function is numerically more stable than naive evaluation via lgamma.
70 *
71 * @tparam T_n type of the first argument
72 * @tparam T_k type of the second argument
73 *
74 * @param n total number of objects.
75 * @param k number of objects chosen.
76 * @return log (n choose k).
77 */
78
79 template <typename T_n, typename T_k,
80 require_all_stan_scalar_t<T_n, T_k>* = nullptr>
binomial_coefficient_log(const T_n n,const T_k k)81 inline return_type_t<T_n, T_k> binomial_coefficient_log(const T_n n,
82 const T_k k) {
83 using T_partials_return = partials_return_t<T_n, T_k>;
84
85 if (is_any_nan(n, k)) {
86 return NOT_A_NUMBER;
87 }
88
89 // Choosing the more stable of the symmetric branches
90 if (n > -1 && k > value_of_rec(n) / 2.0 + 1e-8) {
91 return binomial_coefficient_log(n, n - k);
92 }
93
94 const T_partials_return n_dbl = value_of(n);
95 const T_partials_return k_dbl = value_of(k);
96 const T_partials_return n_plus_1 = n_dbl + 1;
97 const T_partials_return n_plus_1_mk = n_plus_1 - k_dbl;
98
99 static const char* function = "binomial_coefficient_log";
100 check_greater_or_equal(function, "first argument", n, -1);
101 check_greater_or_equal(function, "second argument", k, -1);
102 check_greater_or_equal(function, "(first argument - second argument + 1)",
103 n_plus_1_mk, 0.0);
104
105 operands_and_partials<T_n, T_k> ops_partials(n, k);
106
107 T_partials_return value;
108 if (k_dbl == 0) {
109 value = 0;
110 } else if (n_plus_1 < lgamma_stirling_diff_useful) {
111 value = lgamma(n_plus_1) - lgamma(k_dbl + 1) - lgamma(n_plus_1_mk);
112 } else {
113 value = -lbeta(n_plus_1_mk, k_dbl + 1) - log1p(n_dbl);
114 }
115
116 if (!is_constant_all<T_n, T_k>::value) {
117 // Branching on all the edge cases.
118 // In direct computation many of those would be NaN
119 // But one-sided limits from within the domain exist, all of the below
120 // follows from lim x->0 from above digamma(x) == -Inf
121 //
122 // Note that we have k < n / 2 (see the first branch in this function)
123 // se we can ignore the n == k - 1 edge case.
124 T_partials_return digamma_n_plus_1_mk = digamma(n_plus_1_mk);
125
126 if (!is_constant_all<T_n>::value) {
127 if (n_dbl == -1.0) {
128 if (k_dbl == 0) {
129 ops_partials.edge1_.partials_[0] = 0;
130 } else {
131 ops_partials.edge1_.partials_[0] = NEGATIVE_INFTY;
132 }
133 } else {
134 ops_partials.edge1_.partials_[0]
135 = (digamma(n_plus_1) - digamma_n_plus_1_mk);
136 }
137 }
138 if (!is_constant_all<T_k>::value) {
139 if (k_dbl == 0 && n_dbl == -1.0) {
140 ops_partials.edge2_.partials_[0] = NEGATIVE_INFTY;
141 } else if (k_dbl == -1) {
142 ops_partials.edge2_.partials_[0] = INFTY;
143 } else {
144 ops_partials.edge2_.partials_[0]
145 = (digamma_n_plus_1_mk - digamma(k_dbl + 1));
146 }
147 }
148 }
149
150 return ops_partials.build(value);
151 }
152
153 /**
154 * Enables the vectorised application of the binomial coefficient log function,
155 * when the first and/or second arguments are containers.
156 *
157 * @tparam T1 type of first input
158 * @tparam T2 type of second input
159 * @param a First input
160 * @param b Second input
161 * @return Binomial coefficient log function applied to the two inputs.
162 */
163 template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr>
binomial_coefficient_log(const T1 & a,const T2 & b)164 inline auto binomial_coefficient_log(const T1& a, const T2& b) {
165 return apply_scalar_binary(a, b, [&](const auto& c, const auto& d) {
166 return binomial_coefficient_log(c, d);
167 });
168 }
169
170 } // namespace math
171 } // namespace stan
172 #endif
173