From 2996931a4d4a75cd809d5b4d41142538090c9695 Mon Sep 17 00:00:00 2001
From: Tina Torabi <tntorabii@gmail.com>
Date: Tue, 10 Sep 2024 06:19:20 -0700
Subject: [PATCH] up

---
 src/ellipjc.jl | 77 +++++++++++++++++++++++++++++++++-----------------
 1 file changed, 51 insertions(+), 26 deletions(-)

diff --git a/src/ellipjc.jl b/src/ellipjc.jl
index 4007594..1fafc94 100644
--- a/src/ellipjc.jl
+++ b/src/ellipjc.jl
@@ -14,6 +14,28 @@ where `k` is the elliptic modulus.
 # Returns
 - `(sn, cn, dn)`: Corresponding to the Jacobi elliptic functions evaluated at each element of `u`.
 
+# Examples
+```julia
+L = 0.5
+u = 1.0 + 0.5im
+sn, cn, dn = ellipjc(u, L)
+"""
+"""
+    ellipjc(u::Complex, L::Real)
+
+Compute the Jacobi elliptic functions `sn`, `cn`, and `dn` for a complex argument `u` and
+parameter `m = exp(-2*pi*L)`, where `0 < L < Inf`. The relationship `m = k^2` applies,
+where `k` is the elliptic modulus.
+
+# Arguments
+- `u::Complex`: A complex number or an array of complex numbers. Each element should
+  satisfy the condition `|Re(u)| < K` for real part and `0 < Im(u) < Kp` for imaginary
+  part, where `[K, Kp] = ellipkkp(L)`.
+- `L::Real`: A real scalar defining the modulus squared parameter `m`.
+
+# Returns
+- `(sn, cn, dn)`: Corresponding to the Jacobi elliptic functions evaluated at each element of `u`.
+
 # Examples
 ```julia
 L = 0.5
@@ -21,48 +43,51 @@ u = 1.0 + 0.5im
 sn, cn, dn = ellipjc(u, L)
 """
 function ellipjc(u, L; flag=false)
+    u = u isa Array ? u : [u]
+    sn = similar(u, Complex{Float64})
+    cn = similar(u, Complex{Float64})
+    dn = similar(u, Complex{Float64})
+    
     if !flag
-        _, Kp = ellipkkp(L)
-        high = [(imag(x) > real(Kp) / 2) for x in u]
-        if any(high)
-            u = [(imag(x) > real(Kp) / 2) ? im * Kp .- x : x for x in u]
-        end
+        KK, Kp = ellipkkp(L)
+        high = imag.(u) .> real(Kp) / 2
+        u[high] = im * Kp .- u[high]
         m = exp(-2 * pi * L)
     else
-        high = falses(size(u))  
+        high = falses(size(u))
         m = L
     end
+
     if abs(m) < 6eps(Float64)
-        sinu = [sin(x) for x in u]
-        cosu = [cos(x) for x in u]
-        sn = sinu + m / 4 * (sinu .* cosu - u) .* cosu
-        cn = cosu - m / 4 * (sinu .* cosu - u) .* sinu
-        dn = 1 .+ m / 4 .* (cosu .^ 2 - sinu .^ 2 .- 1)
+        sinu = sin.(u)
+        cosu = cos.(u)
+        sn .= sinu + m / 4 * (sinu .* cosu - u) .* cosu
+        cn .= cosu - m / 4 * (sinu .* cosu - u) .* sinu
+        dn .= 1 .+ m / 4 .* (cosu .^ 2 - sinu .^ 2 .- 1)
     else
         if abs(m) > 1e-3
             kappa = (1 - sqrt(Complex(1 - m))) / (1 + sqrt(Complex(1 - m)))
         else
-            kappa = ComplexElliptic.polyval(Complex{Float64}.([132.0, 42.0, 14.0, 5.0, 2.0, 1.0, 0.0]), Complex{Float64}(m / 4.0))
+            kappa = polyval(Complex{Float64}.([132.0, 42.0, 14.0, 5.0, 2.0, 1.0, 0.0]), Complex{Float64}(m / 4.0))
         end
-        sn1, cn1, dn1 = ellipjc(u / (1 + kappa), kappa ^ 2, flag=true)
+        mu = kappa ^ 2
+        v = u ./ (1 + kappa)
+        sn1, cn1, dn1 = ellipjc(v, mu, flag=true)
 
         denom = 1 .+ kappa .* sn1 .^ 2
-        sn = (1 .+ kappa) .* sn1 ./ denom
-        cn = cn1 .* dn1 ./ denom
-        dn = (1 .- kappa .* sn1 .^ 2) ./ denom
+        sn .= (1 .+ kappa) .* sn1 ./ denom
+        cn .= cn1 .* dn1 ./ denom
+        dn .= (1 .- kappa .* sn1 .^ 2) ./ denom
     end
 
     if any(high)
-        snh = Zygote.Buffer(sn)
-        cnh = Zygote.Buffer(cn)
-        dnh = Zygote.Buffer(dn)
-        snh[:] = sn[:]
-        cnh[:] = cn[:]
-        dnh[:] = dn[:]
-        snh[high] = -1 ./ (sqrt(m) * sn[high])
-        cnh[high] = im * dn[high] ./ (sqrt(m) * sn[high])
-        dnh[high] = im * cn[high] ./ sn[high]
+        snh = sn[high]
+        cnh = cn[high]
+        dnh = dn[high]
+        sn[high] .= -1 ./ (sqrt(m) * snh)
+        cn[high] .= im * dnh ./ (sqrt(m) * snh)
+        dn[high] .= im * cnh ./ snh
     end
 
-    return copy(sn), copy(cn), copy(dn)
+    return sn, cn, dn
 end