Evaluate the incomplete beta complement without cancellation in the lccdf and Student t cdf functions

stan/math/prim/fun/inc_beta.hpp

@@ -29,6 +29,69 @@ inline double inc_beta(double a, double b, double x) {
   return boost::math::ibeta(a, b, x, boost_policy_t<>());
 }
 
+namespace internal {
+
+/**
+ * Return the complement of the regularized incomplete beta function,
+ * 1 - I_x(a, b) = I_{1-x}(b, a), from both x and 1 - x.
+ *
+ * Boost forms 1 - x from its argument, so an argument close to 1 loses
+ * relative precision in 1 - x. This overload calls Boost with whichever
+ * of x and 1 - x is at most 1/2: ibetac(a, b, x) or ibeta(b, a, 1 - x).
+ * Both keep their relative precision when the complement is below eps.
+ *
+ * Below the mean a / (a + b), while I_x(a, b) is at most 1/2, the
+ * complement is formed as 1 - I_x(a, b), which then has full precision.
+ * This also avoids Boost's arcsine case a = b = 1/2, where ibetac
+ * evaluates asin(sqrt(1 - x)) and loses digits for small x.
+ *
+ * @param a first shape, a > 0
+ * @param b second shape, b > 0
+ * @param x argument, 0 <= x <= 1
+ * @param one_m_x 1 - x, computed by the caller without cancellation
+ * @return 1 - I_x(a, b)
+ */
+inline double inc_beta_complement(double a, double b, double x,
+                                  double one_m_x) {
+  check_not_nan("inc_beta_complement", "a", a);
+  check_not_nan("inc_beta_complement", "b", b);
+  check_not_nan("inc_beta_complement", "x", x);
+  if (x > 0.5) {
+    return boost::math::ibeta(b, a, one_m_x, boost_policy_t<>());
+  }
+  if (x < a / (a + b)) {
+    const double inc = boost::math::ibeta(a, b, x, boost_policy_t<>());
+    if (inc <= 0.5) {
+      return 1.0 - inc;
+    }
+  }
+  return boost::math::ibetac(a, b, x, boost_policy_t<>());
+}
+
+/**
+ * Return the complement of the regularized incomplete beta function for
+ * autodiff arguments, by the symmetry relation through inc_beta.
+ *
+ * @tparam T1 type of a
+ * @tparam T2 type of b
+ * @tparam T3 type of x
+ * @tparam T4 type of one_m_x
+ * @param a first shape, a > 0
+ * @param b second shape, b > 0
+ * @param x argument, 0 <= x <= 1; not used
+ * @param one_m_x 1 - x
+ * @return 1 - I_x(a, b)
+ */
+template <typename T1, typename T2, typename T3, typename T4,
+          require_any_not_arithmetic_t<T1, T2, T3, T4>* = nullptr>
+inline return_type_t<T1, T2, T4> inc_beta_complement(const T1& a, const T2& b,
+                                                     const T3& /* x */,
+                                                     const T4& one_m_x) {
+  return inc_beta(b, a, one_m_x);
+}
+
+}  // namespace internal
+
 /**
  * Enables the vectorized application of the inc_beta function, when
  *  any arguments are containers.

stan/math/prim/prob/beta_lccdf.hpp

@@ -7,6 +7,7 @@
 #include <stan/math/prim/fun/digamma.hpp>
 #include <stan/math/prim/fun/grad_reg_inc_beta.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
+#include <stan/math/prim/fun/inc_beta_ddz.hpp>
 #include <stan/math/prim/fun/inv.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
@@ -73,38 +74,71 @@ inline return_type_t<T_y, T_scale_succ, T_scale_fail> beta_lccdf(
   constexpr bool calc_beta = is_autodiff_v<T_scale_fail>;
   for (size_t n = 0; n < N; n++) {
     const auto y_dbl = y_vec.val(n);
+
+    // Explicit results for extreme values
+    // The gradients are technically ill-defined, but treated as zero
+    if (y_dbl >= 1.0) {
+      return ops_partials.build(negative_infinity());
+    }
+
     const auto alpha_dbl = alpha_vec.val(n);
     const auto beta_dbl = beta_vec.val(n);
-    const auto betafunc_dbl = beta(alpha_dbl, beta_dbl);
-    const auto Pn = 1.0 - inc_beta(alpha_dbl, beta_dbl, y_dbl);
+    // The complement 1 - I_y(a, b), evaluated from the smaller of y and
+    // 1 - y. Forming it as a difference loses every digit once I_y(a, b)
+    // rounds to 1, which happens for every complement below eps.
+    const auto Pn = internal::inc_beta_complement(alpha_dbl, beta_dbl, y_dbl,
+                                                  1.0 - y_dbl);
     const auto inv_Pn
         = is_constant_all<T_y, T_scale_succ, T_scale_fail>::value ? 0 : inv(Pn);
 
     ccdf_log += log(Pn);
 
     if constexpr (is_autodiff_v<T_y>) {
-      partials<0>(ops_partials)[n] -= pow(1 - y_dbl, beta_dbl - 1)
-                                      * pow(y_dbl, alpha_dbl - 1) * inv_Pn
-                                      / betafunc_dbl;
+      // inc_beta_ddz is the same density, evaluated without the product of
+      // two powers over a beta function, each of which can underflow on
+      // its own.
+      partials<0>(ops_partials)[n]
+          -= inc_beta_ddz(alpha_dbl, beta_dbl, y_dbl) * inv_Pn;
     }
 
+    // The shape derivatives from the same argument as the value. For
+    // y <= 1/2 they are those of I_y(a, b), which are those of -Pn; on the
+    // reflected arguments they are those of Pn itself, with beta the first
+    // shape and alpha the second.
+    [[maybe_unused]] const bool reflect = y_dbl > 0.5;
     if constexpr (calc_alpha && calc_beta) {
-      const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
-          alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl), digamma(beta_dbl),
-          digamma(alpha_dbl + beta_dbl), 0.0);
-      partials<1>(ops_partials)[n] -= g1 * inv_Pn;
-      partials<2>(ops_partials)[n] -= g2 * inv_Pn;
+      if (reflect) {
+        const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
+            beta_dbl, alpha_dbl, 1.0 - y_dbl, digamma(beta_dbl),
+            digamma(alpha_dbl), digamma(alpha_dbl + beta_dbl), 0.0);
+        partials<1>(ops_partials)[n] += g2 * inv_Pn;
+        partials<2>(ops_partials)[n] += g1 * inv_Pn;
+      } else {
+        const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
+            alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl), digamma(beta_dbl),
+            digamma(alpha_dbl + beta_dbl), 0.0);
+        partials<1>(ops_partials)[n] -= g1 * inv_Pn;
+        partials<2>(ops_partials)[n] -= g2 * inv_Pn;
+      }
     } else if constexpr (calc_alpha) {
+      const auto digamma_sum = digamma(alpha_dbl + beta_dbl);
       partials<1>(ops_partials)[n]
-          -= grad_reg_inc_beta<GRAD_RIB::G1>(alpha_dbl, beta_dbl, y_dbl,
-                                             digamma(alpha_dbl), 0.0,
-                                             digamma(alpha_dbl + beta_dbl), 0.0)
+          += (reflect ? grad_reg_inc_beta<GRAD_RIB::G2>(
+                            beta_dbl, alpha_dbl, 1.0 - y_dbl, 0.0,
+                            digamma(alpha_dbl), digamma_sum, 0.0)
+                      : -grad_reg_inc_beta<GRAD_RIB::G1>(
+                            alpha_dbl, beta_dbl, y_dbl, digamma(alpha_dbl),
+                            0.0, digamma_sum, 0.0))
              * inv_Pn;
     } else if constexpr (calc_beta) {
+      const auto digamma_sum = digamma(alpha_dbl + beta_dbl);
       partials<2>(ops_partials)[n]
-          -= grad_reg_inc_beta<GRAD_RIB::G2>(alpha_dbl, beta_dbl, y_dbl, 0.0,
-                                             digamma(beta_dbl),
-                                             digamma(alpha_dbl + beta_dbl), 0.0)
+          += (reflect ? grad_reg_inc_beta<GRAD_RIB::G1>(
+                            beta_dbl, alpha_dbl, 1.0 - y_dbl,
+                            digamma(beta_dbl), 0.0, digamma_sum, 0.0)
+                      : -grad_reg_inc_beta<GRAD_RIB::G2>(
+                            alpha_dbl, beta_dbl, y_dbl, 0.0, digamma(beta_dbl),
+                            digamma_sum, 0.0))
              * inv_Pn;
     }
   }

stan/math/prim/prob/beta_proportion_lccdf.hpp

@@ -7,6 +7,7 @@
 #include <stan/math/prim/fun/beta.hpp>
 #include <stan/math/prim/fun/grad_reg_inc_beta.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
+#include <stan/math/prim/fun/inc_beta_ddz.hpp>
 #include <stan/math/prim/fun/inv.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
@@ -75,36 +76,64 @@ inline return_type_t<T_y, T_loc, T_prec> beta_proportion_lccdf(
 
   for (size_t n = 0; n < N; n++) {
     const T_partials_return y_dbl = y_vec.val(n);
+
+    // Explicit results for extreme values
+    // The gradients are technically ill-defined, but treated as zero
+    if (y_dbl >= 1.0) {
+      return ops_partials.build(negative_infinity());
+    }
+
     const T_partials_return mu_dbl = mu_vec.val(n);
     const T_partials_return kappa_dbl = kappa_vec.val(n);
     const T_partials_return mukappa_dbl = mu_dbl * kappa_dbl;
     const T_partials_return kappa_mukappa_dbl = kappa_dbl - mukappa_dbl;
-    const T_partials_return betafunc_dbl = beta(mukappa_dbl, kappa_mukappa_dbl);
-    const T_partials_return Pn
-        = 1 - inc_beta(mukappa_dbl, kappa_mukappa_dbl, y_dbl);
+    // The complement 1 - I_y(a, b), evaluated from the smaller of y and
+    // 1 - y. Forming it as a difference loses every digit once I_y(a, b)
+    // rounds to 1, which happens for every complement below eps.
+    const T_partials_return Pn = internal::inc_beta_complement(
+        mukappa_dbl, kappa_mukappa_dbl, y_dbl, 1.0 - y_dbl);
 
     ccdf_log += log(Pn);
 
     const T_partials_return inv_Pn
         = is_constant_all<T_y, T_loc, T_prec>::value ? 0 : inv(Pn);
 
     if constexpr (is_autodiff_v<T_y>) {
-      partials<0>(ops_partials)[n] -= pow(1 - y_dbl, kappa_mukappa_dbl - 1)
-                                      * pow(y_dbl, mukappa_dbl - 1) * inv_Pn
-                                      / betafunc_dbl;
+      // inc_beta_ddz is the same density, evaluated without the product of
+      // two powers over a beta function, each of which can underflow on
+      // its own.
+      partials<0>(ops_partials)[n]
+          -= inc_beta_ddz(mukappa_dbl, kappa_mukappa_dbl, y_dbl) * inv_Pn;
     }
 
     // the mu and kappa partials each need both shape partials
     if constexpr (is_any_autodiff_v<T_loc, T_prec>) {
-      const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
-          mukappa_dbl, kappa_mukappa_dbl, y_dbl, digamma(mukappa_dbl),
-          digamma(kappa_mukappa_dbl), digamma(kappa_dbl), 0.0);
+      // The derivatives with respect to a = mu kappa and b = kappa - mu kappa,
+      // from the same argument as the value
+      T_partials_return dPn_da = 0;
+      T_partials_return dPn_db = 0;
+      if (y_dbl <= 0.5) {
+        // the derivatives of I_y(a, b), which are those of -Pn
+        const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
+            mukappa_dbl, kappa_mukappa_dbl, y_dbl, digamma(mukappa_dbl),
+            digamma(kappa_mukappa_dbl), digamma(kappa_dbl), 0.0);
+        dPn_da = -g1;
+        dPn_db = -g2;
+      } else {
+        // on the reflected arguments, the derivatives of Pn itself
+        const auto [g1, g2] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
+            kappa_mukappa_dbl, mukappa_dbl, 1 - y_dbl,
+            digamma(kappa_mukappa_dbl), digamma(mukappa_dbl),
+            digamma(kappa_dbl), 0.0);
+        dPn_da = g2;
+        dPn_db = g1;
+      }
       if constexpr (is_autodiff_v<T_loc>) {
-        partials<1>(ops_partials)[n] -= kappa_dbl * (g1 - g2) * inv_Pn;
+        partials<1>(ops_partials)[n] += kappa_dbl * (dPn_da - dPn_db) * inv_Pn;
       }
       if constexpr (is_autodiff_v<T_prec>) {
         partials<2>(ops_partials)[n]
-            -= (g1 * mu_dbl + g2 * (1 - mu_dbl)) * inv_Pn;
+            += (dPn_da * mu_dbl + dPn_db * (1 - mu_dbl)) * inv_Pn;
       }
     }
   }

stan/math/prim/prob/binomial_lccdf.hpp

@@ -6,6 +6,7 @@
 #include <stan/math/prim/fun/beta.hpp>
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
+#include <stan/math/prim/fun/inc_beta_ddz.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
 #include <stan/math/prim/fun/scalar_seq_view.hpp>
@@ -89,16 +90,20 @@ inline return_type_t<T_prob> binomial_lccdf(const T_n& n, const T_N& N,
     }
 
     const T_partials_return theta_dbl = theta_vec.val(i);
-    const T_partials_return Pi
-        = 1.0 - inc_beta(N_dbl - n_dbl, n_dbl + 1, 1 - theta_dbl);
+    // The complement by the symmetry relation I_z(a, b) = I_{1-z}(b, a),
+    // which also removes the 1 - theta argument. Forming it as
+    // 1 - I_{1-theta}(N - n, n + 1) loses every digit once that rounds to
+    // 1, which happens for every complement below eps.
+    const T_partials_return Pi = inc_beta(n_dbl + 1, N_dbl - n_dbl, theta_dbl);
 
     P += log(Pi);
 
     if constexpr (is_autodiff_v<T_prob>) {
-      const T_partials_return denom = beta(N_dbl - n_dbl, n_dbl + 1) * Pi;
-      partials<0>(ops_partials)[i] += pow(theta_dbl, n_dbl)
-                                      * pow(1 - theta_dbl, N_dbl - n_dbl - 1)
-                                      / denom;
+      // inc_beta_ddz is the same density, evaluated without the product of
+      // two powers over a beta function, each of which can underflow on
+      // its own.
+      partials<0>(ops_partials)[i]
+          += inc_beta_ddz(n_dbl + 1, N_dbl - n_dbl, theta_dbl) / Pi;
     }
   }
 

stan/math/prim/prob/neg_binomial_lccdf.hpp

@@ -8,6 +8,7 @@
 #include <stan/math/prim/fun/digamma.hpp>
 #include <stan/math/prim/fun/grad_reg_inc_beta.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
+#include <stan/math/prim/fun/inc_beta_ddz.hpp>
 #include <stan/math/prim/fun/inv.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
@@ -76,23 +77,47 @@ inline return_type_t<T_shape, T_inv_scale> neg_binomial_lccdf(
     const auto inv_beta_p1 = inv(beta_dbl + 1);
     const auto p_dbl = beta_dbl * inv_beta_p1;
     const auto d_dbl = square(inv_beta_p1);
-    const auto Pi = 1.0 - inc_beta(alpha_dbl, n_dbl + 1.0, p_dbl);
-    const auto beta_func = beta(n_dbl + 1, alpha_dbl);
+    // The complement 1 - I_p(alpha, n + 1), evaluated from the smaller of
+    // p and 1 - p = 1 / (beta + 1), both formed without cancellation.
+    // Forming it as a difference loses every digit once I_p(alpha, n + 1)
+    // rounds to 1, which happens for every complement below eps.
+    const auto Pi = internal::inc_beta_complement(alpha_dbl, n_dbl + 1.0, p_dbl,
+                                                  inv_beta_p1);
 
     P += log(Pi);
 
-    // alpha is the first shape of inc_beta, so only its partial is needed
+    // Only the alpha partial is needed, from the same argument as the value
     if constexpr (is_autodiff_v<T_shape>) {
-      partials<0>(ops_partials)[i]
-          -= grad_reg_inc_beta<GRAD_RIB::G1>(
-                 alpha_dbl, n_dbl + 1, p_dbl, digamma(alpha_dbl), 0.0,
-                 digamma(n_dbl + alpha_dbl + 1), 0.0)
-             / Pi;
+      const auto digamma_alpha = digamma(alpha_dbl);
+      const auto digamma_sum = digamma(n_dbl + alpha_dbl + 1);
+      if (p_dbl <= 0.5) {
+        // alpha is the first shape of I_p(alpha, n + 1), whose derivative
+        // is that of -Pi
+        partials<0>(ops_partials)[i]
+            -= grad_reg_inc_beta<GRAD_RIB::G1>(alpha_dbl, n_dbl + 1, p_dbl,
+                                               digamma_alpha, 0.0, digamma_sum,
+                                               0.0)
+               / Pi;
+      } else {
+        // on the reflected arguments alpha is the second shape, and the
+        // derivative is that of Pi itself
+        partials<0>(ops_partials)[i]
+            += grad_reg_inc_beta<GRAD_RIB::G2>(n_dbl + 1, alpha_dbl,
+                                               inv_beta_p1, 0.0, digamma_alpha,
+                                               digamma_sum, 0.0)
+               / Pi;
+      }
     }
     if constexpr (is_autodiff_v<T_inv_scale>) {
-      partials<1>(ops_partials)[i] -= d_dbl * pow(1 - p_dbl, n_dbl)
-                                      * pow(p_dbl, alpha_dbl - 1)
-                                      / (beta_func * Pi);
+      // inc_beta_ddz is the same density, evaluated without the product of
+      // two powers over a beta function, each of which can underflow on
+      // its own. The density at p equals the reflected density at 1 - p;
+      // it is evaluated at whichever of the two is at most 1/2, because
+      // inc_beta_ddz forms the complement of its argument internally.
+      const T_partials_return density
+          = p_dbl <= 0.5 ? inc_beta_ddz(alpha_dbl, n_dbl + 1.0, p_dbl)
+                         : inc_beta_ddz(n_dbl + 1.0, alpha_dbl, inv_beta_p1);
+      partials<1>(ops_partials)[i] -= d_dbl * density / Pi;
     }
   }