Add derivatives for ellipk and ellipe (#370)

* Add derivatives for `ellipk` and `ellipe`

* Bump version

* Fix derivatives at 0

Project.toml

@@ -1,6 +1,6 @@
 name = "SpecialFunctions"
 uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
-version = "2.0.0"
+version = "2.1.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/chainrules.jl

@@ -194,6 +194,16 @@ ChainRulesCore.@scalar_rule(expinti(x), exp(x) / x)
 ChainRulesCore.@scalar_rule(sinint(x), sinc(invπ * x))
 ChainRulesCore.@scalar_rule(cosint(x), cos(x) / x)
 
+# elliptic integrals
+ChainRulesCore.@scalar_rule(
+    ellipk(m),
+    iszero(m) ? oftype(Ω, π) / 8 : (ellipe(m) / (1 - m) - Ω) / (2 * m),
+)
+ChainRulesCore.@scalar_rule(
+    ellipe(m),
+    iszero(m) ? -oftype(Ω, π) / 8 : (Ω - ellipk(m)) / (2 * m),
+)
+
 # non-holomorphic functions
 function ChainRulesCore.frule((_, Δν, Δx), ::typeof(besselix), ν::Number, x::Number)
     # primal

src/ellip.jl

@@ -39,7 +39,9 @@ For ``m<0``, followed by
 > <https://www.researchgate.net/publication/267330394>
 As suggested in this paper, the domain is restricted to ``(-\infty,1]``.
 """
-function ellipk(m::Float64)
+ellipk(m::Real) = _ellipk(float(m))
+
+function _ellipk(m::Float64)
     flag_is_m_neg = false
     if m < 0.0
         x               = m / (m-1)         #dealing with negative args
@@ -214,7 +216,9 @@ For ``m<0``, followed by
 > <https://www.researchgate.net/publication/267330394>
 As suggested in this paper, the domain is restricted to ``(-\infty,1]``.
 """
-function ellipe(m::Float64)
+ellipe(m::Real) = _ellipe(float(m))
+
+function _ellipe(m::Float64)
     flag_is_m_neg = false
     if m < 0.0
         x               = m / (m-1)         #dealing with negative args
@@ -346,11 +350,7 @@ function ellipe(m::Float64)
     end
 end
 
-for f in (:ellipk,:ellipe)
-    @eval begin
-        ($f)(x::Float16)        = Float16(($f)(Float64(x)))
-	    ($f)(x::Float32)        = Float32(($f)(Float64(x)))
-        ($f)(x::Real)           = ($f)(float(x))
-	    ($f)(x::AbstractFloat)  = throw(MethodError($f, (x, "")))
-    end
+# Support for Float32 and Float16
+for internalf in (:_ellipk, :_ellipe), T in (:Float16, :Float32)
+    @eval $internalf(x::$T) = $T($internalf(Float64(x)))
 end

test/chainrules.jl

@@ -37,6 +37,11 @@
                     test_scalar(airyaiprimex, x)
                 end
             end
+
+            if x isa Real && x < 1
+                test_scalar(ellipk, x)
+                test_scalar(ellipe, x)
+            end
         end
     end