Feature: Add std::comp_ellint derivatives (#1674)

Fixes #1673

Author: ayush4874

include/clad/Differentiator/BuiltinDerivatives.h

@@ -1213,6 +1213,85 @@ CUDA_HOST_DEVICE void beta_pullback(T x, T y, U d_z, T* d_x, T* d_y) {
 }
 #endif
 
+#if __cplusplus >= 201703L && (defined(__cpp_lib_math_special_funcs) ||        \
+                               defined(__STDCPP_MATH_SPEC_FUNCS__))
+template <typename T, typename dT,
+          typename T_out = decltype(::std::comp_ellint_1(T())),
+          typename dT_out = typename AdjOutType<T_out, dT>::type>
+CUDA_HOST_DEVICE ValueAndPushforward<T_out, dT_out>
+comp_ellint_1_pushforward(T k, dT d_k) {
+  T_out K = ::std::comp_ellint_1(k);
+  T_out E = ::std::comp_ellint_2(k);
+  T_out k_sq = k * k;
+  T_out one = 1.0;
+  T_out derivative = 0.0;
+  if (k_sq < 1.0 && k != 0.0)
+    derivative = (E - (one - k_sq) * K) / (k * (one - k_sq));
+  return {K, static_cast<dT_out>(d_k) * derivative};
+}
+
+template <typename T, typename dT,
+          typename T_out = decltype(::std::comp_ellint_2(T())),
+          typename dT_out = typename AdjOutType<T_out, dT>::type>
+CUDA_HOST_DEVICE ValueAndPushforward<T_out, dT_out>
+comp_ellint_2_pushforward(T k, dT d_k) {
+  T_out K = ::std::comp_ellint_1(k);
+  T_out E = ::std::comp_ellint_2(k);
+  T_out derivative = 0.0;
+  if (k != 0.0)
+    derivative = (E - K) / k;
+  return {E, static_cast<dT_out>(d_k) * derivative};
+}
+
+template <typename T1, typename T2, typename dT1, typename dT2,
+          typename T_out = decltype(::std::comp_ellint_3(T1(), T2())),
+          typename dT_out = typename AdjOutType<T_out, dT1>::type>
+CUDA_HOST_DEVICE ValueAndPushforward<T_out, dT_out>
+comp_ellint_3_pushforward(T1 k, T2 nu, dT1 d_k, dT2 d_nu) {
+  T_out K = ::std::comp_ellint_1(k);
+  T_out E = ::std::comp_ellint_2(k);
+  T_out Pi = ::std::comp_ellint_3(k, nu);
+  T_out k2 = k * k;
+  T_out one = 1.0;
+  T_out grad_k = 0.0;
+  if (k2 != static_cast<T_out>(nu) && k != 0.0 && k2 < 1.0) {
+    T_out term_k = E / (one - k2);
+    grad_k = (k / (k2 - static_cast<T_out>(nu))) * (term_k - Pi);
+  }
+  T_out grad_nu = 0.0;
+  if (nu != 0.0 && nu != 1.0 && k2 != static_cast<T_out>(nu)) {
+    T_out p2 = ((k2 - static_cast<T_out>(nu)) / static_cast<T_out>(nu)) * K;
+    T_out p3 = ((static_cast<T_out>(nu) * static_cast<T_out>(nu) - k2) /
+                static_cast<T_out>(nu)) *
+               Pi;
+    grad_nu = (one / (2.0 * (static_cast<T_out>(nu) - one) *
+                      (k2 - static_cast<T_out>(nu)))) *
+              (E + p2 + p3);
+  }
+  return {Pi, (static_cast<dT_out>(d_k) * grad_k) +
+                  (static_cast<dT_out>(d_nu) * grad_nu)};
+}
+
+template <typename T1, typename T2, typename T3>
+CUDA_HOST_DEVICE void comp_ellint_3_pullback(T1 k, T2 nu, T3 d_out, T1* d_k,
+                                             T2* d_nu) {
+  auto K = ::std::comp_ellint_1(k);
+  auto E = ::std::comp_ellint_2(k);
+  auto Pi = ::std::comp_ellint_3(k, nu);
+  auto k2 = k * k;
+  auto one = 1.0;
+  if (k2 != nu && k != 0.0 && k2 < 1.0) {
+    auto term_k = E / (one - k2);
+    *d_k += d_out * ((k / (k2 - nu)) * (term_k - Pi));
+  }
+  if (nu != 0.0 && nu != one && k2 != nu) {
+    auto p2 = ((k2 - nu) / nu) * K;
+    auto p3 = ((nu * nu - k2) / nu) * Pi;
+    *d_nu += d_out * ((one / (2.0 * (nu - one) * (k2 - nu))) * (E + p2 + p3));
+  }
+}
+#endif
+
 } // namespace std
 
 CUDA_HOST_DEVICE inline ValueAndPushforward<float, float>
@@ -1489,6 +1568,14 @@ using std::beta_pullback;
 using std::beta_pushforward;
 #endif
 
+#if __cplusplus >= 201703L && (defined(__cpp_lib_math_special_funcs) ||        \
+                               defined(__STDCPP_MATH_SPEC_FUNCS__))
+using std::comp_ellint_1_pushforward;
+using std::comp_ellint_2_pushforward;
+using std::comp_ellint_3_pullback;
+using std::comp_ellint_3_pushforward;
+#endif
+
 namespace class_functions {
 template <typename T, typename U>
 void constructor_pullback(ValueAndPushforward<T, U> rhs,

test/Features/stl-cmath.cpp

@@ -99,9 +99,9 @@
 // D assoc_laguerre / f / l    (C++17) associated Laguerre polynomials
 // D assoc_legendre/ f / l     (C++17) associated Legendre polynomials
 // 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)
+// DS comp_ellint_1/ f / l      (C++17) complete elliptic integral (1st kind)
+// DS comp_ellint_2/ f / l      (C++17) complete elliptic integral (2nd kind)
+// DS comp_ellint_3/ f / l      (C++17) complete elliptic integral (3rd kind)
 // D cyl_bessel_i/ f / l       (C++17) modified cylindrical Bessel (regular)
 // D cyl_bessel_j/ f / l       (C++17) cylindrical Bessel functions (1st kind)
 // D cyl_bessel_k/ f / l       (C++17) modified cylindrical Bessel (irregular)
@@ -313,6 +313,25 @@ template<typename T> T f_beta(T x){ return std::beta(x,(T)2.0); } // x in (0, +i
 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); }
 
+#if __cplusplus >= 201703L && (defined(__cpp_lib_math_special_funcs) || defined(__STDCPP_MATH_SPEC_FUNCS__))
+//------------------------ Elliptic integrals -----------------------------
+//
+// Domain: k in (-1, 1)
+template<typename T> T f_comp_ellint_1(T k) { return std::comp_ellint_1(k); }
+float f_comp_ellint_1f(float k) { return std::comp_ellint_1(k); }
+long double f_comp_ellint_1l(long double k) { return std::comp_ellint_1(k); }
+
+// Domain: k in (-1, 1)
+template<typename T> T f_comp_ellint_2(T k) { return std::comp_ellint_2(k); }
+float f_comp_ellint_2f(float k) { return std::comp_ellint_2(k); }
+long double f_comp_ellint_2l(long double k) { return std::comp_ellint_2(k); }
+
+// Domain: k in (-1, 1). Fixed nu = 0.5 for testing.
+template<typename T> T f_comp_ellint_3(T k) { return std::comp_ellint_3(k, (T)0.5); }
+float f_comp_ellint_3f(float k) { return std::comp_ellint_3(k, 0.5f); }
+long double f_comp_ellint_3l(long double k) { return std::comp_ellint_3(k, 0.5L); }
+#endif
+
 int main() {
   // Absolute value
   CHECK(abs);
@@ -376,5 +395,12 @@ int main() {
   CHECK_ALL(erf);
   CHECK_ALL_RANGE(beta, {0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0});
 
+  #if __cplusplus >= 201703L && (defined(__cpp_lib_math_special_funcs) || defined(__STDCPP_MATH_SPEC_FUNCS__))
+  // Elliptic Integrals
+  CHECK_ALL_RANGE(comp_ellint_1, {-0.9, -0.6, -0.3, 0.0, 0.3, 0.6, 0.9});
+  CHECK_ALL_RANGE(comp_ellint_2, {-0.9, -0.6, -0.3, 0.0, 0.3, 0.6, 0.9});
+  CHECK_ALL_RANGE(comp_ellint_3, {-0.9, -0.6, -0.3, 0.0, 0.3, 0.6, 0.9});
+  #endif
+
   return 0;
 }