Generalize logsumexp to complex numbers (#19)

Author: devmotion

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <[email protected]>"]
-version = "0.2.3"
+version = "0.2.4"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

src/logsumexp.jl

@@ -4,8 +4,8 @@ $(SIGNATURES)
 Compute `log(sum(exp, X))` in a numerically stable way that avoids intermediate over- and
 underflow.
 
-`X` should be an iterator of real numbers. The result is computed using a single pass over
-the data.
+`X` should be an iterator of real or complex numbers. The result is computed using a single
+pass over the data.
 
 # References
 
@@ -25,10 +25,10 @@ The result is computed using a single pass over the data.
 
 [Sebastian Nowozin: Streaming Log-sum-exp Computation](http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html)
 """
-logsumexp(X::AbstractArray{<:Real}; dims=:) = _logsumexp(X, dims)
+logsumexp(X::AbstractArray{<:Number}; dims=:) = _logsumexp(X, dims)
 
-_logsumexp(X::AbstractArray{<:Real}, ::Colon) = _logsumexp_onepass(X)
-function _logsumexp(X::AbstractArray{<:Real}, dims)
+_logsumexp(X::AbstractArray{<:Number}, ::Colon) = _logsumexp_onepass(X)
+function _logsumexp(X::AbstractArray{<:Number}, dims)
     # Do not use log(zero(eltype(X))) directly to avoid issues with ForwardDiff (#82)
     FT = float(eltype(X))
     xmax_r = reduce(_logsumexp_onepass_op, X; dims=dims, init=(FT(-Inf), zero(FT)))
@@ -61,36 +61,68 @@ _logsumexp_onepass_reduce(X, ::Base.EltypeUnknown) = reduce(_logsumexp_onepass_o
 ## Reductions for one-pass algorithm: avoid expensive multiplications if numbers are reduced
 
 # reduce two numbers
-function _logsumexp_onepass_op(x1, x2)
-    a = x1 == x2 ? zero(x1 - x2) : -abs(x1 - x2)
-    xmax = x1 > x2 ? oftype(a, x1) : oftype(a, x2)
+function _logsumexp_onepass_op(x1::T, x2::T) where {T<:Number}
+    xmax, a = if x1 == x2
+        # handle `x1 = x2 = ±Inf` correctly
+        x2, zero(x1 - x2)
+    elseif isnan(x1) || isnan(x2)
+        # ensure that `NaN` is propagated correctly for complex numbers
+        z = oftype(x1, NaN)
+        z, exp(z)
+    elseif real(x1) > real(x2)
+        x1, x2 - x1
+    else
+        x2, x1 - x2
+    end
     r = exp(a)
     return xmax, r
 end
+_logsumexp_onepass_op(x1::Number, x2::Number) = _logsumexp_onepass_op(promote(x1, x2)...)
 
 # reduce a number and a partial sum
-function _logsumexp_onepass_op(x, (xmax, r)::Tuple)
-    a = x == xmax ? zero(x - xmax) : -abs(x - xmax)
-    if x > xmax
-        _xmax = oftype(a, x)
-        _r = (r + one(r)) * exp(a)
+_logsumexp_onepass_op(x::Number, (xmax, r)::Tuple{<:Number,<:Number}) =
+    _logsumexp_onepass_op(x, xmax, r)
+_logsumexp_onepass_op((xmax, r)::Tuple{<:Number,<:Number}, x::Number) =
+    _logsumexp_onepass_op(x, xmax, r)
+_logsumexp_onepass_op(x::Number, xmax::Number, r::Number) =
+    _logsumexp_onepass_op(promote(x, xmax)..., r)
+function _logsumexp_onepass_op(x::T, xmax::T, r::Number) where {T<:Number}
+    _xmax, _r = if x == xmax
+        # handle `x = xmax = ±Inf` correctly
+        xmax, r + exp(zero(x - xmax))
+    elseif isnan(x) || isnan(xmax)
+        # ensure that `NaN` is propagated correctly for complex numbers
+        z = oftype(x, NaN)
+        z, r + exp(z)
+    elseif real(x) > real(xmax)
+        x, (r + one(r)) * exp(xmax - x)
     else
-        _xmax = oftype(a, xmax)
-        _r = r + exp(a)
+        xmax, r + exp(x - xmax)
     end
     return _xmax, _r
 end
-_logsumexp_onepass_op(xmax_r::Tuple, x) = _logsumexp_onepass_op(x, xmax_r)
 
 # reduce two partial sums
-function _logsumexp_onepass_op((xmax1, r1)::Tuple, (xmax2, r2)::Tuple)
-    a = xmax1 == xmax2 ? zero(xmax1 - xmax2) : -abs(xmax1 - xmax2)
-    if xmax1 > xmax2
-        xmax = oftype(a, xmax1)
-        r = r1 + (r2 + one(r2)) * exp(a)
+function _logsumexp_onepass_op(
+    (xmax1, r1)::Tuple{<:Number,<:Number}, (xmax2, r2)::Tuple{<:Number,<:Number}
+)
+    return _logsumexp_onepass_op(xmax1, xmax2, r1, r2)
+end
+function _logsumexp_onepass_op(xmax1::Number, xmax2::Number, r1::Number, r2::Number)
+    return _logsumexp_onepass_op(promote(xmax1, xmax2)..., promote(r1, r2)...)
+end
+function _logsumexp_onepass_op(xmax1::T, xmax2::T, r1::R, r2::R) where {T<:Number,R<:Number}
+    xmax, r = if xmax1 == xmax2
+        # handle `xmax1 = xmax2 = ±Inf` correctly
+        xmax2, r2 + (r1 + one(r1)) * exp(zero(xmax1 - xmax2))
+    elseif isnan(xmax1) || isnan(xmax2)
+        # ensure that `NaN` is propagated correctly for complex numbers
+        z = oftype(xmax1, NaN)
+        z, r1 + exp(z)
+    elseif real(xmax1) > real(xmax2)
+        xmax1, r1 + (r2 + one(r2)) * exp(xmax2 - xmax1)
     else
-        xmax = oftype(a, xmax2)
-        r = r2 + (r1 + one(r1)) * exp(a)
+        xmax2, r2 + (r1 + one(r1)) * exp(xmax1 - xmax2)
     end
     return xmax, r
 end

test/basicfuns.jl

@@ -94,18 +94,31 @@ end
     @test logaddexp(2.0, 3.0)     ≈ log(exp(2.0) + exp(3.0))
     @test logaddexp(10002, 10003) ≈ 10000 + logaddexp(2.0, 3.0)
 
-    @test @inferred(logsumexp([1.0])) == 1.0
-    @test @inferred(logsumexp((x for x in [1.0]))) == 1.0
-    @test @inferred(logsumexp([1.0, 2.0, 3.0])) ≈ 3.40760596444438
-    @test @inferred(logsumexp((1.0, 2.0, 3.0))) ≈ 3.40760596444438
-    @test logsumexp([1.0, 2.0, 3.0] .+ 1000.) ≈ 1003.40760596444438
+    for x in ([1.0], Complex{Float64}[1.0])
+        @test @inferred(logsumexp(x)) == 1.0
+        @test @inferred(logsumexp((xi for xi in x))) == 1.0
+    end
 
-    @test @inferred(logsumexp([[1.0, 2.0, 3.0] [1.0, 2.0, 3.0] .+ 1000.]; dims=1)) ≈ [3.40760596444438 1003.40760596444438]
-    @test @inferred(logsumexp([[1.0 2.0 3.0]; [1.0 2.0 3.0] .+ 1000.]; dims=2)) ≈ [3.40760596444438, 1003.40760596444438]
-    @test @inferred(logsumexp([[1.0, 2.0, 3.0] [1.0, 2.0, 3.0] .+ 1000.]; dims=[1,2])) ≈ [1003.4076059644444]
+    for x in ([1.0, 2.0, 3.0], Complex{Float64}[1.0, 2.0, 3.0])
+        @test @inferred(logsumexp(x)) ≈ 3.40760596444438
+        @test logsumexp(x .+ 1000) ≈ 1003.40760596444438
+    end
+
+    for x in ((1.0, 2.0, 3.0), map(complex, (1.0, 2.0, 3.0)))
+        @test @inferred(logsumexp(x)) ≈ 3.40760596444438
+    end
+
+    _x = [[1.0, 2.0, 3.0] [1.0, 2.0, 3.0] .+ 1000.]
+    for x in (_x, complex(_x))
+        @test @inferred(logsumexp(x; dims=1)) ≈ [3.40760596444438 1003.40760596444438]
+        @test @inferred(logsumexp(x; dims=[1, 2])) ≈ [1003.4076059644444]
+        y = copy(x')
+        @test @inferred(logsumexp(y; dims=2)) ≈ [3.40760596444438, 1003.40760596444438]
+    end
 
     # check underflow
     @test logsumexp([1e-20, log(1e-20)]) ≈ 2e-20
+    @test logsumexp(Complex{Float64}[1e-20, log(1e-20)]) ≈ 2e-20
 
     let cases = [([-Inf, -Inf], -Inf),   # correct handling of all -Inf
                  ([-Inf, -Inf32], -Inf), # promotion
@@ -117,6 +130,7 @@ end
         for (arguments, result) in cases
             @test logaddexp(arguments...) ≡ result
             @test logsumexp(arguments) ≡ result
+            @test logsumexp(complex(arguments)) ≡ complex(result)
         end
     end
 
@@ -140,10 +154,26 @@ end
     @test isnan(logsumexp([NaN, 9.0]))
     @test isnan(logsumexp([NaN, Inf]))
     @test isnan(logsumexp([NaN, -Inf]))
+    @test isnan(logsumexp(Complex{Float64}[NaN, 9.0]))
+    @test isnan(logsumexp(Complex{Float64}[NaN, Inf]))
+    @test isnan(logsumexp(Complex{Float64}[NaN, -Inf]))
+    @test isnan(logsumexp(Complex{Float64}[NaN * im, 9.0]))
+    @test isnan(logsumexp(Complex{Float64}[NaN * im, Inf]))
+    @test isnan(logsumexp(Complex{Float64}[NaN * im, -Inf]))
 
     # logsumexp with general iterables (issue #63)
     xs = range(-500, stop = 10, length = 1000)
     @test @inferred(logsumexp(x for x in xs)) == logsumexp(xs)
+    xs = range(-500 + 0.5im, stop = 10 + 30im, length = 1000)
+    @test @inferred(logsumexp(x for x in xs)) == logsumexp(xs)
+
+    # complex numbers
+    xs = randn(Complex{Float64}, 10, 5)
+    @test @inferred(logsumexp(xs)) ≈ log(sum(exp.(xs)))
+    @test @inferred(logsumexp(xs; dims=1)) ≈ log.(sum(exp.(xs); dims=1))
+    @test @inferred(logsumexp(xs; dims=2)) ≈ log.(sum(exp.(xs); dims=2))
+    @test @inferred(logsumexp(xs; dims=[1, 2])) ≈ log(sum(exp.(xs); dims=[1, 2]))
+    @test @inferred(logsumexp(x for x in xs)) == logsumexp(xs)
 end
 
 @testset "softmax" begin