Improve accuracy of log2mexp near log(2)

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <tkpapp@gmail.com>"]
-version = "1.0.5"
+version = "1.0.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/basicfuns.jl

@@ -362,7 +362,50 @@ $(SIGNATURES)
 
 Return `log(2 - exp(x))` evaluated carefully.
 """
-log2mexp(x::Real) = log1p(-expm1(x))
+log2mexp(x::Real) = _log2mexp(float(x)) # ensures that Int/Float64 etc. dispatch to the same algorithm
+
+function _log2mexp(x::Real)
+    # Branch 1 computes log1p(u) with u = -expm1(x), which is better conditioned than
+    # log(1 + u) (branch 2) iff u ≥ -1/2, i.e., x ≤ log(3/2). Both branches lose some
+    # accuracy just past log(3/2), where u resp. 2 - exp(x) cross a power of two.
+    if x <= oftype(x, 0.4054651081081644)
+        return log1p(-expm1(x))
+    else
+        # The domain is x ≤ T(log(2)), as before. T(log(2)) can be larger than log(2) (e.g., Float32)
+        # but it is the value of `log(T(2))` and `log2mexp(T(-Inf))`, so it is mapped to -Inf.
+        hi = oftype(x, IrrationalConstants.logtwo)
+        # Use `>` such that NaN propagates (`min(NaN, 0)` is NaN)
+        if x > hi
+            _log2mexp_domainerror(x)
+        end
+        # log(2 - exp(x)) = log(-2 expm1(x - log(2))) is accurate if x - log(2) is accurate.
+        # For x ≤ hi, x - log(2) > 0 only if x == hi > log(2), which is clamped to 0 (result: -Inf).
+        return log(-2 * expm1(min(_minus_logtwo(x, hi), zero(x))))
+    end
+end
+
+# outlined for performance
+@noinline _log2mexp_domainerror(x) = throw(DomainError(x, "log2mexp(x) is only defined for x ≤ log(2)."))
+
+# x - log(2) with Cody-Waite argument reduction: For hi = T(log(2)) and x > log(3/2),
+# x - hi is exact (Sterbenz lemma), so only lo = log(2) - hi needs to be accurate
+_minus_logtwo(x::Real, hi::Real) = (x - hi) - _logtwo_lo(hi)
+
+# MPFR computes x - l with correct rounding, so it suffices to subtract a more precise log(2)
+function _minus_logtwo(x::BigFloat, ::BigFloat)
+    prec = precision(x) + precision(BigFloat) + 32
+    return x - BigFloat(IrrationalConstants.logtwo; precision = prec)
+end
+
+# lo = log(2) - hi for hi = T(log(2))
+# The hard-coded values for Float16, Float32, and Float64 are equal to the generic fallback
+function _logtwo_lo(hi::Real)
+    l = BigFloat(IrrationalConstants.logtwo; precision = 2 * precision(typeof(hi)) + 32)
+    return oftype(hi, l - hi)
+end
+_logtwo_lo(::Float64) = 2.3190468138462996e-17
+_logtwo_lo(::Float32) = -1.9046542f-9
+_logtwo_lo(::Float16) = Float16(-0.0002122)
 
 """
 $(SIGNATURES)

test/basicfuns.jl

@@ -206,6 +206,64 @@ end
         @test iszero(log2mexp(T(0)))
         @test log2mexp(-T(1)) ≈ log(2 - exp(-T(1)))
     end
+    @test @inferred(log2mexp(0)) === log2mexp(0.0)
+    @test @inferred(log2mexp(1//2)) === log2mexp(0.5)
+
+    # hard-coded values of lo = log(2) - T(log(2))
+    for T in (Float16, Float32, Float64)
+        hi = log(T(2))
+        @test LogExpFunctions._logtwo_lo(hi) === invoke(LogExpFunctions._logtwo_lo, Tuple{Real}, hi)
+    end
+
+    # https://github.com/JuliaStats/LogExpFunctions.jl/issues/102
+    @test ulp_error_maximum(log2mexp, filter(<=(log(big(2))), all_finite_float16())) < 1.5
+    for (T, tol) in ((Float32, 1.5), (Float64, 2))
+        # largest floating point number ≤ log(2)
+        top = log(T(2)) <= log(big(2)) ? log(T(2)) : prevfloat(log(T(2)))
+        @test ulp_error_maximum(log2mexp, (prevfloat(top, k) for k in 0:1999)) < 1
+        @test ulp_error_maximum(log2mexp, range(T(log(1.5)), top; length = 10_000)) < tol
+        @test ulp_error_maximum(log2mexp, range(T(-20), T(0.5); length = 10_000)) < tol
+    end
+
+    # special values and domain boundary
+    for T in (Float16, Float32, Float64)
+        hi = log(T(2))
+        @test isnan(log2mexp(T(NaN)))
+        @test log2mexp(T(-Inf)) === hi
+        @test_throws DomainError log2mexp(T(Inf))
+        @test_throws DomainError log2mexp(T(1))
+        @test_throws DomainError log2mexp(nextfloat(hi))
+        if hi > log(big(2))
+            # Float16 and Float32: T(log(2)) is mapped to the boundary, as before
+            @test log2mexp(hi) === T(-Inf)
+            @test log2mexp(log2mexp(T(-Inf))) === T(-Inf)
+        else
+            # Float64: T(log(2)) is in the domain
+            @test isfinite(log2mexp(hi))
+            @test ulp_error(log2mexp, hi) < 1
+        end
+    end
+
+    # BigFloat close to log(2) and to the branch threshold log(3/2)
+    for prec in (64, 113, 256, 1000)
+        setprecision(BigFloat, prec) do
+            l2 = log(big(2))
+            l15 = log(big(1.5))
+            t = big(0.4054651081081644)
+            xs = [
+                prevfloat(l2, 1), prevfloat(l2, 3), prevfloat(l2, 1000),
+                prevfloat(l15), l15, nextfloat(l15), prevfloat(t), t, nextfloat(t),
+            ]
+            for x in xs
+                expected = setprecision(BigFloat, 4 * prec) do
+                    log(2 - exp(BigFloat(x; precision = 4 * prec)))
+                end
+                y = @inferred(log2mexp(x))
+                @test y isa BigFloat
+                @test abs(y - expected) <= 2 * eps(y)
+            end
+        end
+    end
 end
 
 @testset "logexpm1" begin