Evaluate fdistlogpdf symmetrically in the two degrees of freedom

ext/StatsFunsChainRulesCoreExt.jl

@@ -36,16 +36,25 @@ ChainRulesCore.@scalar_rule(
 ChainRulesCore.@scalar_rule(
     fdistlogpdf(ν1::Real, ν2::Real, x::Number),
     @setup(
-        xν1 = x * ν1,
-        temp1 = xν1 + ν2,
-        a = (x - 1) / temp1,
-        νsum = ν1 + ν2,
-        di = digamma(νsum / 2),
+        di = digamma((ν1 + ν2) / 2),
+        # `r = u / (1 - u)` for the beta variate `u`, split as in `fdistlogpdf`
+        ν1ν2 = ν1 / ν2,
+        r = ν1ν2 * x,
+        invr = inv(r),
+        log1pr = isinf(r) ? log(ν1ν2) + log(x) : log1p(r),
+        # `(x - 1) / (ν1 * x + ν2)`, in the form that stays finite for a large `x`
+        temp1 = ν1 * x + ν2,
+        a = isinf(temp1) ? (1 - inv(x)) / (ν1 + ν2 / x) : (x - 1) / temp1,
+        ν2a = ν2 * a,
     ),
     (
-        (-log1p(ν2 / xν1) - ν2 * a + di - digamma(ν1 / 2)) / 2,
-        (-log1p(xν1 / ν2) + ν1 * a + di - digamma(ν2 / 2)) / 2,
-        ((ν1 - 2) / x - ν1 * νsum / temp1) / 2,
+        # `log(u)` as in `fdistlogpdf`
+        ((isfinite(invr) ? -log1p(invr) : log(ν1ν2) + log(x) - log1pr) - ν2a + di - digamma(ν1 / 2)) / 2,
+        (-log1pr + ν1 * a + di - digamma(ν2 / 2)) / 2,
+        # `(ν1 - 2) / (2 * x) - ν1 * (ν1 + ν2) / (2 * temp1)`, combined into a single quotient
+        # so that the two terms cannot cancel for a large `ν1`. `ν1 * ν2` would overflow for the
+        # largest degrees of freedom, whereas `ν2 * a` is a ratio of the two
+        -(ν1 * ν2a / 2 + 1) / x,
     ),
 )
 

src/distrs/fdist.jl

@@ -5,17 +5,42 @@ fdistpdf(ν1::Real, ν2::Real, x::Real) = exp(fdistlogpdf(ν1, ν2, x))
 
 fdistlogpdf(ν1::Real, ν2::Real, x::Real) = fdistlogpdf(promote(ν1, ν2, x)...)
 function fdistlogpdf(ν1::T, ν2::T, x::T) where {T <: Real}
-    # we ensure that `log(x)` does not error if `x < 0`
+    a = ν1 / 2
+    b = ν2 / 2
+    lbeta = logbeta(a, b)
+    # in terms of the beta variate `u = ν1 * x / (ν1 * x + ν2)` the density is
+    # `u^a * (1 - u)^b / (x * beta(a, b))`; `r` below is `u / (1 - u)`, formed so that a large
+    # `x` does not overflow it
     ν1ν2 = ν1 / ν2
-    y = max(x, 0)
-    val = (xlogy(ν1, ν1ν2) + xlogy(ν1 - 2, y) - xlogy(ν1 + ν2, 1 + ν1ν2 * y)) / 2 - logbeta(ν1 / 2, ν2 / 2)
-    return x < 0 ? oftype(val, -Inf) : val
+    r = ν1ν2 * x
+    invr = inv(r)
+    return if x < 0
+        # outside of the support, where `log(x)` would error as well
+        oftype(lbeta, -Inf)
+    elseif isfinite(invr)
+        # `log(u) = -log1p(1 / r)` and `log(1 - u) = -log1p(r)`, symmetric in the two degrees
+        # of freedom, so no two terms growing like `ν1 * log(ν1)` cancel. `r` overflows for the
+        # largest `x`, where `log1p(r)` is ordinary and both summands below are positive
+        blog1pr = isinf(r) ? xlogy(b, ν1ν2) + xlogy(b, x) : xlog1py(b, r)
+        -xlog1py(a, invr) - blog1pr - log(x) - lbeta
+    else
+        # `r` is zero or subnormal: `x == 0`, `r` underflowing, or `NaN`. The textbook form
+        # covers all three - `log1p(r)` keeps a tiny `r` that `log(1 + r)` would absorb, and
+        # splitting `log(u)` into `log(ν1 / ν2) + log(x)` needs no precision from `r`
+        xlogy(a, ν1ν2) + xlogy(a - 1, x) - xlog1py(a + b, r) - lbeta
+    end
 end
 
 for f in ("cdf", "ccdf", "logcdf", "logccdf")
     ff = Symbol("fdist" * f)
     bf = Symbol("beta" * f)
-    @eval $ff(ν1::T, ν2::T, x::T) where {T <: Real} = $bf(ν1 / 2, ν2 / 2, inv(1 + ν2 / (ν1 * max(0, x))))
+    @eval function $ff(ν1::T, ν2::T, x::T) where {T <: Real}
+        # the beta variate `u = r / (1 + r)`, clamped to the support. Here an overflowing
+        # `ν1 * x` only saturates `u` to 1, so the ratio is formed that way round
+        r = (ν1 * max(0, x)) / ν2
+        u = r > 1 ? inv(1 + inv(r)) : r / (1 + r)
+        return $bf(ν1 / 2, ν2 / 2, u)
+    end
     @eval $ff(ν1::Real, ν2::Real, x::Real) = $ff(promote(ν1, ν2, x)...)
 end
 for f in ("invcdf", "invccdf", "invlogcdf", "invlogccdf")

test/chainrules.jl

@@ -71,3 +71,50 @@
         test_rrule(owens_t, h, a)
     end
 end
+
+@testitem "ChainRules fdistlogpdf with small x" tags = [:autodiff] begin
+    using StatsFuns, Test
+    using ChainRulesTestUtils
+    using Random
+
+    # `∂/∂ν1` contains `log(u)`, which used to overflow to `-Inf` for small `x`. The tangent of
+    # `x` is scaled to `x` itself, since the default tangent is of order one and its finite
+    # differences would leave the domain. Only the `frule` can be tested this way: `test_rrule`
+    # differences each argument on its own and ignores the tangents, but both rules are
+    # generated by the same `@scalar_rule`.
+    Random.seed!(1234)
+    @testset "(ν1, ν2, x) = ($ν1, $ν2, $x)" for (ν1, ν2, x) in
+        ((3.0, 7.0, 1.0e-308), (1.0, 1.0e18, 1.0e-291))
+
+        test_frule(fdistlogpdf, ν1, ν2, x ⊢ x / 10)
+    end
+end
+
+@testitem "ChainRules fdistlogpdf with large ν1" setup = [FDistRef] tags = [:autodiff] begin
+    using StatsFuns, Test
+    using ChainRulesCore
+    using SpecialFunctions: digamma
+
+    # `∂/∂x` was a difference of two terms growing like `ν1 / x`, which cancelled completely:
+    # at `(1e18, 5, 2)` it returned `0.0`. `test_frule` cannot see this, since a finite
+    # difference in `ν1` is lost in the rounding of `ν1` itself, so the pullback is evaluated
+    # directly. It returns the three partials separately, whereas `frule` would sum them.
+    # The last three are where the single quotient's own `ν1 * ν2` would overflow
+    @testset "(ν1, ν2, x) = ($ν1, $ν2, $x)" for (ν1, ν2, x) in (
+            (1.0e8, 5.0, 2.0), (1.0e18, 5.0, 2.0), (1.0e20, 1.0, 0.5),
+            (4.0, 100.0, floatmax(Float64)),
+            (1.0e200, 1.0e300, 1.0e200), (1.0e300, 1.0e200, 2.0), (1.0e200, 1.0e200, 2.0),
+        )
+        _, pullback = rrule(fdistlogpdf, ν1, ν2, x)
+        @test pullback(1.0)[4] ≈ FDistRef.dlogpdf_dx(ν1, ν2, x) rtol = 1.0e-12
+    end
+
+    # `(x - 1) / (ν1 * x + ν2)` was `NaN` at `x = Inf`, taking `∂/∂ν1` and `∂/∂ν2` with it
+    @testset "x = Inf" begin
+        _, pullback = rrule(fdistlogpdf, 3.0, 7.0, Inf)
+        _, ∂ν1, ∂ν2, ∂x = pullback(1.0)
+        @test ∂ν1 ≈ (digamma(5.0) - digamma(1.5) - 7 / 3) / 2
+        @test ∂ν2 == -Inf
+        @test ∂x == 0
+    end
+end

test/misc.jl

@@ -97,4 +97,114 @@
     @testset "tvdistinvcdf: numerical issue" begin
         @test isnan(@inferred(tdistinvcdf(0, 0.975)))
     end
+
+    # https://github.com/JuliaStats/StatsFuns.jl/issues/224
+    @testset "fdist: boundary values" begin
+        # outside the support, where the beta variate is not defined
+        @test @inferred(fdistlogpdf(3, 7, -1.0)) == -Inf
+        @test @inferred(fdistlogpdf(3, 7, -Inf)) == -Inf
+        @test @inferred(fdistpdf(3, 7, -1.0)) == 0
+
+        # the previous formulation of the density subtracted two infinities here
+        @test @inferred(fdistlogpdf(3, 7, Inf)) == -Inf
+        @test @inferred(fdistpdf(3, 7, Inf)) == 0
+        @test @inferred(fdistcdf(3, 7, Inf)) == 1
+        @test @inferred(fdistccdf(3, 7, Inf)) == 0
+
+        # at zero the density behaves like `x^(ν1 / 2 - 1)`; `-0.0` is not `< 0` and takes the
+        # same branch as `0.0`
+        @test @inferred(fdistlogpdf(1, 1, 0.0)) == Inf
+        @test @inferred(fdistlogpdf(2, 1, 0.0)) == 0
+        @test @inferred(fdistlogpdf(10, 3, 0.0)) == -Inf
+        @test @inferred(fdistlogpdf(10, 3, -0.0)) == -Inf
+        @test @inferred(fdistcdf(3, 7, 0.0)) == 0
+
+        @test isnan(@inferred(fdistlogpdf(3, 7, NaN)))
+        @test isnan(@inferred(fdistpdf(3, 7, NaN)))
+        @test isnan(@inferred(fdistcdf(3, 7, NaN)))
+    end
+
+    # `x < 0`, `x == 0` and `x > 0` take three different branches of `fdistlogpdf`, which all
+    # have to return `float(promote_type(...))` for every mix of argument types
+    @testset "fdist: type stability" begin
+        @testset "(ν1, ν2, x) = ($ν1, $ν2, $x)" for (ν1, ν2) in
+                ((3, 7), (3, 7.0f0), (3 // 1, 7), (3.0, Float16(7)), (big(3), 7)),
+                x in (-1, 0, 2)
+
+            T = float(Base.promote_typeof(ν1, ν2, x))
+            @test @inferred(fdistlogpdf(ν1, ν2, x)) isa T
+            @test @inferred(fdistpdf(ν1, ν2, x)) isa T
+        end
+
+        # `betacdf` and friends do not support `BigFloat`
+        @testset "(ν1, ν2, x) = ($ν1, $ν2, $x)" for (ν1, ν2) in
+                ((3, 7), (3, 7.0f0), (3 // 1, 7), (3.0, Float16(7))),
+                x in (-1, 0, 2)
+
+            T = float(Base.promote_typeof(ν1, ν2, x))
+            @test @inferred(fdistcdf(ν1, ν2, x)) isa T
+            @test @inferred(fdistccdf(ν1, ν2, x)) isa T
+            @test @inferred(fdistlogcdf(ν1, ν2, x)) isa T
+            @test @inferred(fdistlogccdf(ν1, ν2, x)) isa T
+        end
+    end
+end
+
+# https://github.com/JuliaStats/StatsFuns.jl/issues/224
+@testitem "fdist: numerical issue" setup = [FDistRef] begin
+    using StatsFuns
+    using Test
+
+    # the density has to stay accurate when either degree of freedom is huge: for large `ν2`
+    # the term `ν1 * x / ν2` is tiny and must not be absorbed by an explicit `1 +`, and for
+    # large `ν1` no two terms growing like `ν1 * log(ν1)` may cancel
+    νs = (1.0e9, 1.0e12, 1.0e15, 1.0e18, 1.0e20)
+    @testset "fdistlogpdf with (ν1, ν2) = ($ν1, $ν2)" for (ν1, ν2) in
+        Iterators.flatten(((ν, 5.0), (1.0, ν)) for ν in νs)
+
+        # `x = ν2 / ν1` is the ridge where the textbook form cancels, the two extremes are
+        # where the ratio `ν1 * x / ν2` under- and overflows
+        @testset "x = $x" for x in
+            (1.0e-300, 1.0e-8, 0.5, 2.0, 6.6, 1.0e8, 1.0e300, 0.99 * ν2 / ν1, ν2 / ν1)
+
+            @test @inferred(fdistlogpdf(ν1, ν2, x))::Float64 ≈ FDistRef.logpdf(ν1, ν2, x) rtol = 1.0e-12
+        end
+    end
+
+    # with both degrees of freedom of `floatmax` scale the product `ν1 * x` overflows where the
+    # ratio `(ν1 / ν2) * x` and the density itself are ordinary
+    @testset "huge (ν1, ν2) = ($ν1, $ν2)" for (ν1, ν2) in ((1.0e200, 1.0e300), (1.0e300, 1.0e200))
+        @testset "x = $x" for x in
+            (1.0e-300, 1.0e-8, 0.5, 2.0, 1.0e8, 1.0e200, floatmax(Float64), ν2 / ν1)
+
+            @test @inferred(fdistlogpdf(ν1, ν2, x))::Float64 ≈ FDistRef.logpdf(ν1, ν2, x) rtol = 1.0e-12
+        end
+    end
+
+    # `ν2 / (ν1 * x)` overflows for `x < ν2 / (ν1 * floatmax(T))` - an ordinary argument in
+    # `Float16` - and used to take the density and the cdf with it
+    @testset "small x = $x" for x in (1.0e-300, 1.0f-38, Float16(1.0e-4))
+        T = typeof(x)
+        ν1, ν2 = one(T), T(7)
+        @test @inferred(fdistlogpdf(ν1, ν2, x))::T ≈ FDistRef.logpdf(ν1, ν2, x) rtol = eps(T)^(3 // 4)
+
+        ref = FDistRef.logcdf_smallx(ν1, ν2, x)
+        @test @inferred(fdistlogcdf(ν1, ν2, x))::T ≈ ref rtol = eps(T)^(2 // 3)
+        # `exp` turns the absolute error of the log-cdf into a relative one
+        @test @inferred(fdistcdf(ν1, ν2, x))::T ≈ T(exp(ref)) rtol = eps(T)^(2 // 3) * max(1, abs(Float64(ref)))
+    end
+
+    # symmetrically, `ν1 * x` overflows for the largest `x`
+    @testset "large x = $x" for x in (floatmax(Float64), floatmax(Float32), floatmax(Float16))
+        T = typeof(x)
+        ν1, ν2 = T(4), T(100)
+        @test @inferred(fdistlogpdf(ν1, ν2, x))::T ≈ T(FDistRef.logpdf(ν1, ν2, x)) rtol = eps(T)^(3 // 4)
+    end
+
+    # the beta variate of the cdf overflows for a huge `ν2` if formed as `ν2 / (ν1 * x)`, and for a
+    # huge `x` if formed as `(ν1 * x) / (ν1 * x + ν2)` - the latter takes `u` to `0`, not `1`
+    @test fdistcdf(1.0, floatmax(Float64), floatmax(Float64)) ==
+        betacdf(0.5, floatmax(Float64) / 2, 0.5)
+    @test fdistcdf(1.0, 1.0e308, 1.0e308) == 1
+    @test fdistcdf(1.0, 1.0e308, 0.9e308) == 1
 end

test/utils.jl

@@ -168,3 +168,37 @@
         end
     end
 end
+
+@testmodule FDistRef begin
+    using SpecialFunctions: logbeta
+
+    # textbook formulas for `fdist`, evaluated in `BigFloat`. They cancel badly in `Float64` -
+    # that is precisely what is being tested - and for degrees of freedom of `floatmax` scale the
+    # cancellation runs to 300 digits, well past the 77 that the default 256 bits give. At 256 bits
+    # `logpdf(1e200, 1e300, 0.5)` is off by a relative 3.6e3 and would validate anything
+    function logpdf(ν1::Real, ν2::Real, x::Real)
+        return setprecision(BigFloat, 1024) do
+            _ν1, _ν2, _x = BigFloat(ν1), BigFloat(ν2), BigFloat(x)
+            a, b = _ν1 / 2, _ν2 / 2
+            a * log(_ν1 / _ν2) + (a - 1) * log(_x) - (a + b) * log1p(_ν1 * _x / _ν2) - logbeta(a, b)
+        end
+    end
+
+    function dlogpdf_dx(ν1::Real, ν2::Real, x::Real)
+        return setprecision(BigFloat, 1024) do
+            _ν1, _ν2, _x = BigFloat(ν1), BigFloat(ν2), BigFloat(x)
+            (_ν1 / 2 - 1) / _x - ((_ν1 + _ν2) / 2) * _ν1 / (_ν1 * _x + _ν2)
+        end
+    end
+
+    # for `x` so small that the beta variate `u = ν1 * x / (ν1 * x + ν2)` is negligible the cdf
+    # is `u^a / (a * beta(a, b)) * (1 + O(u))`, which needs no incomplete beta function
+    function logcdf_smallx(ν1::Real, ν2::Real, x::Real)
+        return setprecision(BigFloat, 1024) do
+            _ν1, _ν2, _x = BigFloat(ν1), BigFloat(ν2), BigFloat(x)
+            a, b = _ν1 / 2, _ν2 / 2
+            u = _ν1 * _x / (_ν1 * _x + _ν2)
+            a * log(u) - log(a) - logbeta(a, b)
+        end
+    end
+end