refactor gamma_lccdf

Author: syclik

stan/math/prim/prob/gamma_lccdf.hpp

@@ -25,10 +25,85 @@
 
 namespace stan {
 namespace math {
+namespace internal {
+  template <typename T>
+  struct Q_eval {
+    T log_Q{0.0};
+    T dlogQ_dalpha{0.0};
+    bool ok{false};
+  };
+
+  /**
+   * Computes log q and d(log q) / d(alpha) using continued fraction.
+   */
+  template <typename T, typename T_shape,
+            bool any_fvar, bool partials_fvar>
+  static inline Q_eval<T> eval_q_cf(const T& alpha,
+                                    const T& beta_y) {
+    Q_eval<T> out;
+    if constexpr (!any_fvar && is_autodiff_v<T_shape>) {
+      auto log_q_result = log_gamma_q_dgamma(value_of_rec(alpha), value_of_rec(beta_y));
+      out.log_Q = log_q_result.log_q;
+      out.dlogQ_dalpha = log_q_result.dlog_q_da;
+    } else {
+      out.log_Q = internal::log_q_gamma_cf(alpha, beta_y);
+      if constexpr (is_autodiff_v<T_shape>) {
+        if constexpr (!partials_fvar) {
+          out.dlogQ_dalpha
+            = grad_reg_inc_gamma(alpha, beta_y, tgamma(alpha),
+                                 digamma(alpha)) / exp(out.log_Q);
+        } else {
+          T alpha_unit = alpha;
+          alpha_unit.d_ = 1;
+          T beta_y_unit = beta_y;
+          beta_y_unit.d_ = 0;
+          T log_Q_fvar = internal::log_q_gamma_cf(alpha_unit, beta_y_unit);
+          out.dlogQ_dalpha = log_Q_fvar.d_;
+        }
+      }
+    }
+
+    out.ok = std::isfinite(value_of_rec(out.log_Q));
+    return out;
+  }
+
+  /**
+   * Computes log q and d(log q) / d(alpha) using log1m.
+   */
+  template <typename T, typename T_shape,
+            bool partials_fvar>
+  static inline Q_eval<T> eval_q_log1m(const T& alpha,
+                                       const T& beta_y) {
+    Q_eval<T> out;
+    out.log_Q = log1m(gamma_p(alpha, beta_y));
+
+    if (!std::isfinite(value_of_rec(out.log_Q))) {
+      out.ok = false;
+      return out;
+    }
+
+    if constexpr (is_autodiff_v<T_shape>) {
+      if constexpr (partials_fvar) {
+        T alpha_unit = alpha;
+        alpha_unit.d_ = 1;
+        T beta_unit = beta_y;
+        beta_unit.d_ = 0;
+        T log_Q_fvar = log1m(gamma_p(alpha_unit, beta_unit));
+        out.dlogQ_dalpha = log_Q_fvar.d_;
+      } else {
+        out.dlogQ_dalpha = -grad_reg_lower_inc_gamma(alpha, beta_y) / exp(out.log_Q);
+      }
+    }
+
+    out.ok = true;
+    return out;
+  }
+}
 
 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) {
+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 std::exp;
   using std::log;
   using T_partials_return = partials_return_t<T_y, T_shape, T_inv_scale>;
@@ -81,91 +156,35 @@ inline return_type_t<T_y, T_shape, T_inv_scale> gamma_lccdf(
       return ops_partials.build(negative_infinity());
     }
 
-    bool use_cf = beta_y > alpha_dbl + 1.0;
-    T_partials_return log_Qn;
-    [[maybe_unused]] T_partials_return dlogQ_dalpha = 0.0;
-
-    // Branch by autodiff type first, then handle use_cf logic inside each path
-    if constexpr (!any_fvar && is_autodiff_v<T_shape>) {
-      // var-only path: use log_gamma_q_dgamma which computes both log_q
-      // and its gradient analytically with double inputs
-      const double beta_y_dbl = value_of_rec(beta_y);
-      const double alpha_dbl_val = value_of_rec(alpha_dbl);
-
-      if (use_cf) {
-        auto log_q_result = log_gamma_q_dgamma(alpha_dbl_val, beta_y_dbl);
-        log_Qn = log_q_result.log_q;
-        dlogQ_dalpha = log_q_result.dlog_q_da;
-      } else {
-        const T_partials_return Pn = gamma_p(alpha_dbl, beta_y);
-        log_Qn = log1m(Pn);
-        const T_partials_return Qn = exp(log_Qn);
-
-        // Check if we need to fallback to continued fraction
-        bool need_cf_fallback
-            = !std::isfinite(value_of_rec(log_Qn)) || Qn <= 0.0;
-        if (need_cf_fallback && beta_y > 0.0) {
-          auto log_q_result = log_gamma_q_dgamma(alpha_dbl_val, beta_y_dbl);
-          log_Qn = log_q_result.log_q;
-          dlogQ_dalpha = log_q_result.dlog_q_da;
-        } else {
-          dlogQ_dalpha = -grad_reg_lower_inc_gamma(alpha_dbl, beta_y) / Qn;
-        }
-      }
-    } else if constexpr (partials_fvar && is_autodiff_v<T_shape>) {
-      // fvar path: use unit derivative trick to compute gradients
-      T_partials_return alpha_unit = alpha_dbl;
-      alpha_unit.d_ = 1;
-      T_partials_return beta_unit = beta_y;
-      beta_unit.d_ = 0;
-
-      if (use_cf) {
-        log_Qn = internal::log_q_gamma_cf(alpha_dbl, beta_y);
-        T_partials_return log_Qn_fvar
-            = internal::log_q_gamma_cf(alpha_unit, beta_unit);
-        dlogQ_dalpha = log_Qn_fvar.d_;
-      } else {
-        const T_partials_return Pn = gamma_p(alpha_dbl, beta_y);
-        log_Qn = log1m(Pn);
-
-        if (!std::isfinite(value_of_rec(log_Qn)) && beta_y > 0.0) {
-          // Fallback to continued fraction
-          log_Qn = internal::log_q_gamma_cf(alpha_dbl, beta_y);
-          T_partials_return log_Qn_fvar
-              = internal::log_q_gamma_cf(alpha_unit, beta_unit);
-          dlogQ_dalpha = log_Qn_fvar.d_;
-        } else {
-          T_partials_return log_Qn_fvar = log1m(gamma_p(alpha_unit, beta_unit));
-          dlogQ_dalpha = log_Qn_fvar.d_;
-        }
-      }
+    const bool use_continued_fraction = beta_y > alpha_dbl + 1.0;
+    internal::Q_eval<T_partials_return> result;
+    if (use_continued_fraction) {
+      result = internal::eval_q_cf<T_partials_return, T_shape,
+                                any_fvar, partials_fvar>(alpha_dbl, beta_y);
     } else {
-      // No alpha derivative needed (alpha is constant or double-only)
-      if (use_cf) {
-        log_Qn = internal::log_q_gamma_cf(alpha_dbl, beta_y);
-      } else {
-        const T_partials_return Pn = gamma_p(alpha_dbl, beta_y);
-        log_Qn = log1m(Pn);
+      result = internal::eval_q_log1m<T_partials_return, T_shape,
+                                   partials_fvar>(alpha_dbl, beta_y);
 
-        if (!std::isfinite(value_of_rec(log_Qn)) && beta_y > 0.0) {
-          log_Qn = internal::log_q_gamma_cf(alpha_dbl, beta_y);
-        }
+      if (!result.ok && beta_y > 0.0) {
+        // Fallback to continued fraction if log1m fails
+        result = internal::eval_q_cf<T_partials_return, T_shape,
+                                  any_fvar, partials_fvar>(alpha_dbl, beta_y);
       }
     }
-    if (!std::isfinite(value_of_rec(log_Qn))) {
+    if (!result.ok) {
       return ops_partials.build(negative_infinity());
     }
-    P += log_Qn;
+
+    P += result.log_Q;
 
     if constexpr (is_autodiff_v<T_y> || is_autodiff_v<T_inv_scale>) {
       const T_partials_return log_y = log(y_dbl);
       const T_partials_return alpha_minus_one = fma(alpha_dbl, log_y, -log_y);
 
-      const T_partials_return log_pdf = alpha_dbl * log(beta_dbl)
-                                        - lgamma(alpha_dbl) + alpha_minus_one
-                                        - beta_y;
+      const T_partials_return log_pdf
+        = alpha_dbl * log(beta_dbl) - lgamma(alpha_dbl) + alpha_minus_one - beta_y;
 
-      const T_partials_return hazard = exp(log_pdf - log_Qn);  // f/Q
+      const T_partials_return hazard = exp(log_pdf - result.log_Q);  // f/Q
 
       if constexpr (is_autodiff_v<T_y>) {
         partials<0>(ops_partials)[n] -= hazard;
@@ -175,7 +194,7 @@ inline return_type_t<T_y, T_shape, T_inv_scale> gamma_lccdf(
       }
     }
     if constexpr (is_autodiff_v<T_shape>) {
-      partials<1>(ops_partials)[n] += dlogQ_dalpha;
+      partials<1>(ops_partials)[n] += result.dlogQ_dalpha;
     }
   }
   return ops_partials.build(P);