Improved numerical stability of binomial_coefficient_log

stan/math/prim/fun/binomial_coefficient_log.hpp

@@ -2,9 +2,14 @@
 #define STAN_MATH_PRIM_FUN_BINOMIAL_COEFFICIENT_LOG_HPP
 
 #include <stan/math/prim/meta.hpp>
-#include <stan/math/prim/fun/inv.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/multiply_log.hpp>
+#include <stan/math/prim/fun/value_of.hpp>
 
 namespace stan {
 namespace math {
@@ -13,22 +18,24 @@ namespace math {
  * Return the log of the binomial coefficient for the specified
  * arguments.
  *
- * The binomial coefficient, \f${N \choose n}\f$, read "N choose n", is
- * defined for \f$0 \leq n \leq N\f$ by
+ * The binomial coefficient, \f${n \choose k}\f$, read "n choose k", is
+ * defined for \f$0 \leq k \leq n\f$ by
  *
- * \f${N \choose n} = \frac{N!}{n! (N-n)!}\f$.
+ * \f${n \choose k} = \frac{n!}{k! (n-k)!}\f$.
  *
  * This function uses Gamma functions to define the log
- * and generalize the arguments to continuous N and n.
+ * and generalize the arguments to continuous n and k.
+ *
+ * \f$ \log {n \choose k}
+ * = \log \ \Gamma(n+1) - \log \Gamma(k+1) - \log \Gamma(n-k+1)\f$.
  *
- * \f$ \log {N \choose n}
- * = \log \ \Gamma(N+1) - \log \Gamma(n+1) - \log \Gamma(N-n+1)\f$.
  *
    \f[
    \mbox{binomial\_coefficient\_log}(x, y) =
    \begin{cases}
-     \textrm{error} & \mbox{if } y > x \textrm{ or } y < 0\\
-     \ln\Gamma(x+1) & \mbox{if } 0\leq y \leq x \\
+     \textrm{error} & \mbox{if } y > x + 1 \textrm{ or } y < -1 \textrm{ or } x
+ < -1\\
+     \ln\Gamma(x+1) & \mbox{if } -1 < y < x + 1 \\
      \quad -\ln\Gamma(y+1)& \\
      \quad -\ln\Gamma(x-y+1)& \\[6pt]
      \textrm{NaN} & \mbox{if } x = \textrm{NaN or } y = \textrm{NaN}
@@ -38,7 +45,8 @@ namespace math {
    \f[
    \frac{\partial\, \mbox{binomial\_coefficient\_log}(x, y)}{\partial x} =
    \begin{cases}
-     \textrm{error} & \mbox{if } y > x \textrm{ or } y < 0\\
+     \textrm{error} & \mbox{if } y > x + 1 \textrm{ or } y < -1 \textrm{ or } x
+ < -1\\
      \Psi(x+1) & \mbox{if } 0\leq y \leq x \\
      \quad -\Psi(x-y+1)& \\[6pt]
      \textrm{NaN} & \mbox{if } x = \textrm{NaN or } y = \textrm{NaN}
@@ -48,32 +56,95 @@ namespace math {
    \f[
    \frac{\partial\, \mbox{binomial\_coefficient\_log}(x, y)}{\partial y} =
    \begin{cases}
-     \textrm{error} & \mbox{if } y > x \textrm{ or } y < 0\\
+     \textrm{error} & \mbox{if } y > x + 1 \textrm{ or } y < -1 \textrm{ or } x
+ < -1\\
      -\Psi(y+1) & \mbox{if } 0\leq y \leq x \\
      \quad +\Psi(x-y+1)& \\[6pt]
      \textrm{NaN} & \mbox{if } x = \textrm{NaN or } y = \textrm{NaN}
    \end{cases}
    \f]
  *
- * @tparam T_N type of the first argument
- * @tparam T_n type of the second argument
- * @param N total number of objects.
- * @param n number of objects chosen.
- * @return log (N choose n).
+ *  This function is numerically more stable than naive evaluation via lgamma.
+ *
+ * @tparam T_n type of the first argument
+ * @tparam T_k type of the second argument
+ *
+ * @param n total number of objects.
+ * @param k number of objects chosen.
+ * @return log (n choose k).
  */
-template <typename T_N, typename T_n>
-inline return_type_t<T_N, T_n> binomial_coefficient_log(const T_N N,
-                                                        const T_n n) {
-  const double CUTOFF = 1000;
-  if (N - n < CUTOFF) {
-    const T_N N_plus_1 = N + 1;
-    return lgamma(N_plus_1) - lgamma(n + 1) - lgamma(N_plus_1 - n);
+
+template <typename T_n, typename T_k>
+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 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);
+
+  operands_and_partials<T_n, T_k> ops_partials(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 {
-    return_type_t<T_N, T_n> N_minus_n = N - n;
-    const double one_twelfth = inv(12);
-    return multiply_log(n, N_minus_n) + multiply_log((N + 0.5), N / N_minus_n)
-           + one_twelfth / N - n - one_twelfth / N_minus_n - lgamma(n + 1);
+    value = -lbeta(n_plus_1_mk, k_dbl + 1) - log1p(n_dbl);
   }
+
+  if (!is_constant_all<T_n, T_k>::value) {
+    // 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 (!is_constant_all<T_n>::value) {
+      if (n_dbl == -1.0) {
+        if (k_dbl == 0) {
+          ops_partials.edge1_.partials_[0] = 0;
+        } else {
+          ops_partials.edge1_.partials_[0] = NEGATIVE_INFTY;
+        }
+      } else {
+        ops_partials.edge1_.partials_[0]
+            = (digamma(n_plus_1) - digamma_n_plus_1_mk);
+      }
+    }
+    if (!is_constant_all<T_k>::value) {
+      if (k_dbl == 0 && n_dbl == -1.0) {
+        ops_partials.edge2_.partials_[0] = NEGATIVE_INFTY;
+      } else if (k_dbl == -1) {
+        ops_partials.edge2_.partials_[0] = INFTY;
+      } else {
+        ops_partials.edge2_.partials_[0]
+            = (digamma_n_plus_1_mk - digamma(k_dbl + 1));
+      }
+    }
+  }
+
+  return ops_partials.build(value);
 }
 
 }  // namespace math

test/unit/math/mix/fun/binomial_coefficient_log_test.cpp

@@ -6,5 +6,9 @@ TEST(mathMixScalFun, binomialCoefficientLog) {
   };
   stan::test::expect_ad(f, 3, 2);
   stan::test::expect_ad(f, 24.0, 12.0);
+  stan::test::expect_ad(f, 1.0, 0.0);
+  stan::test::expect_ad(f, 0.0, 1.0);
+  stan::test::expect_ad(f, -0.3, 0.5);
+
   stan::test::expect_common_nonzero_binary(f);
 }

test/unit/math/prim/fun/binomial_coefficient_log_test.cpp

@@ -1,13 +1,14 @@
 #include <stan/math/prim.hpp>
+#include <test/unit/math/expect_near_rel.hpp>
 #include <gtest/gtest.h>
 #include <cmath>
-#include <limits>
 
 template <typename T_N, typename T_n>
 void test_binom_coefficient(const T_N& N, const T_n& n) {
   using stan::math::binomial_coefficient_log;
   EXPECT_FLOAT_EQ(lgamma(N + 1) - lgamma(n + 1) - lgamma(N - n + 1),
-                  binomial_coefficient_log(N, n));
+                  binomial_coefficient_log(N, n))
+      << "N = " << N << ", n = " << n;
 }
 
 TEST(MathFunctions, binomial_coefficient_log) {
@@ -19,6 +20,13 @@ TEST(MathFunctions, binomial_coefficient_log) {
 
   EXPECT_FLOAT_EQ(29979.16, binomial_coefficient_log(100000, 91116));
 
+  EXPECT_EQ(binomial_coefficient_log(-1, 0), 0);  // Needed for neg_binomial_2
+  EXPECT_EQ(binomial_coefficient_log(50, 0), 0);
+  EXPECT_EQ(binomial_coefficient_log(10000, 0), 0);
+
+  EXPECT_EQ(binomial_coefficient_log(10, 11), stan::math::NEGATIVE_INFTY);
+  EXPECT_EQ(binomial_coefficient_log(10, -1), stan::math::NEGATIVE_INFTY);
+
   for (int n = 0; n < 1010; ++n) {
     test_binom_coefficient(1010, n);
     test_binom_coefficient(1010.0, n);
@@ -32,9 +40,27 @@ TEST(MathFunctions, binomial_coefficient_log) {
 }
 
 TEST(MathFunctions, binomial_coefficient_log_nan) {
-  double nan = std::numeric_limits<double>::quiet_NaN();
+  double nan = stan::math::NOT_A_NUMBER;
 
   EXPECT_TRUE(std::isnan(stan::math::binomial_coefficient_log(2.0, nan)));
   EXPECT_TRUE(std::isnan(stan::math::binomial_coefficient_log(nan, 2.0)));
   EXPECT_TRUE(std::isnan(stan::math::binomial_coefficient_log(nan, nan)));
 }
+
+TEST(MathFunctions, binomial_coefficient_log_errors_edge_cases) {
+  using stan::math::INFTY;
+  using stan::math::binomial_coefficient_log;
+
+  EXPECT_NO_THROW(binomial_coefficient_log(10, 11));
+  EXPECT_THROW(binomial_coefficient_log(10, 11.01), std::domain_error);
+  EXPECT_THROW(binomial_coefficient_log(10, -1.1), std::domain_error);
+  EXPECT_THROW(binomial_coefficient_log(-1, 0.3), std::domain_error);
+  EXPECT_NO_THROW(binomial_coefficient_log(-0.5, 0.49));
+  EXPECT_NO_THROW(binomial_coefficient_log(10, -0.9));
+
+  EXPECT_FLOAT_EQ(binomial_coefficient_log(0, -1), -INFTY);
+  EXPECT_FLOAT_EQ(binomial_coefficient_log(-1, 0), 0);
+  EXPECT_FLOAT_EQ(binomial_coefficient_log(-1, -0.3), INFTY);
+  EXPECT_FLOAT_EQ(binomial_coefficient_log(0.3, -1), -INFTY);
+  EXPECT_FLOAT_EQ(binomial_coefficient_log(5.0, 6.0), -INFTY);
+}