[Derivatives] Implement forward-mode derivative for std::comp_ellint_1

include/clad/Differentiator/BuiltinDerivatives.h

@@ -1062,6 +1062,84 @@ 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;
 }
 
+#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>
@@ -1326,6 +1404,12 @@ using std::min_pushforward;
 using std::pow_pullback;
 using std::pow_pushforward;
 using std::sqrt_pushforward;
+#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>

test/Features/stl-cmath.cpp

@@ -66,6 +66,11 @@
 // D tgamma/ tgammaf/ tgammal          (C++11) gamma function
 // D lgamma/ lgammaf/ lgammal          (C++11) log gamma
 //
+//------------------------ Elliptic integrals -----------------------------
+// D comp_ellint_1 / comp_ellint_1f / comp_ellint_1l (C++17) complete elliptic integral of the first kind
+// D comp_ellint_2 / comp_ellint_2f / comp_ellint_2l (C++17) complete elliptic integral of the second kind
+// D comp_ellint_3 (C++17) complete elliptic integral of the third kind
+//
 //---------------------- Nearest integer operations ----------------------------
 // N ceil/ ceilf/ ceill                (C++11) smallest integer >= x
 // N floor/ floorf/ floorl             (C++11) largest integer <= x
@@ -295,6 +300,25 @@ DEFINE_FUNCTIONS(atanh) // x in [-1,1]
 //
 DEFINE_FUNCTIONS(erf)  // x in (-inf,+inf)
 
+#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);
@@ -352,6 +376,13 @@ int main() {
 
   // Error / Gamma functions
   CHECK_ALL(erf);
+
+  #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;
-}
+}