cleanup for gamma_lccdf

Author: SteveBronder

stan/math/prim/fun/log_gamma_q_dgamma.hpp

@@ -5,6 +5,7 @@
 #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>
@@ -48,18 +49,13 @@ 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 = 1e-16,
                                               int max_steps = 250) {
-  using stan::math::lgamma;
-  using stan::math::log;
-  using stan::math::value_of;
-  using std::fabs;
   using T_return = return_type_t<T_a, T_z>;
+  const auto log_prefactor = a * log(z) - z - lgamma(a);
 
-  const T_return log_prefactor = a * log(z) - z - lgamma(a);
-
-  T_return b = z + 1.0 - a;
-  T_return C = (fabs(value_of(b)) >= EPSILON) ? b : T_return(EPSILON);
-  T_return D = 0.0;
-  T_return f = C;
+  auto b = z + 1.0 - a;
+  auto C = (fabs(value_of(b)) >= EPSILON) ? b : std::decay_t<decltype(b)>(EPSILON);
+  auto D = 0.0;
+  auto f = C;
 
   for (int i = 1; i <= max_steps; ++i) {
     T_a an = -i * (i - a);
@@ -75,15 +71,14 @@ inline return_type_t<T_a, T_z> log_q_gamma_cf(const T_a& a, const T_z& z,
     }
 
     D = inv(D);
-    T_return delta = C * D;
+    auto delta = C * D;
     f *= delta;
 
     const double delta_m1 = value_of(fabs(value_of(delta) - 1.0));
     if (delta_m1 < precision) {
       break;
     }
   }
-
   return log_prefactor - log(f);
 }
 
@@ -109,45 +104,36 @@ inline return_type_t<T_a, T_z> log_q_gamma_cf(const T_a& a, const T_z& z,
 template <typename T_a, typename T_z>
 inline log_gamma_q_result<return_type_t<T_a, T_z>> log_gamma_q_dgamma(
     const T_a& a, const T_z& z, double precision = 1e-16, int max_steps = 250) {
-  using std::exp;
-  using std::log;
   using T_return = return_type_t<T_a, T_z>;
-
-  const double a_dbl = value_of(a);
-  const double z_dbl = value_of(z);
-
-  log_gamma_q_result<T_return> result;
-
+  const auto a_dbl = value_of(a);
+  const auto z_dbl = value_of(z);
   // For z > a + 1, use continued fraction for better numerical stability
   if (z_dbl > a_dbl + 1.0) {
-    result.log_q = internal::log_q_gamma_cf(a_dbl, z_dbl, precision, max_steps);
-
+    log_gamma_q_result<T_return> result{internal::log_q_gamma_cf(a_dbl, z_dbl, precision, max_steps), 0.0};
     // For gradient, use: d/da log(Q) = (1/Q) * dQ/da
     // grad_reg_inc_gamma computes dQ/da
-    const double Q_val = exp(result.log_q);
-    const double dQ_da
+    const auto Q_val = exp(result.log_q);
+    const auto dQ_da
         = grad_reg_inc_gamma(a_dbl, z_dbl, tgamma(a_dbl), digamma(a_dbl));
     result.dlog_q_da = 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_dbl, z_dbl);
-    result.log_q = log1m(P_val);
-
+    const auto P_val = gamma_p(a_dbl, z_dbl);
+    log_gamma_q_result<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 double Q_val = exp(result.log_q);
+    const auto Q_val = exp(result.log_q);
     if (Q_val > 0) {
-      const double dQ_da
+      const auto dQ_da
           = grad_reg_inc_gamma(a_dbl, z_dbl, tgamma(a_dbl), digamma(a_dbl));
       result.dlog_q_da = dQ_da / Q_val;
     } else {
       // Fallback if Q rounds to zero - use asymptotic approximation
       result.dlog_q_da = log(z_dbl) - digamma(a_dbl);
     }
+    return result;
   }
-
-  return result;
 }
 
 }  // namespace math

stan/math/prim/prob/gamma_lccdf.hpp

@@ -22,23 +22,22 @@
 #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 {
 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;
+template <bool any_fvar, bool partials_fvar, typename T_shape, typename T>
+static inline auto eval_q_cf(const T& alpha, const T& beta_y) {
+  Q_eval<return_type_t<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));
@@ -62,23 +61,23 @@ static inline Q_eval<T> eval_q_cf(const T& alpha, const T& beta_y) {
     }
   }
 
-  out.ok = std::isfinite(value_of_rec(out.log_Q));
-  return out;
+  if (std::isfinite(value_of_rec(out.log_Q))) {
+    return std::optional<Q_eval<return_type_t<T>>>{out};
+  } else {
+    return std::optional<Q_eval<return_type_t<T>>>{std::nullopt};
+  }
 }
 
 /**
  * Computes log q and d(log q) / d(alpha) using log1m.
  */
-template <typename T, typename T_shape, bool partials_fvar>
+template <bool partials_fvar, typename T_shape, typename T>
 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;
@@ -92,8 +91,6 @@ static inline Q_eval<T> eval_q_log1m(const T& alpha, const T& beta_y) {
           = -grad_reg_lower_inc_gamma(alpha, beta_y) / exp(out.log_Q);
     }
   }
-
-  out.ok = true;
   return out;
 }
 }  // namespace internal
@@ -154,26 +151,21 @@ inline return_type_t<T_y, T_shape, T_inv_scale> gamma_lccdf(
     }
 
     const bool use_continued_fraction = beta_y > alpha_dbl + 1.0;
-    internal::Q_eval<T_partials_return> result;
+    std::optional<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);
+      result = internal::eval_q_cf<any_fvar, partials_fvar, T_shape>(alpha_dbl, beta_y);
     } else {
-      result
-          = internal::eval_q_log1m<T_partials_return, T_shape, partials_fvar>(
-              alpha_dbl, beta_y);
-
-      if (!result.ok && beta_y > 0.0) {
+      result = internal::eval_q_log1m<partials_fvar, T_shape>(alpha_dbl, beta_y);
+      if (!result && 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);
+        result = internal::eval_q_cf<any_fvar, partials_fvar, T_shape>(alpha_dbl, beta_y);
       }
     }
-    if (!result.ok) {
+    if (!result) {
       return ops_partials.build(negative_infinity());
     }
 
-    P += result.log_Q;
+    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);
@@ -183,7 +175,7 @@ inline return_type_t<T_y, T_shape, T_inv_scale> gamma_lccdf(
                                         - lgamma(alpha_dbl) + alpha_minus_one
                                         - beta_y;
 
-      const T_partials_return hazard = exp(log_pdf - result.log_Q);  // 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;
@@ -193,7 +185,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] += result.dlogQ_dalpha;
+      partials<1>(ops_partials)[n] += result->dlogQ_dalpha;
     }
   }
   return ops_partials.build(P);