Improved Matlab compatibility

README.md

@@ -76,6 +76,10 @@ also provided. `ellipj(u,m)` is equivalent to `sn(u,m), cn(u,m), dn(u,m)`,
 but faster if you want all three. Likewise, `ellipke(m)` is equivalent to
 `K(m), E(m)`, but faster if you want both.
 
+Additionally, you may also input arrays of the compatible size to calculate elliptic integrals 
+for multiple values together. Please check the documentation of `ellipke` and `ellipj` for details.
+Example usages are also available in `test/runtests.jl" (see the "MATLAB compatibility" testset).
+
 ```jlcon
 julia> import Elliptic
 
@@ -84,6 +88,17 @@ julia> k,e = Elliptic.ellipke(0.5)
 
 julia> sn,cn,dn = Elliptic.ellipj(0.672, 0.36)
 (0.6095196917919022,0.792770928653356,0.9307281387786907)
+
+julia> using Elliptic
+
+julia> ellipke([0, 0.3, 0.7, 1.0])
+([1.5707963267948968, 1.7138894481787914, 2.0753631352924686, Inf], 
+ [1.5707963267948968, 1.4453630644126656, 1.2416705679458224, 1.0])
+
+julia> ellipj([-1.5, 1.5], 0.23)
+([-0.9881951524605244, 0.9881951524605244], 
+ [0.15320032850330667, 0.15320032850330667], 
+ [0.8805669641488431, 0.8805669641488431])
 ```
 
 Installation

src/Elliptic.jl

@@ -36,11 +36,14 @@ function _E(sinphi::Float64, m::Float64)
 end
 E(phi::Real, m::Real) = E(Float64(phi), Float64(m))
 
-
 """
-`ellipke(m::Real)`
-returns `(K(m), E(m))` for scalar `0 ≤ m ≤ 1`
+    ellipke(M)
+
+Complete elliptic integrals of first and second kind. Each element of `M`, say `m`, should satisfy `m ≤ 1.0`.
+A tuple `(K, E)` is returned. If `M` is an array, then `K` and `E` share the same size with `M`.
 """
+function ellipke end
+
 function ellipke(m::Float64)
     if m < 1.
         y = 1. - m
@@ -58,6 +61,15 @@ function ellipke(m::Float64)
 end
 ellipke(x::Real) = ellipke(Float64(x))
 
+function ellipke(M::AbstractArray{T}) where T <: Real
+    K = similar(M, Float64)
+    E = similar(M, Float64)
+    for i in eachindex(M)
+        @inbounds K[i], E[i] = ellipke(M[i])
+    end
+    return K, E
+end
+
 E(m::Float64) = ellipke(m)[2]
 E(x::Float32) = Float32(E(Float64(x)))
 E(x::Real) = E(Float64(x))
@@ -122,6 +134,18 @@ end
 Pi(n::Real, phi::Real, m::Real) = Pi(Float64(n), Float64(phi), Float64(m))
 Π = Pi
 
+"""
+    ellipj(U, M)
+
+Compute the Jacobi elliptic functions `SN`, `CN`, and `DN` evaluated for corresponding 
+elements of argument `U` and parameter `M`.
+Note that `M` can only take values in range `[0, 1]` (both ends closed).
+If both `U` and `M` are arrays, they must be the same size, and each element of the returned tuple `(SN, CN, DN)`
+also has the same size. If either `U` or `M` is an array, then the other (a scalar) is broadcasted, and 
+each element of the returned tuple `(SN, CN, DN)` has the same size as the array.
+"""
+function ellipj end
+
 function ellipj(u::Float64, m::Float64, tol::Float64)
     phi = Jacobi.am(u, m, tol)
     s = sin(phi)
@@ -132,4 +156,16 @@ end
 ellipj(u::Float64, m::Float64) = ellipj(u, m, eps(Float64))
 ellipj(u::Real, m::Real) = ellipj(Float64(u), Float64(m))
 
+function ellipj(U::AbstractArray{TU}, M::AbstractArray{TM}) where {TU <: Real,TM <: Real}
+    @assert size(U) == size(M) "U and M must be the same size"
+    SN = similar(U, Float64)
+    CN = similar(U, Float64)
+    DN = similar(U, Float64)
+    for i in eachindex(U)
+        @inbounds SN[i], CN[i], DN[i] = ellipj(U[i], M[i])
+    end
+    return SN, CN, DN
+end
+ellipj(U::AbstractArray{TU}, m::Real) where TU <: Real = ellipj(U, fill(Float64(m), size(U)))
+ellipj(u::Real, M::AbstractArray{TM}) where TM <: Real = ellipj(fill(Float64(u), size(M)), M)
 end # module

test/runtests.jl

@@ -717,4 +717,32 @@ end
     end
 end
 
+@testset "MATLAB compatibility" begin
+    # compare the results with the MATLAB results
+    Base.isapprox(x::Tuple, y::Tuple; kws...) = isapprox(collect(x), collect(y); kws...)
+    @testset "ellipke" begin
+        @test ellipke(0.5) ≈ (1.854074677301372, 1.350643881047675)
+        @test ellipke([0, 0.3, 0.7, 1.0]) ≈ ([1.570796326794897, 1.713889448178791, 2.075363135292469, Inf],
+            [1.570796326794897, 1.445363064412665, 1.241670567945823, 1.0])
+    end
+    @testset "ellipj" begin
+        @test ellipj(1.2, 0.5) ≈ (0.887715488619278, 0.460392453527896, 0.778447561260692)
+        S, C, D = ellipj(-1.0, [0.0, 0.7, 1.0])
+        @test S ≈ [-0.841470984807897,  -0.786762346175161,  -0.761594155955765]
+        @test C ≈ [0.540302305868140,   0.617256033296522,   0.648054273663885]
+        @test D ≈ [1.0,  0.752797122370078,   0.648054273663885]
+        S, C, D = ellipj([-1.5, 1.5], 0.23)
+        @test S ≈ [-0.988195152460524,  0.988195152460524]
+        @test C ≈ [0.153200328503307,   0.153200328503307]
+        @test D ≈ [0.880566964148843,   0.880566964148843]
+        S, C, D = ellipj([-1.5 1.5; 0.1 9.9], [0.1 0.3; 0.5 0.9])
+        @test S ≈ [-0.994313629062765   0.983958114459644
+                    0.099750685474625  -0.391507114124818]
+        @test C ≈ [0.106491347348199   0.178399632816390
+                    0.995012462609058   0.920175080943652]
+        @test D ≈ [0.949280801821044   0.842346679637729
+                    0.997509348514424   0.928466456922753]
+    end
+end
+
 include("jacobi_tests.jl")