Make elliptic integrals differentiable

[archermarx]

This PR addresses #28 by removing all type annotations for the elliptic integrals so that they are auto-differentiable. I have checked using Zygote.jl and ForwardDiff.jl and have added a suite of tests to verify several elliptic integral differentiation identities.

[henry2004y]

I also wish to have this compatibility. Otherwise I may switch to SpecialFunctions for equivalent calculations.

[archermarx]

Hi @henry2004y, check out https://github.com/dominic-chang/JacobiElliptic.jl by @dominic-chang. It has differentiable and GPU-compatible version of several of the special functions in this package.

[henry2004y]

Hi @henry2004y, check out https://github.com/dominic-chang/JacobiElliptic.jl by @dominic-chang. It has differentiable and GPU-compatible version of several of the special functions in this package.

Thanks for sharing! Is it also true that JacobiElliptic.jl is faster than Elliptic.jl and SpecialFunctions.jl?

[archermarx]

Not sure. I suspect it's about as fast.

[henry2004y]

Benchmark Script

using Pkg
Pkg.activate(; temp = true)
Pkg.add(["Chairmarks", "SpecialFunctions", "JacobiElliptic"])

using Chairmarks
using SpecialFunctions
using JacobiElliptic

# Define a range of values to test, avoiding singularities if necessary
ks = 0.1:0.1:0.9
ms = ks .^ 2

println("Benchmarking Elliptic K (Complete Elliptic Integral of the First Kind)")
println("SpecialFunctions.ellipk vs JacobiElliptic.K")

# Warmup to ensure compilation
ellipk.(ms)
K.(ms)

b_sf_k = @be ellipk.($ms)
b_je_k = @be K.($ms)

println("\nSpecialFunctions.ellipk:")
display(b_sf_k)
println("\nJacobiElliptic.K:")
display(b_je_k)

println("\n---------------------------------------------------")

println("Benchmarking Elliptic E (Complete Elliptic Integral of the Second Kind)")
println("SpecialFunctions.ellipe vs JacobiElliptic.E")

# Warmup
ellipe.(ms)
E.(ms)

b_sf_e = @be ellipe.($ms)
b_je_e = @be E.($ms)

println("\nSpecialFunctions.ellipe:")
display(b_sf_e)
println("\nJacobiElliptic.E:")
display(b_je_e)

Benchmarking Elliptic K (Complete Elliptic Integral of the First Kind)
SpecialFunctions.ellipk vs JacobiElliptic.K

SpecialFunctions.ellipk:
Benchmark: 951 samples with 513 evaluations
 min    49.123 ns (2 allocs: 128 bytes)
 median 54.971 ns (2 allocs: 128 bytes)
 mean   338.571 ns (2 allocs: 128 bytes, 0.11% gc time)
 max    267.867 μs (2 allocs: 128 bytes, 99.95% gc time)

JacobiElliptic.K:
Benchmark: 3192 samples with 66 evaluations
 min    413.636 ns (2 allocs: 128 bytes)
 median 434.848 ns (2 allocs: 128 bytes)
 mean   444.154 ns (2 allocs: 128 bytes)
 max    5.708 μs (2 allocs: 128 bytes)

---------------------------------------------------
Benchmarking Elliptic E (Complete Elliptic Integral of the Second Kind)
SpecialFunctions.ellipe vs JacobiElliptic.E

SpecialFunctions.ellipe:
Benchmark: 2646 samples with 475 evaluations
 min    46.737 ns (2 allocs: 128 bytes)
 median 53.263 ns (2 allocs: 128 bytes)
 mean   76.577 ns (2 allocs: 128 bytes, 0.08% gc time)
 max    34.021 μs (2 allocs: 128 bytes, 99.56% gc time)

JacobiElliptic.E:
Benchmark: 3079 samples with 29 evaluations
 min    975.862 ns (2 allocs: 128 bytes)
 median 1.028 μs (2 allocs: 128 bytes)
 mean   1.051 μs (2 allocs: 128 bytes)
 max    14.186 μs (2 allocs: 128 bytes)

It seems that it's not even close... @archermarx

[archermarx]

Hi @henry2004y, I am not involved in the maintenance of JacobiElliptic.jl, so I'd recommend putting in a PR on that repo.

A while ago, I also made a fork of this package to make it differentiable: https://github.com/archermarx/elliptic2.jl. I do not maintain this repository anymore and offer no guarantees about its performance or stability, but it's another option.