log1pexp

src/basicfuns.jl

@@ -152,9 +152,61 @@ Return `log(1+exp(x))` evaluated carefully for largish `x`.
 
 This is also called the ["softplus"](https://en.wikipedia.org/wiki/Rectifier_(neural_networks))
 transformation, being a smooth approximation to `max(0,x)`. Its inverse is [`logexpm1`](@ref).
+
+See:
+ * Martin Maechler (2012) [“Accurately Computing log(1 − exp(− |a|))”](http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf)
 """
-log1pexp(x::Real) = x < 18.0 ? log1p(exp(x)) : x < 33.3 ? x + exp(-x) : oftype(exp(-x), x)
-log1pexp(x::Float32) = x < 9.0f0 ? log1p(exp(x)) : x < 16.0f0 ? x + exp(-x) : oftype(exp(-x), x)
+function log1pexp(_x::Real)
+    x = float(_x)
+    x0, x1, x2 = _log1pexp_thresholds(x)
+    if x < x0
+        return exp(x)
+    elseif x < x1
+        return log1p(exp(x))
+    elseif x < x2
+        return x + exp(-x)
+    else
+        return x
+    end
+end
+
+#= The precision of BigFloat cannot be computed from the type only and computing
+thresholds is slow. Therefore prefer version without thresholds in this case. =#
+log1pexp(x::BigFloat) = x > 0 ? x + log1p(exp(-x)) : log1p(exp(x))
+
+#=
+Returns thresholds x0, x1, x2 such that:
+
+    * log1pexp(x) ≈ exp(x) for x ≤ x0
+    * log1pexp(x) ≈ log1p(exp(x)) for x0 < x ≤ x1
+    * log1pexp(x) ≈ x + exp(-x) for x1 < x ≤ x2
+    * log1pexp(x) ≈ x for x > x2
+
+where the tolerances of the approximations are on the order of eps(typeof(x)).
+For types for which `precision(x)` depends only on the type of `x`, the compiler
+should optimize away all computations done here.
+=#
+@inline function _log1pexp_thresholds(x::Real)
+    prec = precision(x)
+    logtwo = oftype(x, IrrationalConstants.logtwo)
+    x0 = -prec * logtwo
+    x1 = (prec - 1) * logtwo / 2
+    x2 = -x0 - log(-x0) * (1 + 1 / x0) # approximate root of e^-x == x * ϵ/2 via asymptotics of Lambert's W function
+    return (x0, x1, x2)
+end
+
+#=
+For common types we hard-code the thresholds to make absolutely sure they are not
+recomputed each time. Also, _log1pexp_thresholds is not completely elided by the
+compiler in Julia 1.0 / 1.6. We round the outputs of the non-hardcoded version:
+
+* _log1pexp_thresholds(::Float64) = (-36.7368005696771, 18.021826694558577, 33.23111882352963)
+* _log1pexp_thresholds(::Float32) = (-16.635532f0, 7.9711924f0, 13.993f0)
+* _log1pexp_thresholds(::Float16) = (-7.625, 3.467, 5.86)
+=#
+@inline _log1pexp_thresholds(::Float64) = (-37e0, 18e0, 33e0) # same as Maechler 2012
+@inline _log1pexp_thresholds(::Float32) = (-17f0, 8f0, 14f0)
+@inline _log1pexp_thresholds(::Float16) = (Float16(-7.6), Float16(3.5), Float16(5.9))
 
 """
 $(SIGNATURES)

test/basicfuns.jl

@@ -110,15 +110,39 @@ end
 # log1pexp, log1mexp, log2mexp & logexpm1
 
 @testset "log1pexp" begin
-    @test log1pexp(2.0)    ≈ log(1.0 + exp(2.0))
-    @test log1pexp(-2.0)   ≈ log(1.0 + exp(-2.0))
-    @test log1pexp(10000)  ≈ 10000.0
-    @test log1pexp(-10000) ≈ 0.0
-
-    @test log1pexp(2f0)      ≈ log(1f0 + exp(2f0))
-    @test log1pexp(-2f0)     ≈ log(1f0 + exp(-2f0))
-    @test log1pexp(10000f0)  ≈ 10000f0
-    @test log1pexp(-10000f0) ≈ 0f0
+    @testset "log1pexp($x::$T)" for T in (Float16, Float32, Float64, BigFloat), x in 1:40
+        @test (@inferred log1pexp(+log(T(x)))) ≈ T(log1p(big(x)))
+        @test (@inferred log1pexp(-log(T(x)))) ≈ T(log1p(1/big(x)))
+    end
+
+    # special values
+    @test (@inferred log1pexp(0)) ≈ log(2)
+    @test (@inferred log1pexp(0f0)) ≈ log(2)
+    @test (@inferred log1pexp(big(0))) ≈ log(2)
+    @test (@inferred log1pexp(+1)) ≈ log1p(ℯ)
+    @test (@inferred log1pexp(-1)) ≈ log1p(ℯ) - 1
+
+    # large arguments
+    @test (@inferred log1pexp(1e4)) ≈ 1e4
+    @test (@inferred log1pexp(1f4)) ≈ 1f4
+    @test iszero(@inferred log1pexp(-1e4))
+    @test iszero(@inferred log1pexp(-1f4))
+
+    # compare to accurate but slower implementation
+    correct_log1pexp(x::Real) = x > 0 ? x + log1p(exp(-x)) : log1p(exp(x))
+    # large range needed to cover all branches, for all floats (from Float16 to BigFloat)
+    @testset "log1pexp($x::$T)" for T in (Float16, Float32, Float64, BigFloat), x in -300:300
+        @test (@inferred log1pexp(T(x))) ≈ T(correct_log1pexp(big(x)))
+    end
+
+    # test BigFloat with multiple precisions
+    @testset "log1pexp($x) prec=$prec" for prec in (10, 20, 50, 100), x in -300:300
+        expected = correct_log1pexp(big(x))
+        setprecision(prec) do
+            actual = @inferred log1pexp(big(float(x)))
+            @test actual ≈ expected
+        end
+    end
 end
 
 @testset "log1mexp" begin

test/chainrules.jl

@@ -57,14 +57,11 @@
         test_rrule(logcosh, x)
     end
 
-    # test all branches of `log1pexp`
-    for x in (-20.9, 15.4, 41.5)
-        test_frule(log1pexp, x)
-        test_rrule(log1pexp, x)
-    end
-    for x in (8.3f0, 12.5f0, 21.2f0)
-        test_frule(log1pexp, x; rtol=1f-3, atol=1f-3)
-        test_rrule(log1pexp, x; rtol=1f-3, atol=1f-3)
+    @testset "log1pexp" begin
+        for absx in (0, 1, 2, 10, 15, 20, 40), x in (-absx, absx)
+            test_scalar(log1pexp, Float64(x))
+            test_scalar(log1pexp, Float32(x); rtol=1f-3, atol=1f-3)
+        end
     end
 
     for x in (-10.2, -3.3, -0.3)