Add derivative support for std::beta and introduce digamma function utility (#1733)

* Add analytical derivative support for std::beta

* feat: add trigamma function

Author: leetcodez

include/clad/Differentiator/BuiltinDerivatives.h

@@ -1062,6 +1062,112 @@ CUDA_HOST_DEVICE void hypot_pullback(T x, T y, U d_z, T* d_x, T* d_y) {
   *d_y += (y / h) * d_z;
 }
 
+// 7. Special Functions
+#if __cplusplus >= 201703L
+template <typename T> CUDA_HOST_DEVICE inline T clad_beta_primal(T x, T y) {
+#if defined(__cpp_lib_math_special_functions)
+  return ::std::beta(x, y);
+#else
+  return ::std::tgamma(x) * ::std::tgamma(y) / ::std::tgamma(x + y);
+#endif
+}
+// Computes the digamma function psi(x) using a standard asymptotic expansion
+// NOTE: This is a numerical approximation
+// Reference: Wolfram MathWorld, Digamma Function
+// https://mathworld.wolfram.com/DigammaFunction.html
+
+template <typename T> CUDA_HOST_DEVICE inline T clad_digamma(T x) {
+  if (x <= 0.0) {
+    if (x == ::std::floor(x))
+      return (T)NAN;
+    return clad_digamma(1.0 - x) -
+           ::std::acos((T)-1.0) / ::std::tan(::std::acos((T)-1.0) * x);
+  }
+  T result = 0.0;
+  while (x < 8.0) {
+    result -= 1.0 / x;
+    x += 1.0;
+  }
+  T inv_x = 1.0 / x;
+  T inv_x2 = inv_x * inv_x;
+  result +=
+      ::std::log(x) - 0.5 * inv_x -
+      inv_x2 * (1.0 / 12.0 -
+                inv_x2 * (1.0 / 120.0 -
+                          inv_x2 * (1.0 / 252.0 - inv_x2 * (1.0 / 240.0))));
+  return result;
+}
+
+// NOTE: Like digamma this uses a truncated asymptotic expansion.
+// Used as a custom derivative for digamma
+// Reference: Wolfram MathWorld, Trigamma Function
+// https://mathworld.wolfram.com/TrigammaFunction.html
+
+template <typename T> CUDA_HOST_DEVICE inline T clad_trigamma(T x) {
+  if (x <= 0.0) {
+    if (x == ::std::floor(x))
+      return (T)NAN;
+    T pi = ::std::acos((T)-1.0);
+    T csc = 1.0 / ::std::sin(pi * x);
+    return -clad_trigamma(1.0 - x) + (pi * pi * csc * csc);
+  }
+
+  T result = 0.0;
+  while (x < 8.0) {
+    result += 1.0 / (x * x);
+    x += 1.0;
+  }
+
+  T inv_x = 1.0 / x;
+  T inv_x2 = inv_x * inv_x;
+
+  result += inv_x + 0.5 * inv_x2 +
+            inv_x2 * inv_x *
+                (1.0 / 6.0 -
+                 inv_x2 * (1.0 / 30.0 -
+                           inv_x2 * (1.0 / 42.0 - inv_x2 * (1.0 / 30.0))));
+  return result;
+}
+
+template <typename T, typename dT>
+CUDA_HOST_DEVICE ValueAndPushforward<T, dT> digamma_pushforward(T x, dT d_x) {
+  T psi = clad_digamma(x);
+  dT pushforward = 0;
+  if (d_x)
+    pushforward += clad_trigamma(x) * d_x;
+  return {psi, pushforward};
+}
+
+template <typename T, typename U>
+CUDA_HOST_DEVICE void digamma_pullback(T x, U d_z, T* d_x) {
+  if (d_x)
+    *d_x += clad_trigamma(x) * d_z;
+}
+
+template <typename T, typename dT>
+CUDA_HOST_DEVICE ValueAndPushforward<T, dT> beta_pushforward(T x, T y, dT d_x,
+                                                             dT d_y) {
+  T b = clad_beta_primal(x, y);
+  T psi_xy = clad_digamma(x + y);
+  dT pushforward = 0;
+  if (d_x)
+    pushforward += b * (clad_digamma(x) - psi_xy) * d_x;
+  if (d_y)
+    pushforward += b * (clad_digamma(y) - psi_xy) * d_y;
+  return {b, pushforward};
+}
+
+template <typename T, typename U>
+CUDA_HOST_DEVICE void beta_pullback(T x, T y, U d_z, T* d_x, T* d_y) {
+  T b = clad_beta_primal(x, y);
+  T psi_xy = clad_digamma(x + y);
+  if (d_x)
+    *d_x += b * (clad_digamma(x) - psi_xy) * d_z;
+  if (d_y)
+    *d_y += b * (clad_digamma(y) - psi_xy) * d_z;
+}
+#endif
+
 } // namespace std
 
 CUDA_HOST_DEVICE inline ValueAndPushforward<float, float>
@@ -1327,6 +1433,12 @@ using std::pow_pullback;
 using std::pow_pushforward;
 using std::sqrt_pushforward;
 
+// 7. Special Functions
+#if __cplusplus >= 201703L
+using std::beta_pullback;
+using std::beta_pushforward;
+#endif
+
 namespace class_functions {
 template <typename T, typename U>
 void constructor_pullback(ValueAndPushforward<T, U> rhs,

test/Features/stl-cmath.cpp

@@ -98,7 +98,7 @@
 //----------------------- Mathematical special functions -----------------------
 // D assoc_laguerre / f / l    (C++17) associated Laguerre polynomials
 // D assoc_legendre/ f / l     (C++17) associated Legendre polynomials
-// D beta/ betaf/ betal        (C++17) beta function
+// DS beta/ betaf/ betal        (C++17) beta function
 // D comp_ellint_1/ f / l      (C++17) complete elliptic integral (1st kind)
 // D comp_ellint_2/ f / l      (C++17) complete elliptic integral (2nd kind)
 // D comp_ellint_3/ f / l      (C++17) complete elliptic integral (3rd kind)
@@ -123,6 +123,14 @@
 #include <cmath>
 #include <iostream>
 #include <iomanip>
+#if !defined(__cpp_lib_math_special_functions)
+namespace std {
+  template <typename T>
+  inline T beta(T x, T y) {
+    return std::tgamma(x) * std::tgamma(y) / std::tgamma(x + y);
+  }
+}
+#endif
 
 template <typename T>
 T get_tolerance() {
@@ -295,6 +303,10 @@ DEFINE_FUNCTIONS(atanh) // x in [-1,1]
 //
 DEFINE_FUNCTIONS(erf)  // x in (-inf,+inf)
 
+template<typename T> T f_beta(T x){ return std::beta(x,(T)2.0); } // x in (0, +inf)
+inline float f_betaf(float x){ return std::beta(x, 2.0f); }
+inline long double f_betal(long double x){ return std::beta(x, 2.0L); }
+
 int main() {
   // Absolute value
   CHECK(abs);
@@ -352,6 +364,7 @@ int main() {
 
   // Error / Gamma functions
   CHECK_ALL(erf);
+  CHECK_ALL_RANGE(beta, {0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0});
 
   return 0;
 }

test/FirstDerivative/BuiltinDerivatives.C

@@ -6,6 +6,17 @@
 
 #include "clad/Differentiator/Differentiator.h"
 #include "../TestUtils.h"
+#include <cmath>
+
+#if !defined(__cpp_lib_math_special_functions)
+namespace std {
+  // Mock std::beta on Apple platforms so Clad's AST has a target to differentiate
+  inline double beta(double x, double y) {
+    return std::tgamma(x) * std::tgamma(y) / std::tgamma(x + y);
+  }
+}
+#endif
+
 extern "C" int printf(const char* fmt, ...);
 
 
@@ -515,6 +526,14 @@ double f_custom_min(double x, double y) { return std::min(x, y, std::greater<dou
 // CHECK-NEXT:      return _t0.pushforward;
 // CHECK-NEXT:  }
 
+double f_beta(double x, double y) { return std::beta(x, y); }
+// CHECK: double f_beta_darg0(double x, double y) {
+// CHECK-NEXT:     double _d_x = 1;
+// CHECK-NEXT:     double _d_y = 0;
+// CHECK-NEXT:     {{.*}}ValueAndPushforward<double, double> _t0 = {{.*}}beta_pushforward(x, y, _d_x, _d_y);
+// CHECK-NEXT:     return _t0.pushforward;
+// CHECK-NEXT: }
+
 int main () { //expected-no-diagnostics
   float f_result[2];
   double d_result[2];
@@ -695,5 +714,8 @@ int main () { //expected-no-diagnostics
   auto d_custom_min =  clad::differentiate(f_custom_min, 0);
   printf("Result is = %.6f\n", d_custom_min.execute(2, 3)); // CHECK-EXEC: Result is = 0.000000
 
+  auto d_beta = clad::differentiate(f_beta, 0);
+  printf("Result is = %.6f\n", d_beta.execute(2.0, 3.0)); // CHECK-EXEC: Result is = -0.090278
+
   return 0;
 }