Merge pull request #3275 from stan-dev/fix-gamma-lccdf-v3

syclik · web-flow · commit ac8c21a5e902 · 2026-04-16T09:53:00.000-04:00
Fix gamma lccdf
diff --git a/stan/math/fwd/meta/is_fvar.hpp b/stan/math/fwd/meta/is_fvar.hpp
@@ -21,8 +21,5 @@ struct is_fvar<T,
                std::enable_if_t<internal::is_fvar_impl<std::decay_t<T>>::value>>
     : std::true_type {};
 
-template <typename T>
-inline constexpr bool is_fvar_v = is_fvar<T>::value;
-
 }  // namespace stan
 #endif
diff --git a/stan/math/prim/fun/log_gamma_q_dgamma.hpp b/stan/math/prim/fun/log_gamma_q_dgamma.hpp
@@ -0,0 +1,134 @@
+#ifndef STAN_MATH_PRIM_FUN_LOG_GAMMA_Q_DGAMMA_HPP
+#define STAN_MATH_PRIM_FUN_LOG_GAMMA_Q_DGAMMA_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/digamma.hpp>
+#include <stan/math/prim/fun/exp.hpp>
+#include <stan/math/prim/fun/fabs.hpp>
+#include <stan/math/prim/fun/gamma_p.hpp>
+#include <stan/math/prim/fun/gamma_q.hpp>
+#include <stan/math/prim/fun/grad_reg_inc_gamma.hpp>
+#include <stan/math/prim/fun/inv.hpp>
+#include <stan/math/prim/fun/lgamma.hpp>
+#include <stan/math/prim/fun/log.hpp>
+#include <stan/math/prim/fun/log1m.hpp>
+#include <stan/math/prim/fun/tgamma.hpp>
+#include <stan/math/prim/fun/value_of.hpp>
+#include <stan/math/prim/fun/value_of_rec.hpp>
+#include <cmath>
+
+namespace stan {
+namespace math {
+
+namespace internal {
+
+constexpr double LOG_Q_GAMMA_CF_PRECISION = 1.49012e-12;
+
+/**
+ * Compute log(Q(a,z)) using continued fraction expansion for upper incomplete
+ * gamma function.
+ *
+ * @tparam T_a Type of shape parameter a (double or fvar types)
+ * @tparam T_z Type of value parameter z (double or fvar types)
+ * @param a Shape parameter
+ * @param z Value at which to evaluate
+ * @param precision Convergence threshold, default of sqrt(machine_epsilon)
+ * @param max_steps Maximum number of continued fraction iterations
+ * @return log(Q(a,z)) with the return type of T_a and T_z
+ */
+template <typename T_a, typename T_z>
+inline return_type_t<T_a, T_z> log_q_gamma_cf(const T_a& a, const T_z& z,
+                                              double precision
+                                              = LOG_Q_GAMMA_CF_PRECISION,
+                                              int max_steps = 250) {
+  using T_return = return_type_t<T_a, T_z>;
+  const T_return log_prefactor = a * log(z) - z - lgamma(a);
+
+  T_return b_init = z + 1.0 - a;
+  T_return C = (fabs(value_of_rec(b_init)) >= EPSILON)
+                   ? b_init
+                   : std::decay_t<decltype(b_init)>(EPSILON);
+  T_return D = 0.0;
+  T_return f = C;
+  for (int i = 1; i <= max_steps; ++i) {
+    T_a an = -i * (i - a);
+    const T_return b = b_init + 2.0 * i;
+    D = b + an * D;
+    D = (fabs(value_of_rec(D)) >= EPSILON) ? D
+                                           : std::decay_t<decltype(D)>(EPSILON);
+    C = b + an / C;
+    C = (fabs(value_of_rec(C)) >= EPSILON) ? C
+                                           : std::decay_t<decltype(C)>(EPSILON);
+    D = inv(D);
+    const T_return delta = C * D;
+    f *= delta;
+    const double delta_m1 = fabs(value_of_rec(delta) - 1.0);
+    if (delta_m1 < precision) {
+      break;
+    }
+  }
+  return log_prefactor - log(f);
+}
+
+}  // namespace internal
+
+/**
+ * Compute log(Q(a,z)) and its gradient with respect to a using continued
+ * fraction expansion, where Q(a,z) = Gamma(a,z) / Gamma(a) is the regularized
+ * upper incomplete gamma function.
+ *
+ * This uses a continued fraction representation for numerical stability when
+ * computing the upper incomplete gamma function in log space, along with
+ * analytical gradient computation.
+ *
+ * @tparam T_a type of the shape parameter
+ * @tparam T_z type of the value parameter
+ * @param a shape parameter (must be positive)
+ * @param z value parameter (must be non-negative)
+ * @param precision convergence threshold, default of sqrt(machine_epsilon)
+ * @param max_steps maximum iterations for continued fraction
+ * @return structure containing log(Q(a,z)) and d/da log(Q(a,z))
+ */
+template <typename T_a, typename T_z>
+inline std::pair<return_type_t<T_a, T_z>, return_type_t<T_a, T_z>>
+log_gamma_q_dgamma(const T_a& a, const T_z& z,
+                   double precision = internal::LOG_Q_GAMMA_CF_PRECISION,
+                   int max_steps = 250) {
+  using T_return = return_type_t<T_a, T_z>;
+  const double a_val = value_of(a);
+  const double z_val = value_of(z);
+  // For z > a + 1, use continued fraction for better numerical stability
+  if (z_val > a_val + 1.0) {
+    std::pair<T_return, T_return> result{
+        internal::log_q_gamma_cf(a_val, z_val, precision, max_steps), 0.0};
+    // For gradient, use: d/da log(Q) = (1/Q) * dQ/da
+    // grad_reg_inc_gamma computes dQ/da
+    const T_return Q_val = exp(result.first);
+    const double dQ_da
+        = grad_reg_inc_gamma(a_val, z_val, tgamma(a_val), digamma(a_val));
+    result.second = dQ_da / Q_val;
+    return result;
+  } else {
+    // For z <= a + 1, use log1m(P(a,z)) for better numerical accuracy
+    const double P_val = gamma_p(a_val, z_val);
+    std::pair<T_return, T_return> result{log1m(P_val), 0.0};
+    // Gradient: d/da log(Q) = (1/Q) * dQ/da
+    // grad_reg_inc_gamma computes dQ/da
+    const T_return Q_val = exp(result.first);
+    if (Q_val > 0) {
+      const double dQ_da
+          = grad_reg_inc_gamma(a_val, z_val, tgamma(a_val), digamma(a_val));
+      result.second = dQ_da / Q_val;
+    } else {
+      // Fallback if Q rounds to zero - use asymptotic approximation
+      result.second = log(z_val) - digamma(a_val);
+    }
+    return result;
+  }
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif
diff --git a/stan/math/prim/meta/is_fvar.hpp b/stan/math/prim/meta/is_fvar.hpp
@@ -14,6 +14,9 @@ namespace stan {
 template <typename T, typename = void>
 struct is_fvar : std::false_type {};
 
+template <typename T>
+inline constexpr bool is_fvar_v = is_fvar<T>::value;
+
 /** \ingroup type_trait
  * Specialization for pointers returns the underlying value the pointer is
  * pointing to.
diff --git a/stan/math/prim/prob/gamma_lccdf.hpp b/stan/math/prim/prob/gamma_lccdf.hpp
@@ -6,28 +6,102 @@
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/digamma.hpp>
 #include <stan/math/prim/fun/exp.hpp>
-#include <stan/math/prim/fun/gamma_q.hpp>
+#include <stan/math/prim/fun/fma.hpp>
+#include <stan/math/prim/fun/gamma_p.hpp>
 #include <stan/math/prim/fun/grad_reg_inc_gamma.hpp>
+#include <stan/math/prim/fun/grad_reg_lower_inc_gamma.hpp>
+#include <stan/math/prim/fun/lgamma.hpp>
 #include <stan/math/prim/fun/log.hpp>
+#include <stan/math/prim/fun/log1m.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
 #include <stan/math/prim/fun/scalar_seq_view.hpp>
 #include <stan/math/prim/fun/size.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/fun/tgamma.hpp>
-#include <stan/math/prim/fun/value_of.hpp>
+#include <stan/math/prim/fun/value_of_rec.hpp>
+#include <stan/math/prim/fun/log_gamma_q_dgamma.hpp>
 #include <stan/math/prim/functor/partials_propagator.hpp>
 #include <cmath>
+#include <optional>
 
 namespace stan {
 namespace math {
+namespace internal {
+
+/**
+ * Computes log q and d(log q) / d(alpha) using continued fraction.
+ */
+template <bool any_fvar, bool partials_fvar, typename T_shape, typename T1,
+          typename T2>
+inline std::optional<std::pair<return_type_t<T1, T2>, return_type_t<T1, T2>>>
+eval_q_cf(const T1& alpha, const T2& beta_y) {
+  using scalar_t = return_type_t<T1, T2>;
+  using ret_t = std::pair<scalar_t, scalar_t>;
+  if constexpr (!any_fvar && is_autodiff_v<T_shape>) {
+    std::pair<double, double> log_q_result
+        = log_gamma_q_dgamma(value_of(alpha), value_of(beta_y));
+    if (likely(std::isfinite(log_q_result.first))) {
+      return std::optional{log_q_result};
+    } else {
+      return std::optional<ret_t>{std::nullopt};
+    }
+  } else {
+    ret_t out{internal::log_q_gamma_cf(alpha, beta_y), 0.0};
+    if (unlikely(!std::isfinite(value_of_rec(out.first)))) {
+      return std::optional<ret_t>{std::nullopt};
+    }
+    if constexpr (is_autodiff_v<T_shape>) {
+      if constexpr (!partials_fvar) {
+        out.second
+            = grad_reg_inc_gamma(alpha, beta_y, tgamma(alpha), digamma(alpha))
+              / exp(out.first);
+      } else {
+        auto alpha_unit = alpha;
+        alpha_unit.d_ = 1;
+        auto beta_y_unit = beta_y;
+        beta_y_unit.d_ = 0;
+        auto log_Q_fvar = internal::log_q_gamma_cf(alpha_unit, beta_y_unit);
+        out.second = log_Q_fvar.d_;
+      }
+    }
+    return std::optional{out};
+  }
+}
+
+/**
+ * Computes log q and d(log q) / d(alpha) using log1m.
+ */
+template <bool partials_fvar, typename T_shape, typename T1, typename T2>
+inline std::optional<std::pair<return_type_t<T1, T2>, return_type_t<T1, T2>>>
+eval_q_log1m(const T1& alpha, const T2& beta_y) {
+  using scalar_t = return_type_t<T1, T2>;
+  using ret_t = std::pair<scalar_t, scalar_t>;
+  ret_t out{log1m(gamma_p(alpha, beta_y)), 0.0};
+  if (unlikely(!std::isfinite(value_of_rec(out.first)))) {
+    return std::optional<ret_t>{std::nullopt};
+  }
+  if constexpr (is_autodiff_v<T_shape>) {
+    if constexpr (partials_fvar) {
+      auto alpha_unit = alpha;
+      alpha_unit.d_ = 1;
+      auto beta_unit = beta_y;
+      beta_unit.d_ = 0;
+      auto log_Q_fvar = log1m(gamma_p(alpha_unit, beta_unit));
+      out.second = log_Q_fvar.d_;
+    } else {
+      out.second = -grad_reg_lower_inc_gamma(alpha, beta_y) / exp(out.first);
+    }
+  }
+  return std::optional{out};
+}
+}  // namespace internal
 
 template <typename T_y, typename T_shape, typename T_inv_scale>
 inline return_type_t<T_y, T_shape, T_inv_scale> gamma_lccdf(
     const T_y& y, const T_shape& alpha, const T_inv_scale& beta) {
-  using T_partials_return = partials_return_t<T_y, T_shape, T_inv_scale>;
   using std::exp;
   using std::log;
-  using std::pow;
+  using T_partials_return = partials_return_t<T_y, T_shape, T_inv_scale>;
   using T_y_ref = ref_type_t<T_y>;
   using T_alpha_ref = ref_type_t<T_shape>;
   using T_beta_ref = ref_type_t<T_inv_scale>;
@@ -51,61 +125,70 @@ inline return_type_t<T_y, T_shape, T_inv_scale> gamma_lccdf(
   scalar_seq_view<T_y_ref> y_vec(y_ref);
   scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
   scalar_seq_view<T_beta_ref> beta_vec(beta_ref);
-  size_t N = max_size(y, alpha, beta);
-
-  // Explicit return for extreme values
-  // The gradients are technically ill-defined, but treated as zero
-  for (size_t i = 0; i < stan::math::size(y); i++) {
-    if (y_vec.val(i) == 0) {
-      // LCCDF(0) = log(P(Y > 0)) = log(1) = 0
-      return ops_partials.build(0.0);
-    }
-  }
+  const size_t N = max_size(y, alpha, beta);
+
+  constexpr bool is_y_fvar = is_fvar_v<scalar_type_t<T_y>>;
+  constexpr bool is_shape_fvar = is_fvar_v<scalar_type_t<T_shape>>;
+  constexpr bool is_beta_fvar = is_fvar_v<scalar_type_t<T_inv_scale>>;
+  constexpr bool any_fvar = is_y_fvar || is_shape_fvar || is_beta_fvar;
+  constexpr bool partials_fvar = is_fvar_v<T_partials_return>;
 
   for (size_t n = 0; n < N; n++) {
     // Explicit results for extreme values
     // The gradients are technically ill-defined, but treated as zero
-    if (y_vec.val(n) == INFTY) {
-      // LCCDF(∞) = log(P(Y > ∞)) = log(0) = -∞
+    const T_partials_return y_val = y_vec.val(n);
+    if (y_val == 0.0) {
+      continue;
+    }
+    if (y_val == INFTY) {
       return ops_partials.build(negative_infinity());
     }
 
-    const T_partials_return y_dbl = y_vec.val(n);
-    const T_partials_return alpha_dbl = alpha_vec.val(n);
-    const T_partials_return beta_dbl = beta_vec.val(n);
-    const T_partials_return beta_y_dbl = beta_dbl * y_dbl;
+    const T_partials_return alpha_val = alpha_vec.val(n);
+    const T_partials_return beta_val = beta_vec.val(n);
 
-    // Qn = 1 - Pn
-    const T_partials_return Qn = gamma_q(alpha_dbl, beta_y_dbl);
-    const T_partials_return log_Qn = log(Qn);
+    const T_partials_return beta_y = beta_val * y_val;
+    if (beta_y == INFTY) {
+      return ops_partials.build(negative_infinity());
+    }
+    std::optional<std::pair<T_partials_return, T_partials_return>> result;
+    if (beta_y > alpha_val + 1.0) {
+      result = internal::eval_q_cf<any_fvar, partials_fvar, T_shape>(alpha_val,
+                                                                     beta_y);
+    } else {
+      result
+          = internal::eval_q_log1m<partials_fvar, T_shape>(alpha_val, beta_y);
+      if (!result && beta_y > 0.0) {
+        // Fallback to continued fraction if log1m fails
+        result = internal::eval_q_cf<any_fvar, partials_fvar, T_shape>(
+            alpha_val, beta_y);
+      }
+    }
+    if (unlikely(!result)) {
+      return ops_partials.build(negative_infinity());
+    }
 
-    P += log_Qn;
+    P += result->first;
 
-    if constexpr (is_any_autodiff_v<T_y, T_inv_scale>) {
-      const T_partials_return log_y_dbl = log(y_dbl);
-      const T_partials_return log_beta_dbl = log(beta_dbl);
-      const T_partials_return log_pdf
-          = alpha_dbl * log_beta_dbl - lgamma(alpha_dbl)
-            + (alpha_dbl - 1.0) * log_y_dbl - beta_y_dbl;
-      const T_partials_return common_term = exp(log_pdf - log_Qn);
+    if constexpr (is_autodiff_v<T_y> || is_autodiff_v<T_inv_scale>) {
+      const T_partials_return log_y = log(y_val);
+      const T_partials_return alpha_minus_one = fma(alpha_val, log_y, -log_y);
+
+      const T_partials_return log_pdf = alpha_val * log(beta_val)
+                                        - lgamma(alpha_val) + alpha_minus_one
+                                        - beta_y;
+
+      const T_partials_return hazard = exp(log_pdf - result->first);  // f/Q
 
       if constexpr (is_autodiff_v<T_y>) {
-        // d/dy log(1-F(y)) = -f(y)/(1-F(y))
-        partials<0>(ops_partials)[n] -= common_term;
+        partials<0>(ops_partials)[n] -= hazard;
       }
       if constexpr (is_autodiff_v<T_inv_scale>) {
-        // d/dbeta log(1-F(y)) = -y*f(y)/(beta*(1-F(y)))
-        partials<2>(ops_partials)[n] -= y_dbl / beta_dbl * common_term;
+        partials<2>(ops_partials)[n] -= (y_val / beta_val) * hazard;
       }
     }
-
     if constexpr (is_autodiff_v<T_shape>) {
-      const T_partials_return digamma_val = digamma(alpha_dbl);
-      const T_partials_return gamma_val = tgamma(alpha_dbl);
-      // d/dalpha log(1-F(y)) = grad_upper_inc_gamma / (1-F(y))
-      partials<1>(ops_partials)[n]
-          += grad_reg_inc_gamma(alpha_dbl, beta_y_dbl, gamma_val, digamma_val)
-             / Qn;
+      partials<1>(ops_partials)[n] += result->second;
     }
   }
   return ops_partials.build(P);
diff --git a/test/unit/math/prim/prob/gamma_lccdf_test.cpp b/test/unit/math/prim/prob/gamma_lccdf_test.cpp
diff --git a/test/unit/math/rev/prob/gamma_lccdf_test.cpp b/test/unit/math/rev/prob/gamma_lccdf_test.cpp
