diff --git a/docs/src/index.md b/docs/src/index.md index 4dead14b..edd7f75b 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -40,6 +40,9 @@ libraries. | [`besselix(nu,z)`](@ref SpecialFunctions.besselix) | scaled modified Bessel function of the first kind of order `nu` at `z` | | [`besselk(nu,z)`](@ref SpecialFunctions.besselk) | modified [Bessel function](https://en.wikipedia.org/wiki/Bessel_function) of the second kind of order `nu` at `z` | | [`besselkx(nu,z)`](@ref SpecialFunctions.besselkx) | scaled modified Bessel function of the second kind of order `nu` at `z` | +| [`ellipj(u,m)`](@ref SpecialFunctions.ellipj) | Jacobi elliptic functions `sn,cn,dn` | +| `jpq(u,m)` | Jacobi elliptic function `pq` | +| [`ellipK(m)`](@ref SpecialFunctions.ellipj) | Complete elliptic integral of the first kind | ## Installation diff --git a/docs/src/special.md b/docs/src/special.md index f8bb48b1..43b74ddc 100644 --- a/docs/src/special.md +++ b/docs/src/special.md @@ -46,4 +46,6 @@ SpecialFunctions.besselk SpecialFunctions.besselkx SpecialFunctions.eta SpecialFunctions.zeta +SpecialFunctions.ellipj +SpecialFunctions.K ``` diff --git a/src/SpecialFunctions.jl b/src/SpecialFunctions.jl index 5994a194..f438f4ec 100644 --- a/src/SpecialFunctions.jl +++ b/src/SpecialFunctions.jl @@ -71,7 +71,12 @@ end export sinint, cosint +export ellipj, + jss,jsc,jsd,jsn,jcs,jcc,jcd,jcn,jds,jdc,jdd,jdn,jns,jnc,jnd,jnn, + ellipK, ellipiK + include("bessel.jl") +include("elliptic.jl") include("erf.jl") include("sincosint.jl") include("gamma.jl") diff --git a/src/elliptic.jl b/src/elliptic.jl new file mode 100644 index 00000000..fcaeffe4 --- /dev/null +++ b/src/elliptic.jl @@ -0,0 +1,245 @@ +# References: +# [1] Abramowitz, Stegun: Handbook of Mathematical Functions (1965) +# [2] ellipjc.m in Toby Driscoll's Schwarz-Christoffel Toolbox + + +#------------------------------------------------ +# Descending and Ascending Landen Transformation + +descstep(m) = m/(1+sqrt(1-m))^2 + +@generated function shrinkm(m,::Val{N}) where {N} + # [1, Sec 16.12] / https://dlmf.nist.gov/22.7.i + quote + f = one(m) + Base.Cartesian.@nexprs $N i->begin + k_i = descstep(m) + m = k_i^2 + f *= 1+k_i + end + return (Base.Cartesian.@ntuple $N k), f, m + end +end + +function ellipj_smallm(u,m) + # [1, Sec 16.13] / https://dlmf.nist.gov/22.10.ii + if VERSION < v"0.7-" + sinu = sin(u) + cosu = cos(u) + else + sinu, cosu = sincos(u) + end + sn = sinu - m*(u-sinu*cosu)*cosu/4 + cn = cosu + m*(u-sinu*cosu)*sinu/4 + dn = 1 - m*sinu^2/2; + return sn,cn,dn +end +function ellipj_largem(u,m1) + # [1, Sec 16.15] / https://dlmf.nist.gov/22.10.ii + sinhu = sinh(u) + coshu = cosh(u) + tanhu = sinhu/coshu + sechu = 1/coshu + sn = tanhu + m1*(sinhu*coshu-u)*sechu^2/4 + cn = sechu - m1*(sinhu*coshu-u)*tanhu*sechu/4 + dn = sechu + m1*(sinhu*coshu+u)*tanhu*sechu/4 + return sn,cn,dn +end + +function ellipj_growm(sn,cn,dn, k) + # [1, Sec 16.12] / https://dlmf.nist.gov/22.7.i + for kk in reverse(k) + sn,cn,dn = (1+kk)*sn/(1+kk*sn^2), + cn*dn/(1+kk*sn^2), + (1-kk*sn^2)/(1+kk*sn^2) + # ^ Use [1, 16.9.1]. Idea taken from [2] + end + return sn,cn,dn +end +function ellipj_shrinkm(sn,cn,dn, k::NTuple{N,<:Any}) where {N} + # [1, Sec 16.14] / https://dlmf.nist.gov/22.7.ii + for kk in reverse(k) + sn,cn,dn = (1+kk)*sn*cn/dn, + (cn^2-kk*sn^2)/dn, # Use [1, 16.9.1] + (cn^2+kk*sn^2)/dn # Use [1, 16.9.1] + end + return sn,cn,dn +end + +function ellipj_viasmallm(u,m,::Val{N}) where {N} + k,f,m = shrinkm(m,Val{N}()) + sn,cn,dn = ellipj_smallm(u/f,m) + return ellipj_growm(sn,cn,dn,k) +end +function ellipj_vialargem(u,m,::Val{N}) where {N} + k,f,m1 = shrinkm(1-m,Val{N}()) + sn,cn,dn = ellipj_largem(u/f,m1) + return ellipj_shrinkm(sn,cn,dn,k) +end + + +#---------------- +# Pick algorithm + +function ellipj_dispatch(u,m, ::Val{N}) where {N} + if abs(m) <= 1 && real(m) <= 0.5 + return ellipj_viasmallm(u,m, Val{N}()) + elseif abs(1-m) <= 1 + return ellipj_vialargem(u,m, Val{N}()) + elseif imag(m) == 0 && real(m) < 0 + # [1, Sec 16.10] + sn,cn,dn = ellipj_dispatch(u*sqrt(1-m),-m/(1-m), Val{N}()) + return sn/(dn*sqrt(1-m)), cn/dn, 1/dn + else + # [1, Sec 16.11] + sn,cn,dn = ellipj_dispatch(u*sqrt(m),1/m, Val{N}()) + return sn/sqrt(m), dn, cn + end +end + +Base.@pure puresqrt(x::Float64) = sqrt(x) +Base.@pure function nsteps(m,ε) + i = 0 + while abs(m) > ε + m = descstep(m)^2 + i += 1 + end + return i +end +Base.@pure nsteps(ε,::Type{<:Real}) = nsteps(0.5,ε) # Guarantees convergence in [-1,0.5] +Base.@pure nsteps(ε,::Type{<:Complex}) = nsteps(0.5+sqrt(3)/2im,ε) # This is heuristic. +function ellipj_nsteps(u,m) + # Compute the number of Landen steps required to reach machine precision. + # For all FloatXX types, this can be done at compile time, while for + # BigFloat this has to be done at runtime. + T = promote_type(typeof(u),typeof(m)) + ε = puresqrt(Float64(eps(real(typeof(m))))) + N = nsteps(ε,typeof(m)) + return ellipj_dispatch(u,m,Val{N}())::NTuple{3,T} +end + + +#----------------------------------- +# Type promotion and special values + +function ellipj_check(u,m) + if isfinite(u) && isfinite(m) + return ellipj_nsteps(u,m) + else + T = promote_type(typeof(u),typeof(m)) + return (T(NaN),T(NaN),T(NaN)) + end +end + +ellipj(u::Real,m::Real) = ellipj_check(promote(float(u),float(m))...) +function ellipj(u::Complex,m::Real) + T = promote_type(real.(typeof.(float.((u,m))))...) + return ellipj_check(convert(Complex{T},u), convert(T,m)) +end +ellipj(u,m::Complex) = ellipj_check(promote(float(u),float(m))...) + +""" + ellipj(u,m) -> sn,cn,dn + +Jacobi elliptic functions `sn`, `cn` and `dn`. + +Convenience function `jpq(u,m)` with `p,q ∈ {s,c,d,n}` are also +provided, but this function is more efficient if more than one elliptic +function with the same arguments is required. +""" +function ellipj end + + +#----------------------- +# Convenience functions + +chars = ("s","c","d") +for (i,p) in enumerate(chars) + pn = Symbol("j"*p*"n") + np = Symbol("jn"*p) + @eval begin + $pn(u,m) = ellipj(u,m)[$i] + $np(u,m) = 1/$pn(u,m) + end +end +for p in (chars...,"n") + pp = Symbol("j"*p*p) + @eval $pp(u,m) = one(promote_type(typeof.(float.((u,m)))...)) +end + +for p in chars, q in chars + p == q && continue + pq = Symbol("j"*p*q) + pn = Symbol(p*"n") + qn = Symbol(q*"n") + @eval begin + function $pq(u::Number,m::Number) + sn,cn,dn = ellipj(u,m) + return $pn/$qn + end + end +end + +for p in (chars...,"n"), q in (chars...,"n") + pq = Symbol("j"*p*q) + @eval begin +""" + $(string($pq))(u,m) + +Jacobi elliptic function `$($p)$($q)`. + +See also `ellipj(u,m)` if more than one Jacobi elliptic function +with the same arguments is required. +""" +function $pq end + end +end + + +#---------------------------------------------- +# Complete elliptic integral of the first kind + +function ellipiK_agm(m) + # [1, Sec 17.6] + T = typeof(m) + m == 0 && return T(Inf) + isnan(m) && return T(NaN) + a,b = one(m),sqrt(m) + while abs(a-b) > 2*eps(abs(a)) + a,b = (a+b)/2,sqrt(a*b) + end + return T(π)/(2*a) # https://github.com/JuliaLang/julia/issues/26324 +end +ellipiK(m) = ellipiK_agm(float(m)) +ellipK(m::Real) = ellipiK(1-m) +function ellipK(m::Complex) + # Make sure we hit the "right" branch of sqrt if imag(m) == 0. + # Here, "right" is defined as being consistent with mpmath. + if imag(m) == 0 + return ellipiK(complex(1-real(m),imag(m))) + else + return ellipiK(1-m) + end +end + +""" + ellipiK(m1) + +Evaluate `ellipK(1-m1)` with better precision for small values of `m1`. +""" +function ellipiK end + +doc""" + ellipK(m) + +Complete elliptic integral of the first kind ``K``. + +```math +\begin{aligned} + K(m) + &= \int_0^1 \big((1-t^2)\,(1-mt^2)\big)^{-1/2} \, dt \\ + &= \int_0^{π/2} (1-m \sin^2\theta)^{-1/2} \, d\theta +\end{aligned} +``` +""" +function ellipK end diff --git a/test/elliptic.jl b/test/elliptic.jl new file mode 100644 index 00000000..dcd6e2b5 --- /dev/null +++ b/test/elliptic.jl @@ -0,0 +1,146 @@ +@testset "elliptic" begin + +setprecision(BigFloat,256) + +function Base.read(s::IO, ::Type{BigFloat}) + setprecision(BigFloat, 256) do + x = BigFloat() + x.sign = Cint(read(s,Int8)) + unsafe_read(s, reinterpret(Ptr{UInt8},x.d), 32) + x.exp = Clong(ntoh(read(s,Int64))) + + if x.sign == 0 + return big(0) + elseif x.sign == 2 + return big(-Inf) + elseif x.sign == 3 + return big(Inf) + elseif x.sign == 4 + return big(NaN) + else + @assert unsafe_load(reinterpret(Ptr{UInt8},x.d), 32) >= 128 + # ^ MPFR crashes if the first bit in the mantissa is not set + return x + end + end +end +Base.read(s::IO, ::Type{Complex{BigFloat}}) = + complex(read(s,BigFloat),read(s,BigFloat)) + +@testset "io" begin + open("elliptic/io.bin","r") do f + @test ntoh(read(f,Int64)) == 0 + @test ntoh(read(f,Int64)) == 1 + @test ntoh(read(f,Int64)) == -1 + @test read(f,BigFloat) == 0 + @test read(f,BigFloat) == 1 + @test read(f,BigFloat) == 2 + @test read(f,BigFloat) == -1 + @test read(f,Complex{BigFloat}) == im + @test read(f,BigFloat) == sqrt(big(2)) + @test read(f,BigFloat) == Inf + @test read(f,BigFloat) == -Inf + @test isnan(read(f,BigFloat)) + end +end + +@testset "ellipj" begin + + @testset "type stability" begin + realtypes = (Int,Float32,Float64,BigFloat) + types = (realtypes..., complex.(realtypes)...) + @testset "typeof(u) = $U, typeof(m) = $M" for U in types, M in types + @inferred ellipj(zero(U),zero(M)) + end + end + + @testset "special values" begin + vals = (0.0,Inf,NaN) + @testset "u = $u, m = $m" for u in vals, m in vals + if !isfinite(u) || !isfinite(m) + @test all(isnan.(ellipj(u,m))) + end + end + end + + @testset "mpmath" begin + open("elliptic/ellipj.bin","r") do f + npts = ntoh(read(f,Int64)) + for i = 1:npts + u = read(f, Complex{BigFloat}) + m = read(f, Complex{BigFloat}) + snref = read(f, Complex{BigFloat}) + cnref = read(f, Complex{BigFloat}) + dnref = read(f, Complex{BigFloat}) + sn,cn,dn = ellipj(u,m) + atol = sqrt(eps(BigFloat))*maximum(abs.((snref,cnref,dnref))) + @test sn ≈ snref atol=atol + @test cn ≈ cnref atol=atol + @test dn ≈ dnref atol=atol + end + end + end + + @testset "precision ($T)" for T in (Float16,Float32,Float64) + n = 16 + s = @. exp(2T(π)*im*(0:n-1)/n) + r = T[eps(T), π*eps(T), sqrt(eps(T)), π*sqrt(eps(T)), 1/π, 0.5, 3/π, 1] + x = Complex{T}[0; vec(s.*r')] + for u = x + for m0 = (0,1) + for m = m0 .+ x + sn,cn,dn = ellipj(u,m) + snref,cnref,dnref = ellipj(big(u),big(m)) + + rtol = 0 + atol = 10*eps(T)*maximum(abs.((snref,cnref,dnref))) + @test sn ≈ snref rtol=rtol atol = atol + @test cn ≈ cnref rtol=rtol atol = atol + @test dn ≈ dnref rtol=rtol atol = atol + end + end + end + end + +end + +@testset "ellipK" begin + @testset "type stability" begin + realtypes = (Int,Float32,Float64,BigFloat) + types = (realtypes..., complex.(realtypes)...) + @testset "typeof(m) = $M" for M in types + @inferred ellipK(zero(M)) + end + end + + @testset "special values" begin + @test isnan(ellipK(NaN)) + @test ellipK(Inf+0im) == 0.0 + @test ellipK(-Inf) == 0.0 + end + + @testset "mpmath" begin + open("elliptic/ellipK.bin","r") do f + npts = ntoh(read(f,Int64)) + for i = 1:npts + m = read(f, Complex{BigFloat}) + ellipKref = read(f, Complex{BigFloat}) + @test ellipK(m) ≈ ellipKref + end + end + end + + @testset "precision ($T)" for T in (Float16,Float32,Float64) + n = 16 + s = @. exp(2T(π)*im*(0:n-1)/n) + r = T[eps(T), π*eps(T), sqrt(eps(T)), π*sqrt(eps(T)), 1/T(π), 0.5, 1, π, 1/sqrt(eps(T)), 1/eps(T)] + x = Complex{T}[0; vec(s.*r')] + for m0 = (0,1) + for m = m0 .+ x + @test ellipK(m) ≈ ellipK(big(m)) rtol=3*eps(T) + end + end + end +end + +end diff --git a/test/elliptic/ellipK.bin b/test/elliptic/ellipK.bin new file mode 100644 index 00000000..5eb25ba1 Binary files /dev/null and b/test/elliptic/ellipK.bin differ diff --git a/test/elliptic/ellipj.bin b/test/elliptic/ellipj.bin new file mode 100644 index 00000000..2313602f Binary files /dev/null and b/test/elliptic/ellipj.bin differ diff --git a/test/elliptic/elliptic.py b/test/elliptic/elliptic.py new file mode 100644 index 00000000..22903430 --- /dev/null +++ b/test/elliptic/elliptic.py @@ -0,0 +1,140 @@ +# Create test set for elliptic functions + +import mpmath + +mpmath.mp.prec = 256 + +def isplusinf(x): return x == mpmath.mp.inf or (x > 0 and x.exp + x.bc >= 2**63) +def isminusinf(x): return isplusinf(-x) +def isinf(x): return isplusinf(x) or isminusinf(x) + +def parse_int8(x): + return x.to_bytes(1, "big", signed=True) + +def parse_int64(x): + return x.to_bytes(8, "big", signed=True) + +def reverse_bits(x,nbits): + y = 0 + for i in range(nbits): + y <<= 1 + y += (x&1) + x >>= 1 + return y + +def parse_mantissa(x): + if x == 0: + return bytes(32) + m = reverse_bits(x.man, x.man.bit_length()) + b = m.to_bytes(32,"big",signed=False) + return bytes(map(lambda b: reverse_bits(b,8), b)) + +def parse_exponent(x): + return parse_int64(x.exp+x.bc) + +def parse_mpf(x): + s = bytes() + if x < 0 and not isminusinf(x): + s = parse_int8(-1) + if x == 0: + s = parse_int8(0) + if x > 0 and not isplusinf(x): + s = parse_int8(1) + if isminusinf(x): + s = parse_int8(2) + x = mpmath.mpf(0) + if isplusinf(x): + s = parse_int8(3) + x = mpmath.mpf(0) + if mpmath.isnan(x): + s = parse_int8(4) + x = mpmath.mpf(0) + return s + parse_mantissa(x) + parse_exponent(x) + +def parse_mpc(x): + return parse_mpf(x.real) + parse_mpf(x.imag) + + +def test_io(): + with open("io.bin","wb") as f: + f.write(parse_int64(0)) + f.write(parse_int64(1)) + f.write(parse_int64(-1)) + f.write(parse_mpf(mpmath.mpf( 0))) + f.write(parse_mpf(mpmath.mpf( 1))) + f.write(parse_mpf(mpmath.mpf( 2))) + f.write(parse_mpf(mpmath.mpf(-1))) + f.write(parse_mpc(mpmath.mpc(1j))) + f.write(parse_mpf(mpmath.sqrt(mpmath.mpf(2)))) + f.write(parse_mpf(mpmath.mp.inf)) + f.write(parse_mpf(-mpmath.mp.inf)) + f.write(parse_mpf(mpmath.mp.nan)) +test_io() + + +def dump_testset(name,u,m): + with open(name,"wb") as f: + f.write(parse_int64(len(u)*len(m))) + for ui in u: + for mi in m: + sn,cn,dn = ellipj(ui,mi) + f.write(parse_mpc(ui)) + f.write(parse_mpc(mi)) + f.write(parse_mpc(sn)) + f.write(parse_mpc(cn)) + f.write(parse_mpc(dn)) + +def ellipj(u,m): + sn = mpmath.ellipfun("sn",u,m=m) + cn = mpmath.ellipfun("cn",u,m=m) + dn = mpmath.ellipfun("dn",u,m=m) + if u == 0: sn = 0*sn # https://github.com/fredrik-johansson/mpmath/issues/386 + return sn,cn,dn + +def test_ellipj(): + s = [ + mpmath.mpc(1,0), + mpmath.sqrt(mpmath.mpc(0,1)), + mpmath.mpc(0,1), + mpmath.mpc(-1,0), + ] + r = [ + mpmath.mpf(mpmath.mp.eps), + mpmath.sqrt(mpmath.mp.eps), + 1/mpmath.mp.e, + mpmath.mpf("1"), + mpmath.mp.e, + ] + x = [mpmath.mpc(0)] + [si*ri for si in s for ri in r] + dump_testset("ellipj.bin",x, x+[1+mi for mi in x]) +test_ellipj() + + +def dump_testset(name,m): + with open(name,"wb") as f: + f.write(parse_int64(len(m))) + for mi in m: + K = mpmath.ellipk(mi) + f.write(parse_mpc(mi)) + f.write(parse_mpc(K)) + +def test_ellipK(): + ns = 16 + s = [ + mpmath.exp(2*mpmath.mp.pi*1j*i/ns) + for i in range(ns) + ] + r = [ + mpmath.mpf(mpmath.mp.eps), + mpmath.sqrt(mpmath.mp.eps), + 1/mpmath.mp.pi, + 1/mpmath.mp.e, + mpmath.mpf("0.5"), + mpmath.mpf("1"), + mpmath.mpf("2"), + mpmath.mp.e, + mpmath.mp.pi, + ] + m = [mpmath.mpc(0)] + [si*ri for si in s for ri in r] + dump_testset("ellipK.bin",m + [1+mi for mi in m]) +test_ellipK() diff --git a/test/elliptic/io.bin b/test/elliptic/io.bin new file mode 100644 index 00000000..595697ef Binary files /dev/null and b/test/elliptic/io.bin differ diff --git a/test/runtests.jl b/test/runtests.jl index c82a4cfa..a361b0e9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -558,3 +558,5 @@ end end @test sprint(showerror, SF.AmosException(1)) == "AmosException with id 1: input error." + +include("elliptic.jl")