Adding computation of complete elliptic integrals into amrex::Math so…

Docs/sphinx_documentation/source/Basics.rst

@@ -596,7 +596,7 @@ the Heaviside step function, ``heaviside(x1,x2)`` that gives ``0``, ``x2``,
 ``1``, for ``x1 < 0``, ``x1 = 0`` and ``x1 > 0``, respectively.
 It supports the Bessel function of the first kind of order ``n``
 ``jn(n,x)``. Complete elliptic integrals of the first and second kind, ``comp_ellint_1`` and ``comp_ellint_2``,
-are supported only for gcc and CPUs.
+are supported.
 There is ``if(a,b,c)`` that gives ``b`` or ``c`` depending on the value of
 ``a``.  A number of comparison operators are supported, including ``<``,
 ``>``, ``==``, ``!=``, ``<=``, and ``>=``.  The Boolean results from

Src/Base/AMReX_Math.H

@@ -5,6 +5,7 @@
 #include <AMReX_GpuQualifiers.H>
 #include <AMReX_Extension.H>
 #include <AMReX_INT.H>
+#include <AMReX_REAL.H>
 #include <cmath>
 #include <cstdlib>
 #include <type_traits>
@@ -225,6 +226,70 @@ std::uint64_t umulhi (std::uint64_t a, std::uint64_t b)
 }
 #endif
 
+template <typename T>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T comp_ellint_1 (T m)
+{
+    // Computing K based on DLMF
+    // https://dlmf.nist.gov/19.8
+    T tol = 1e-12;
+    T a0 = 1.0;
+    T g0 = std::sqrt(1.0 - m);
+    T a, g;
+
+    // Find Arithmetic Geometric mean
+    while(std::abs(a0 - g0) > tol) {
+        a = 0.5*(a0 + g0);
+        g = std::sqrt(a0 * g0);
+
+        a0 = a;
+        g0 = g;
+    }
+
+    return 0.5*pi<T>()/a;
+}
+
+template <typename T>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T comp_ellint_2 (T m)
+{
+    // Computing E based on DLMF
+    // https://dlmf.nist.gov/19.8
+    T Kcomp = amrex::Math::comp_ellint_1<T>(m);
+    T tol = 1e-12;
+
+    // Step Zero
+    T a0 = 1.0;
+    T g0 = std::sqrt(1.0 - m);
+    T cn = std::sqrt(a0*a0 - g0*g0);
+
+    // Step 1
+    int n = 1;
+    T a = 0.5 * (a0 + g0);
+    T g = std::sqrt(a0*g0);
+    cn = 0.25*cn*cn/a;
+
+    T sum_val = a*a;
+    a0 = a;
+    g0 = g;
+
+    while(std::abs(cn*cn) > tol) {
+        // Compute coefficients for this iteration
+        a = 0.5 * (a0 + g0);
+        g = std::sqrt(a0*g0);
+        cn = 0.25*cn*cn/a;
+
+        n++;
+        sum_val -= std::pow(2,n-1)*cn*cn;
+
+        // Save a and g for next iteration
+        a0 = a;
+        g0 = g;
+    }
+
+    return Kcomp*sum_val;
+}
+
 /***************************************************************************************************
  * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: BSD-3-Clause

Src/Base/Parser/AMReX_Parser_Y.H

@@ -349,28 +349,24 @@ AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
 T parser_math_atanh (T a) { return std::atanh(a); }
 
 template <typename T>
-AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
-T parser_math_comp_ellint_1 (T a)
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T parser_math_comp_ellint_1 (T k)
 {
 #if defined(__GNUC__) && !defined(__clang__) && !defined(__CUDA_ARCH__) && !defined(__NVCOMPILER)
-    return std::comp_ellint_1(a);
+    return std::comp_ellint_1(k);
 #else
-    amrex::ignore_unused(a);
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(false, "parser: comp_ellint_1 only supported with gcc and cpu");
-    return 0.0;
+    return amrex::Math::comp_ellint_1<T>(k);
 #endif
 }
 
 template <typename T>
-AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
-T parser_math_comp_ellint_2 (T a)
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T parser_math_comp_ellint_2 (T k)
 {
 #if defined(__GNUC__) && !defined(__clang__) && !defined(__CUDA_ARCH__) && !defined(__NVCOMPILER)
-    return std::comp_ellint_2(a);
+    return std::comp_ellint_2(k);
 #else
-    amrex::ignore_unused(a);
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(false, "parser: comp_ellint_2 only supported with gcc and cpu");
-    return 0.0;
+    return amrex::Math::comp_ellint_2<T>(k);
 #endif
 }