Incomplete gamma gradient roots and log-space cdfs for six distributions

stan/math/fwd/fun.hpp

@@ -3,6 +3,7 @@
 
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/fwd/fun/Eigen_NumTraits.hpp>
+#include <stan/math/fwd/fun/Eigen_SpecialFunctions.hpp>
 
 #include <stan/math/fwd/fun/abs.hpp>
 #include <stan/math/fwd/fun/accumulator.hpp>

stan/math/fwd/fun/Eigen_SpecialFunctions.hpp

@@ -0,0 +1,58 @@
+#ifndef STAN_MATH_FWD_FUN_EIGEN_SPECIALFUNCTIONS_HPP
+#define STAN_MATH_FWD_FUN_EIGEN_SPECIALFUNCTIONS_HPP
+
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/fwd/core/fvar.hpp>
+#include <stan/math/fwd/fun/digamma.hpp>
+#include <stan/math/fwd/fun/lgamma.hpp>
+#include <unsupported/Eigen/SpecialFunctions>
+
+namespace Eigen {
+namespace internal {
+
+/**
+ * Specializations of the three `Eigen::internal` helpers that
+ * `Eigen::internal::igamma_generic_impl` restricts to `float` and
+ * `double`: `lgamma_impl`, `digamma_impl` and `cephes_helper`. With these,
+ * `Eigen::numext::igamma_der_a` accepts `fvar<T>` at any autodiff order.
+ *
+ * These names are Eigen internals. An Eigen upgrade that renames them
+ * fails here at compile time.
+ */
+template <typename T>
+struct lgamma_impl<stan::math::fvar<T>> {
+  EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE stan::math::fvar<T> run(
+      const stan::math::fvar<T>& x) {
+    return stan::math::lgamma(x);
+  }
+};
+
+template <typename T>
+struct digamma_impl<stan::math::fvar<T>> {
+  EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE stan::math::fvar<T> run(
+      const stan::math::fvar<T>& x) {
+    return stan::math::digamma(x);
+  }
+};
+
+/**
+ * The Cephes constants are properties of the underlying floating point
+ * format, not of the autodiff type, so they carry a zero tangent.
+ */
+template <typename T>
+struct cephes_helper<stan::math::fvar<T>> {
+  EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE stan::math::fvar<T> machep() {
+    return stan::math::fvar<T>(cephes_helper<double>::machep());
+  }
+  EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE stan::math::fvar<T> big() {
+    return stan::math::fvar<T>(cephes_helper<double>::big());
+  }
+  EIGEN_DEVICE_FUNC static EIGEN_STRONG_INLINE stan::math::fvar<T> biginv() {
+    return stan::math::fvar<T>(cephes_helper<double>::biginv());
+  }
+};
+
+}  // namespace internal
+}  // namespace Eigen
+
+#endif

stan/math/fwd/fun/gamma_q.hpp

@@ -5,71 +5,49 @@
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/fun/digamma.hpp>
 #include <stan/math/fwd/fun/exp.hpp>
-#include <stan/math/fwd/fun/fabs.hpp>
+#include <stan/math/fwd/fun/lgamma.hpp>
 #include <stan/math/fwd/fun/log.hpp>
-#include <stan/math/fwd/fun/pow.hpp>
 #include <stan/math/fwd/fun/tgamma.hpp>
 #include <stan/math/prim/fun/gamma_q.hpp>
+#include <stan/math/prim/fun/grad_reg_inc_gamma.hpp>
 #include <cmath>
 
 namespace stan {
 namespace math {
 
+/*
+ * The derivative with respect to the first argument is grad_reg_inc_gamma.
+ * The derivative with respect to the second argument is minus the gamma
+ * density, evaluated in log space so that neither pow nor tgamma can
+ * overflow; this is the same form that gamma_p uses.
+ */
+
 template <typename T>
 inline fvar<T> gamma_q(const fvar<T>& x1, const fvar<T>& x2) {
   T u = gamma_q(x1.val_, x2.val_);
 
-  T S = 0;
-  T s = 1;
-  T l = log(x2.val_);
-  T g = tgamma(x1.val_);
-  T dig = digamma(x1.val_);
-
-  int k = 0;
-  T delta = s / (x1.val_ * x1.val_);
-
-  while (fabs(delta) > 1e-6) {
-    S += delta;
-    ++k;
-    s *= -x2.val_ / k;
-    delta = s / ((k + x1.val_) * (k + x1.val_));
-  }
-
-  T der1 = (1.0 - u) * (dig - l) + exp(x1.val_ * l) * S / g;
-  T der2 = -exp(-x2.val_) * pow(x2.val_, x1.val_ - 1.0) / g;
+  T der1
+      = grad_reg_inc_gamma(x1.val_, x2.val_, tgamma(x1.val_), digamma(x1.val_));
+  T der2 = -exp(-x2.val_ + (x1.val_ - 1.0) * log(x2.val_) - lgamma(x1.val_));
 
   return fvar<T>(u, x1.d_ * der1 + x2.d_ * der2);
 }
 
 template <typename T>
 inline fvar<T> gamma_q(const fvar<T>& x1, double x2) {
   T u = gamma_q(x1.val_, x2);
-  T S = 0;
-  double s = 1;
-  double l = log(x2);
-  T g = tgamma(x1.val_);
-  T dig = digamma(x1.val_);
 
-  int k = 0;
-  T delta = s / (x1.val_ * x1.val_);
-
-  while (fabs(delta) > 1e-6) {
-    S += delta;
-    ++k;
-    s *= -x2 / k;
-    delta = s / ((k + x1.val_) * (k + x1.val_));
-  }
-
-  T der1 = (1.0 - u) * (dig - l) + exp(x1.val_ * l) * S / g;
+  T der1 = grad_reg_inc_gamma(x1.val_, x2, tgamma(x1.val_), digamma(x1.val_));
 
   return fvar<T>(u, x1.d_ * der1);
 }
 
 template <typename T>
 inline fvar<T> gamma_q(double x1, const fvar<T>& x2) {
   T u = gamma_q(x1, x2.val_);
-  double g = tgamma(x1);
-  T der2 = -exp(-x2.val_) * pow(x2.val_, x1 - 1.0) / g;
+
+  T der2 = -exp(-x2.val_ + (x1 - 1.0) * log(x2.val_) - lgamma(x1));
+
   return fvar<T>(u, x2.d_ * der2);
 }
 }  // namespace math

stan/math/prim/fun/grad_reg_inc_gamma.hpp

@@ -2,115 +2,49 @@
 #define STAN_MATH_PRIM_FUN_GRAD_REG_INC_GAMMA_HPP
 
 #include <stan/math/prim/meta.hpp>
-#include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/constants.hpp>
-#include <stan/math/prim/fun/exp.hpp>
-#include <stan/math/prim/fun/gamma_p.hpp>
-#include <stan/math/prim/fun/gamma_q.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/is_any_nan.hpp>
-#include <stan/math/prim/fun/is_inf.hpp>
-#include <stan/math/prim/fun/log.hpp>
-#include <stan/math/prim/fun/multiply_log.hpp>
-#include <cmath>
+#include <unsupported/Eigen/SpecialFunctions>
 #include <limits>
 
 namespace stan {
 namespace math {
 
 /**
- * Gradient of the regularized incomplete gamma functions igamma(a, z)
+ * Gradient of the regularized incomplete gamma function igamma(a, z) with
+ * respect to the shape parameter, d/da Q(a, z), where Q is the upper
+ * regularized incomplete gamma function.
  *
- * For small z, the gradient is computed via the series expansion;
- * for large z, the series is numerically inaccurate due to cancellation
- * and the asymptotic expansion is used.
+ * Delegates to `Eigen::numext::igamma_der_a`, which returns d/da of the
+ * lower P(a, z); the sign is flipped here. For autodiff scalar types the
+ * Cephes helpers that Eigen needs are specialized in
+ * `stan/math/fwd/fun/Eigen_SpecialFunctions.hpp` and
+ * `stan/math/rev/fun/Eigen_SpecialFunctions.hpp`.
+ *
+ * `g`, `dig`, `precision` and `max_steps` are accepted for signature
+ * compatibility and are not used.
  *
  * @tparam T1 type of the shape parameter
  * @tparam T2 type of the location parameter
  * @param a shape parameter, a > 0
  * @param z location z >= 0
- * @param g stan::math::tgamma(a) (precomputed value)
- * @param dig boost::math::digamma(a) (precomputed value)
- * @param precision required precision; applies to series expansion only
- * @param max_steps number of steps to take.
- * @throw throws std::domain_error if not converged after max_steps
- *   or increment overflows to inf.
- *
- * For the asymptotic expansion, the gradient is given by:
-   \f[
-   \begin{array}{rcl}
-   \Gamma(a, z) & = & z^{a-1}e^{-z} \sum_{k=0}^N \frac{(a-1)_k}{z^k} \qquad , z
- \gg a\\
-   Q(a, z) & = & \frac{z^{a-1}e^{-z}}{\Gamma(a)} \sum_{k=0}^N
- \frac{(a-1)_k}{z^k}\\
-   (a)_k & = & (a)_{k-1}(a-k)\\
-   \frac{d}{da} (a)_k & = & (a)_{k-1} + (a-k)\frac{d}{da} (a)_{k-1}\\
-   \frac{d}{da}Q(a, z) & = & (log(z) - \psi(a)) Q(a, z)\\
-   && + \frac{z^{a-1}e^{-z}}{\Gamma(a)} \sum_{k=0}^N \left(\frac{d}{da}
- (a-1)_k\right) \frac{1}{z^k} \end{array} \f]
+ * @param g unused
+ * @param dig unused
+ * @param precision unused
+ * @param max_steps unused
+ * @return d/da of the upper regularized incomplete gamma function
  */
 template <typename T1, typename T2>
 inline return_type_t<T1, T2> grad_reg_inc_gamma(T1 a, T2 z, T1 g, T1 dig,
                                                 double precision = 1e-6,
                                                 int max_steps = 1e5) {
-  using std::exp;
-  using std::fabs;
-  using std::log;
   using TP = return_type_t<T1, T2>;
 
   if (is_any_nan(a, z, g, dig)) {
     return std::numeric_limits<TP>::quiet_NaN();
   }
 
-  T2 l = log(z);
-  if (z >= a && z >= 8) {
-    // asymptotic expansion http://dlmf.nist.gov/8.11#E2
-    TP S = 0;
-    T1 a_minus_one_minus_k = a - 1;
-    T1 fac = a_minus_one_minus_k;  // falling_factorial(a-1, k)
-    T1 dfac = 1;                   // d/da[falling_factorial(a-1, k)]
-    T2 zpow = z;                   // z ** k
-    TP delta = dfac / zpow;
-
-    for (int k = 1; k < 10; ++k) {
-      a_minus_one_minus_k -= 1;
-
-      S += delta;
-
-      zpow *= z;
-      dfac = a_minus_one_minus_k * dfac + fac;
-      fac *= a_minus_one_minus_k;
-      delta = dfac / zpow;
-
-      if (is_inf(delta)) {
-        throw_domain_error("grad_reg_inc_gamma", "is not converging", "", "");
-      }
-    }
-
-    return gamma_q(a, z) * (l - dig) + exp(-z + (a - 1) * l) * S / g;
-  } else {
-    // gradient of series expansion http://dlmf.nist.gov/8.7#E3
-    TP S = 0;
-    TP log_s = 0.0;
-    double s_sign = 1.0;
-    T2 log_z = log(z);
-    TP log_delta = log_s - multiply_log(2, a);
-    for (int k = 1; k <= max_steps; ++k) {
-      S += s_sign >= 0.0 ? exp(log_delta) : -exp(log_delta);
-      log_s += log_z - log(k);
-      s_sign = -s_sign;
-      log_delta = log_s - multiply_log(2, k + a);
-      if (is_inf(log_delta)) {
-        throw_domain_error("grad_reg_inc_gamma", "is not converging", "", "");
-      }
-      if (log_delta <= log(precision)) {
-        return gamma_p(a, z) * (dig - l) + exp(a * l) * S / g;
-      }
-    }
-    throw_domain_error(
-        "grad_reg_inc_gamma", "k (internal counter)", max_steps, "exceeded ",
-        " iterations, gamma function gradient did not converge.");
-    return INFTY;
-  }
+  return -Eigen::numext::igamma_der_a(TP(a), TP(z));
 }
 
 }  // namespace math