Backport #221 to release-1: avoid Int overflow in signrank and wilcox

.github/workflows/ci.yml

@@ -4,9 +4,11 @@ on:
   pull_request:
     branches:
       - master
+      - 'release-*'
   push:
     branches:
       - master
+      - 'release-*'
     tags: '*'
 
 concurrency:

Project.toml

@@ -1,6 +1,6 @@
 name = "StatsFuns"
 uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
-version = "1.5.2"
+version = "1.5.3"
 
 [deps]
 HypergeometricFunctions = "34004b35-14d8-5ef3-9330-4cdb6864b03a"

src/StatsFuns.jl

@@ -3,7 +3,7 @@ module StatsFuns
 using Base: Math.@horner
 using Reexport: @reexport
 using SpecialFunctions: beta_inc, beta_inc_inv, digamma,
-    erfc, erfcinv, erfcx, gamma_inc, gamma_inc_inv, logbeta, loggamma
+    erfc, erfcinv, erfcx, gamma_inc, gamma_inc_inv, logabsbinomial, logbeta, loggamma
 
 # reexports
 @reexport using IrrationalConstants:

src/distrs/signrank.jl

@@ -11,15 +11,28 @@ the number of subsets of {1,2,...,n-1} summing to W.
 This can be calculated bottom up using dynamic programming.
 
 The i'th element of DP in the j'th outer loop iteration represents:
-the number of ways {1,2,...,j} can sum to W-i+1.
+the probability that a subset of {1,2,...,j} sums to W-i+1,
+i.e. the number of such subsets divided by the 2^j subsets in total.
+
+The counts themselves grow like 2^n and overflow `Int` well before the
+distribution stops being of interest, so the recursion is normalised as it goes
+rather than divided by 2^n at the end. Halving is exact in binary floating
+point, so while every count stays below 2^53 this gives exactly the counts
+divided by 2^n; beyond that the additions round, but the entries stay in [0,1]
+and can never overflow.
  =#
 
-@inline function signrankDP(n, W)
-    DP = zeros(Int, W + 1)
-    DP[W + 1] = 1
+function signrankDP(n, W)
+    DP = zeros(Float64, W + 1)
+    DP[W + 1] = 1.0
     for j in 1:n
         for i in 1:(W + 1 - j)
-            DP[i] += DP[i + j]
+            DP[i] = (DP[i] + DP[i + j]) / 2
+        end
+        # Sums smaller than j cannot contain j, so their counts are unchanged and
+        # the corresponding probabilities only pick up the factor 1/2
+        for i in max(1, W + 2 - j):(W + 1)
+            DP[i] /= 2
         end
     end
     return DP
@@ -38,7 +51,7 @@ function signrankpdf(n::Int, W::Int)
         return signrankpdf(n, W2)
     end
     DP = signrankDP(n, W)
-    return ldexp(float(DP[1]), -n)
+    return DP[1]
 end
 
 function signranklogpdf(n::Int, W::Union{Float64, Int})
@@ -58,7 +71,7 @@ function signrankcdf(n::Int, W::Int)
         return 1.0 - signrankcdf(n, W2)
     end
     DP = signrankDP(n, W)
-    return sum(Base.Fix2(ldexp, -n) ∘ float, DP)
+    return sum(DP)
 end
 
 function signranklogcdf(n::Int, W::Union{Float64, Int})

src/distrs/wilcox.jl

@@ -34,14 +34,20 @@ H. B. Mann, D. R. Whitney. "On a Test of Whether one of Two Random Variables is
 A. Löffler: "Über eine Partition der nat. Zahlen und ihre Anwendung beim U-Test." Wissenschaftliche Zeitschrift der Martin-Luther-Universität Halle-Wittenberg; Mathematisch-Naturwissenschaftliche Reihe, XXXII'83 M, Heft 5, 87–89; available as https://upload.wikimedia.org/wikipedia/commons/f/f5/LoefflerWilcoxonMannWhitneyTest.pdf
 =#
 
-@inline function wilcox_partitions(nx::Int, ny::Int, U::Int)
+#=
+The recurrence below is linear and homogeneous in `pₘ,ₙ`, so seeding it with
+`1 / binomial(nx + ny, nx)` instead of `1` yields the probabilities directly.
+This matters because the counts themselves grow like `binomial(nx + ny, nx)` and overflow
+`Int` from `nx + ny = 68` onwards, whereas the probabilities are bounded by 1.
+=#
+function wilcox_probabilities(nx::Int, ny::Int, U::Int)
     # This internal function expects 0 <= U <= nx * ny / 2
     if !(0 <= U <= (nx * ny) / 2)
-        throw(ArgumentError("`wilcox_partitions(nx, ny, U)` is only implemented for 0 <= U <= (nx * ny) / 2"))
+        throw(ArgumentError("`wilcox_probabilities(nx, ny, U)` is only implemented for 0 <= U <= (nx * ny) / 2"))
     end
 
-    # Due to symmetry, `wilcox_partitions(nx, ny, U) = wilcox_partitions(ny, nx, U)`
-    # Hence for simplicity we only consider the case `wilcox_partitions(min(nx, ny), max(nx, ny), U)` here
+    # Due to symmetry, `wilcox_probabilities(nx, ny, U) = wilcox_probabilities(ny, nx, U)`
+    # Hence for simplicity we only consider the case `wilcox_probabilities(min(nx, ny), max(nx, ny), U)` here
     m, n = minmax(nx, ny)
 
     # Compute σ(k) = ∑_{d|k} ϵ(d) d where
@@ -60,18 +66,30 @@ A. Löffler: "Über eine Partition der nat. Zahlen und ihre Anwendung beim U-Tes
         end
     end
 
-    # Recursively compute the number of partitions pₘ,ₙ(a) for 0 <= a <= U
-    partitions = Vector{Int}(undef, U + 1)
-    partitions[1] = 1
+    # Recursively compute pₘ,ₙ(a) / binomial(nx + ny, nx) for 0 <= a <= U, in units of
+    # 2^shift. The logarithm of the coefficient is used since `binomial(nx + ny, nx)` itself
+    # overflows `Int` from `nx + ny = 68` onwards.
+    #
+    # The seed is the probability of the least likely outcome, so it leaves the normal range
+    # long before the probabilities of interest do: it is subnormal from `nx + ny = 1030` and
+    # zero from `nx + ny = 1090`, and since the recurrence is homogeneous a seed of zero
+    # would silently zero every probability. Working in units of 2^shift avoids that. The
+    # recurrence is homogeneous so the choice of unit cannot affect the result, `ldexp`
+    # removes it from the scalar result exactly, and `shift` is zero whenever the seed is
+    # normal.
+    logbinom = first(logabsbinomial(nx + ny, nx))
+    shift = max(0, ceil(Int, (logbinom - 700) / logtwo))
+    probabilities = Vector{Float64}(undef, U + 1)
+    probabilities[1] = exp(shift * logtwo - logbinom)
     for a in 1:U
-        p = 0
+        p = 0.0
         for i in 1:a
-            p += partitions[i] * sigmas[a + 1 - i]
+            p += probabilities[i] * sigmas[a + 1 - i]
         end
-        partitions[a + 1] = p ÷ a
+        probabilities[a + 1] = p / a
     end
 
-    return partitions
+    return probabilities, shift
 end
 
 function wilcoxpdf(nx::Int, ny::Int, U::Float64)
@@ -83,8 +101,8 @@ function wilcoxpdf(nx::Int, ny::Int, U::Int)
         return 0.0
     end
     U = min(U, max_U - U)
-    partitions = wilcox_partitions(nx, ny, U)
-    return partitions[end] / binomial(nx + ny, nx)
+    probabilities, shift = wilcox_probabilities(nx, ny, U)
+    return ldexp(probabilities[end], -shift)
 end
 
 function wilcoxlogpdf(nx::Int, ny::Int, U::Union{Float64, Int})
@@ -102,8 +120,8 @@ function wilcoxcdf(nx::Int, ny::Int, U::Int)
         return 1.0
     else
         U2 = max_U - U - 1
-        partitions = wilcox_partitions(nx, ny, min(U, U2))
-        p = sum(float, partitions) / binomial(nx + ny, nx)
+        probabilities, shift = wilcox_probabilities(nx, ny, min(U, U2))
+        p = ldexp(sum(probabilities), -shift)
         return U2 < U ? 1.0 - p : p
     end
 end

test/rmath.jl

@@ -1,5 +1,6 @@
 using StatsFuns
 using StatsFuns: RFunctions
+using Rmath: Rmath
 using Test
 
 include("utils.jl")
@@ -27,17 +28,20 @@ function rmathcomp(basename, params, X::AbstractArray, rtol = _default_rtol(para
     end
 
     #=
-    R version of signrank has system variation
+    signrank and wilcox are implemented natively rather than by delegating to Rmath, so
+    their inverse functions are not expected to reproduce the Rmath quantile functions
+    bit for bit. The inverses locate a discontinuity by searching for the first argument
+    at which the cdf reaches `q`, and the `q` tested here are themselves cdf values, so
+    the comparison is decided by the last bit of the cdf and is not stable across
+    platforms. R's own signrank already varies this way:
     julia> psignrank(18,10,false,true) # windows
     -0.2076393647782445
     julia> psignrank(18,10,false,true) # linux
     -0.20763936477824452
-    This slight difference causes test failures for the inverse functions,
-    due to a slight shift in the location of the discontinuity.
-
-    This also holds true for wilcox.
+    As for the other natively implemented discrete distributions, the inverse functions
+    are instead tested by round tripping through our own cdf further down.
     =#
-    test_inv = (basename != "signrank" && basename != "wilcox") || !Sys.islinux()
+    test_inv = basename != "signrank" && basename != "wilcox"
 
     if has_pdf
         pdf = string(basename, "pdf")
@@ -455,6 +459,26 @@ end
     @test isnan(signrankinvlogccdf.(50, signranklogccdf.(50, 1275)))
     @test isnan(signrankinvlogccdf.(50, signranklogccdf.(50, 1276)))
 
+    # The subset counts grow like 2^n, so accumulating them in `Int` silently wrapped
+    # around and returned invalid probabilities from n = 72 onwards (#219)
+    @testset "signrank: large n" begin
+        @testset "n = $n" for n in (71, 72, 73, 74, 80, 100, 120, 200)
+            # The distribution is symmetric about n * (n + 1) / 4
+            W = (n * (n + 1)) ÷ 4
+            @test signrankpdf(n, W) ≈ Rmath.dsignrank(W, n, false)
+            @test signrankcdf(n, W) ≈ Rmath.psignrank(W, n, true, false)
+            @test signrankccdf(n, W) ≈ Rmath.psignrank(W, n, false, false)
+            # Deep in the left tail, where the probabilities are tiny but still normal
+            W2 = W ÷ 3
+            @test signrankpdf(n, W2) ≈ Rmath.dsignrank(W2, n, false)
+            @test signrankcdf(n, W2) ≈ Rmath.psignrank(W2, n, true, false)
+        end
+
+        @testset "pdf sums to one, n = $n" for n in (70, 80, 100)
+            @test sum(signrankpdf(n, W) for W in 0:((n * (n + 1)) ÷ 2)) ≈ 1
+        end
+    end
+
     rmathcomp_tests(
         "srdist", [
             ((1, 2), (0.0:0.2:5.0)),
@@ -488,6 +512,54 @@ end
     @test isnan(wilcoxinvlogccdf.(10, 10, wilcoxlogccdf.(10, 10, 100)))
     @test isnan(wilcoxinvlogccdf.(10, 10, wilcoxlogccdf.(10, 10, 101)))
 
+    # Accumulating the partition counts in `Int` used to break down in two separate ways:
+    # the counts themselves wrapped around and returned silently invalid probabilities (#220),
+    # and their normaliser `binomial(nx + ny, nx)` threw an `OverflowError` from
+    # nx + ny = 68 onwards (#219). The silent limit is the lower of the two, and it is set by
+    # `binomial(nx + ny, nx)` rather than by `nx + ny`: the Löffler recurrence subtracts as
+    # well as adds, so its intermediates overshoot the counts and wrap first. Unbalanced
+    # groups therefore stayed correct much further out, which is why the cases below are not
+    # all balanced.
+    @testset "wilcox: large nx and ny" begin
+        @testset "nx = $nx, ny = $ny" for (nx, ny) in (
+                (32, 32), (33, 33), (34, 34), (40, 40), (50, 50),
+                (30, 36), (20, 48), (20, 60), (10, 325), (10, 326), (5, 200), (150, 40),
+            )
+            # The distribution is symmetric about nx * ny / 2
+            U = (nx * ny) ÷ 2
+            @test wilcoxpdf(nx, ny, U) ≈ Rmath.dwilcox(U, nx, ny, false)
+            @test wilcoxcdf(nx, ny, U) ≈ Rmath.pwilcox(U, nx, ny, true, false)
+            @test wilcoxccdf(nx, ny, U) ≈ Rmath.pwilcox(U, nx, ny, false, false)
+            # Deep in the left tail, where the probabilities are tiny but still normal
+            U2 = U ÷ 3
+            @test wilcoxpdf(nx, ny, U2) ≈ Rmath.dwilcox(U2, nx, ny, false)
+            @test wilcoxcdf(nx, ny, U2) ≈ Rmath.pwilcox(U2, nx, ny, true, false)
+        end
+
+        @testset "pdf sums to one, nx = $nx, ny = $ny" for (nx, ny) in ((30, 30), (33, 33), (34, 34), (50, 50), (5, 200))
+            @test sum(wilcoxpdf(nx, ny, U) for U in 0:(nx * ny)) ≈ 1
+        end
+
+        # `1 / binomial(nx + ny, nx)` is the probability of the least likely outcome and
+        # leaves the normal range from nx + ny = 1030. It seeds the recurrence, which is
+        # homogeneous, so once it reaches zero every probability follows it and the results
+        # are silently zero rather than merely imprecise. The recurrence therefore works in
+        # units of a power of two, which these check has not disturbed anything.
+        @testset "seed below the normal range, nx = ny = $m" for m in (515, 520, 530, 535)
+            # the pdf at U = 0 is exactly the seed, so this pins the scaling down directly
+            @test wilcoxpdf(m, m, 0) ≈ Float64(1 / BigFloat(binomial(big(2m), big(m)), precision = 512))
+        end
+
+        @testset "seed at zero without scaling" begin
+            # `1 / binomial(1080, 524)` is zero in `Float64`, so every one of these was
+            # exactly 0.0 before the recurrence was rescaled
+            @test 0 < wilcoxpdf(524, 556, 20000) < 1
+            @test 0 < wilcoxcdf(524, 556, 20000) < 1
+            @test wilcoxpdf(524, 556, 20000) <= wilcoxcdf(524, 556, 20000)
+            @test wilcoxcdf(524, 556, 20000) < wilcoxcdf(524, 556, 20001)
+        end
+    end
+
     # Note: Convergence fails in srdist with cdf values below 0.16 with df = 10, k = 5.
     # Reduced df or k allows convergence. This test documents this behavior.
     x = 0.15