Faster, more correct complex^complex (#24570)

Author: stevengj

base/complex.jl

@@ -648,98 +648,89 @@ end
 exp10(z::Complex) = exp10(float(z))
 
 # _cpow helper function to avoid method ambiguity with ^(::Complex,::Real)
-function _cpow(z::Complex{T}, p::Union{T,Complex{T}})::Complex{T} where T<:AbstractFloat
-    if p == 2 #square
-        zr, zi = reim(z)
-        x = (zr-zi)*(zr+zi)
-        y = 2*zr*zi
-        if isnan(x)
-            if isinf(y)
-                x = copysign(zero(T),zr)
-            elseif isinf(zi)
-                x = convert(T,-Inf)
-            elseif isinf(zr)
-                x = convert(T,Inf)
+function _cpow(z::Union{T,Complex{T}}, p::Union{T,Complex{T}}) where {T<:AbstractFloat}
+    if isreal(p)
+        pᵣ = real(p)
+        if isinteger(pᵣ) && abs(pᵣ) < typemax(Int32)
+            # |p| < typemax(Int32) serves two purposes: it prevents overflow
+            # when converting p to Int, and it also turns out to be roughly
+            # the crossover point for exp(p*log(z)) or similar to be faster.
+            if iszero(pᵣ) # fix signs of imaginary part for z^0
+                zer = flipsign(copysign(zero(T),pᵣ), imag(z))
+                return Complex(one(T), zer)
             end
-        elseif isnan(y) && isinf(x)
-            y = copysign(zero(T), y)
-        end
-        Complex(x,y)
-    elseif z!=0
-        if p!=0 && isinteger(p)
-            rp = real(p)
-            if rp < 0
-                return power_by_squaring(inv(z), convert(Integer, -rp))
-            else
-                return power_by_squaring(z, convert(Integer, rp))
-            end
-        end
-        exp(p*log(z))
-    elseif p!=0 #0^p
-        zero(z) #CHECK SIGNS
-    else #0^0
-        zer = copysign(zero(T),real(p))*copysign(zero(T),imag(z))
-        Complex(one(T), zer)
-    end
-end
-function _cpow(z::Complex{T}, p::Union{T,Complex{T}}) where T<:Real
-    if isinteger(p)
-        rp = real(p)
-        if rp < 0
-            return power_by_squaring(inv(float(z)), convert(Integer, -rp))
-        else
-            return power_by_squaring(float(z), convert(Integer, rp))
-        end
-    end
-    pr, pim = reim(p)
-    zr, zi = reim(z)
-    r = abs(z)
-    rp = r^pr
-    theta = atan2(zi, zr)
-    ntheta = pr*theta
-    if pim != 0 && r != 0
-        rp = rp*exp(-pim*theta)
-        ntheta = ntheta + pim*log(r)
-    end
-    sinntheta, cosntheta = sincos(ntheta)
-    re, im = rp*cosntheta, rp*sinntheta
-    if isinf(rp)
-        if isnan(re)
-            re = copysign(zero(re), cosntheta)
-        end
-        if isnan(im)
-            im = copysign(zero(im), sinntheta)
-        end
-    end
-
-    # apply some corrections to force known zeros
-    if pim == 0
-        if isinteger(pr)
-            if zi == 0
-                im = copysign(zero(im), im)
-            elseif zr == 0
-                if isinteger(0.5*pr) # pr is even
-                    im = copysign(zero(im), im)
+            ip = convert(Int, pᵣ)
+            if isreal(z)
+                zᵣ = real(z)
+                if ip < 0
+                    iszero(z) && return Complex(T(NaN),T(NaN))
+                    re = power_by_squaring(inv(zᵣ), -ip)
+                    im = -imag(z)
                 else
-                    re = copysign(zero(re), re)
+                    re = power_by_squaring(zᵣ, ip)
+                    im = imag(z)
                 end
+                # slightly tricky to get the correct sign of zero imag. part
+                return Complex(re, ifelse(iseven(ip) & signbit(zᵣ), -im, im))
+            else
+                return ip < 0 ? power_by_squaring(inv(z), -ip) : power_by_squaring(z, ip)
             end
-        else
-            dr = pr*2
-            if isinteger(dr) && zi == 0
-                if zr < 0
-                    re = copysign(zero(re), re)
+        elseif isreal(z)
+            # (note: if both z and p are complex with ±0.0 imaginary parts,
+            #  the sign of the ±0.0 imaginary part of the result is ambiguous)
+            if iszero(real(z))
+                return pᵣ > 0 ? complex(z) : Complex(T(NaN),T(NaN)) # 0 or NaN+NaN*im
+            elseif real(z) > 0
+                return Complex(real(z)^pᵣ, z isa Real ? ifelse(real(z) < 1, -imag(p), imag(p)) : flipsign(imag(z), pᵣ))
+            else
+                zᵣ = real(z)
+                rᵖ = (-zᵣ)^pᵣ
+                if isfinite(pᵣ)
+                    # figuring out the sign of 0.0 when p is a complex number
+                    # with zero imaginary part and integer/2 real part could be
+                    # improved here, but it's not clear if it's worth it…
+                    return rᵖ * complex(cospi(pᵣ), flipsign(sinpi(pᵣ),imag(z)))
                 else
-                    im = copysign(zero(im), im)
+                    iszero(rᵖ) && return zero(Complex{T}) # no way to get correct signs of 0.0
+                    return Complex(T(NaN),T(NaN)) # non-finite phase angle or NaN input
                 end
             end
+        else
+            rᵖ = abs(z)^pᵣ
+            ϕ = pᵣ*angle(z)
         end
+    elseif isreal(z)
+        iszero(z) && return real(p) > 0 ? complex(z) : Complex(T(NaN),T(NaN)) # 0 or NaN+NaN*im
+        zᵣ = real(z)
+        pᵣ, pᵢ = reim(p)
+        if zᵣ > 0
+            rᵖ = zᵣ^pᵣ
+            ϕ = pᵢ*log(zᵣ)
+        else
+            r = -zᵣ
+            θ = copysign(T(π),imag(z))
+            rᵖ = r^pᵣ * exp(-pᵢ*θ)
+            ϕ = pᵣ*θ + pᵢ*log(r)
+        end
+    else
+        pᵣ, pᵢ = reim(p)
+        r = abs(z)
+        θ = angle(z)
+        rᵖ = r^pᵣ * exp(-pᵢ*θ)
+        ϕ = pᵣ*θ + pᵢ*log(r)
     end
 
-    Complex(re, im)
+    if isfinite(ϕ)
+        return rᵖ * cis(ϕ)
+    else
+        iszero(rᵖ) && return zero(Complex{T}) # no way to get correct signs of 0.0
+        return Complex(T(NaN),T(NaN)) # non-finite phase angle or NaN input
+    end
 end
+_cpow(z, p) = _cpow(float(z), float(p))
 ^(z::Complex{T}, p::Complex{T}) where T<:Real = _cpow(z, p)
 ^(z::Complex{T}, p::T) where T<:Real = _cpow(z, p)
+^(z::T, p::Complex{T}) where T<:Real = _cpow(z, p)
 
 ^(z::Complex, n::Bool) = n ? z : one(z)
 ^(z::Complex, n::Integer) = z^Complex(n)
@@ -755,6 +746,10 @@ function ^(z::Complex{T}, p::S) where {T<:Real,S<:Real}
     P = promote_type(T,S)
     return Complex{P}(z) ^ P(p)
 end
+function ^(z::T, p::Complex{S}) where {T<:Real,S<:Real}
+    P = promote_type(T,S)
+    return P(z) ^ Complex{P}(p)
+end
 
 function sin(z::Complex{T}) where T
     F = float(T)

base/mathconstants.jl

@@ -82,7 +82,7 @@ catalan = 0.9159655941772...
 catalan
 
 # loop over types to prevent ambiguities for ^(::Number, x)
-for T in (AbstractIrrational, Rational, Integer, Number)
+for T in (AbstractIrrational, Rational, Integer, Number, Complex)
     Base.:^(::Irrational{:ℯ}, x::T) = exp(x)
 end
 Base.literal_pow(::typeof(^), ::Irrational{:ℯ}, ::Val{p}) where {p} = exp(p)

test/complex.jl

@@ -998,3 +998,74 @@ end
     end
     @test (2+0im)^(-21//10) === (2//1+0im)^(-21//10) === 2^-2.1 - 0.0im
 end
+
+@testset "more cpow" begin
+    # for testing signs of zeros, it is useful to convert ±0.0 to ±1e-15
+    zero2small(r::Real) = iszero(r) ? copysign(1e-15, r) : r
+    zero2small(z::Complex) = complex(zero2small(real(z)), zero2small(imag(z)))
+    ≋(x::Real, y::Real) = x*y == 0 ? abs(x) < 1e-8 && abs(y) < 1e-8 && signbit(x)==signbit(y) : isfinite(x) ? x ≈ y : isequal(x, y)
+    ≋(x::Complex, y::Complex) = real(x) ≋ real(y) && imag(x) ≋ imag(y)
+    ≟(x,y) = isequal(x,y)
+
+    # test z^p for positive/negative/zero real and imaginary parts of z and p:
+    v=(-2.7,-3.0,-2.0,-0.0,+0.0,2.0,3.0,2.7)
+    for zr=v, zi=v, pr=v, pi=v
+        z = complex(zr,zi)
+        p = iszero(pi) ? pr : complex(pr,pi)
+        if isinteger(p)
+            c = zero2small(z)^Integer(pr)
+        else
+            c = exp(zero2small(p) * log(zero2small(z)))
+        end
+        if !iszero(z*p) # z==0 or p==0 is tricky, check it separately
+            @test z^p ≋ c
+            if isreal(p)
+                @test z^(p + 1e-15im) ≈ z^(p - 1e-15im) ≈ c
+                if isinteger(p)
+                    @test isequal(z^Integer(pr), z^p)
+                end
+            elseif (zr != 0 || !signbit(zr)) && (zi != 0 || !signbit(zi))
+                @test isequal((Complex{Int}(z*10)//10)^p, z^p)
+            end
+        end
+    end
+
+    @test 2 ^ (0.3 + 0.0im) === 2.0 ^ (0.3 + 0.0im) === conj(2.0 ^ (0.3 - 0.0im)) ≋  2.0 ^ (0.3 + 1e-15im)
+    @test 0.2 ^ (0.3 + 0.0im) === conj(0.2 ^ (0.3 - 0.0im)) ≋  0.2 ^ (0.3 + 1e-15im)
+    @test (0.0 - 0.0im)^2.0 === (0.0 - 0.0im)^2 === (0.0 - 0.0im)^1.1 === (0.0 - 0.0im) ^ (1.1 + 2.3im) === 0.0 - 0.0im
+    @test (0.0 - 0.0im)^-2.0 ≟ (0.0 - 0.0im)^-2 ≟ (0.0 - 0.0im)^-1.1 ≟ (0.0 - 0.0im) ^ (-1.1 + 2.3im) ≟ NaN + NaN*im
+    @test (1.0+0.0)^(1.2+0.7im) === 1.0 + 0.0im
+    @test (-1.0+0.0)^(2.0+0.7im) ≈ exp(-0.7π)
+    @test (-4.0+0.0im)^1.5 === (-4.0)^(1.5+0.0im) === (-4)^(1.5+0.0im) === (-4)^(3//2+0im) === 0.0 - 8.0im
+
+    # issue #24515:
+    @test (Inf + Inf*im)^2.0 ≟ (Inf + Inf*im)^2 ≟ NaN + Inf*im
+    @test (0+0im)^-3.0 ≟ (0+0im)^-3 ≟ NaN + NaN*im
+    @test (1.0+0.0im)^1e300 === 1.0 + 0.0im
+    @test Inf^(-Inf + 0.0im) == (Inf + 0.0im)^(-Inf - 0.0im) == (Inf - 0.0im)^(-Inf - 0.0im) == (Inf - 0.0im)^-Inf == 0
+
+    # NaN propagation
+    @test (0 + NaN*im)^1 ≟ (0 + NaN*im)^1.0 ≟ (0 + NaN*im)^(1.0+0im) ≟ 0.0 + NaN*im
+    @test (0 + NaN*im)^2 ≟ (0 + NaN*im)^2.0 ≟ (0 + NaN*im)^(2.0+0im) ≟ NaN + NaN*im
+    @test (NaN + 0im)^2.0 ≟ (NaN + 0im)^(2.0+0im) ≟ (2+0im)^NaN ≟ NaN + 0im
+    @test (NaN + 0im)^2.5 ≟ NaN^(2.5+0im) ≟ (NaN + NaN*im)^2.5 ≟ (-2+0im)^NaN ≟ (2+0im)^(1+NaN*im) ≟ NaN + NaN*im
+
+    # more Inf cases:
+    @test (Inf + 0im)^Inf === Inf^(Inf + 0im) === (Inf + 0im)^(Inf + 0im) == Inf + 0im
+    @test (-Inf + 0im)^(0.7 + 0im) === (-Inf + 1im)^(0.7 + 0im) === conj((-Inf - 1im)^(0.7 + 0im)) === -Inf + Inf*im
+    @test (-Inf + 0.0im) ^ 3.1 === conj((-Inf - 0.0im) ^ 3.1) === -Inf - Inf*im
+    @test (3.0+0.0im)^(Inf + 1im) === (3.0-0.0im)^(Inf + 1im) === conj((3.0+0.0im)^(Inf - 1im)) === Inf + Inf*im
+
+    # The following cases should arguably give Inf + Inf*im, but currently
+    # give partial NaNs instead.  Marking as broken for now (since Julia 0.4 at least),
+    # in the hope that someday we can fix these corner cases.  (Python gets them wrong too.)
+    @test_broken (Inf + 1im)^3 === (Inf + 1im)^3.0 === (Inf + 1im)^(3+0im) === Inf + Inf*im
+    @test_broken (Inf + 1im)^3.1 === (Inf + 1im)^(3.1+0im) === Inf + Inf*im
+
+    # cases where phase angle is non-finite yield NaN + NaN*im:
+    @test NaN + NaN*im ≟ Inf ^ (2 + 3im) ≟ (Inf + 1im) ^ (2 + 3im) ≟ (Inf*im) ^ (2 + 3im) ≟
+          3^(Inf*im) ≟ (-3)^(Inf + 0im) ≟ (-3)^(Inf + 1im) ≟ (3+1im)^Inf ≟
+          (3+1im)^(Inf + 1im) ≟ (1e200+1e-200im)^Inf ≟ (1e200+1e-200im)^(Inf+1im)
+
+    @test @inferred(2.0^(3.0+0im)) === @inferred((2.0+0im)^(3.0+0im)) === @inferred((2.0+0im)^3.0) === 8.0+0.0im
+end