add owens_t function

[stevengj]

Closes #99, with a port of the well-tested and fairly standard Float64 implementation from Boost. Update: including a chain rule for the derivatives.

(Various other algorithms were proposed in that issue, but I think it makes sense to start with a trusted baseline.)

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 75.39%. Comparing base (f8d7f5b) to head (888ab98).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #214      +/-   ##
==========================================
+ Coverage   72.39%   75.39%   +3.00%     
==========================================
  Files          22       23       +1     
  Lines        1007     1130     +123     
==========================================
+ Hits          729      852     +123     
  Misses        278      278              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[stevengj]

Unfortunately, the Boost tests don't seem to cover all of the branches, so I'll need to add more.

[stevengj]

CI failure seems unrelated:

in expression starting at D:\a\StatsFuns.jl\StatsFuns.jl\test\chainrules.jl:6
in expression starting at D:\a\StatsFuns.jl\StatsFuns.jl\test\runtests.jl:3
non-integer x = 1.100000non-integer x = 1.100000convergence failed in 'qtukey'

But that reminds me that it is probably worth adding a chain rule for owens_t, since there are simple formulas for the derivatives with respect to both arguments.

[stevengj]

Again, unrelated CI failure?

RNG of the outermost testset: Xoshiro(0x2ef1b0668e912b14, 0x91795677434389ce, 0xb8c8de0393cd7263, 0xb2f189b4e6006935, 0x63e4cdd08561e4e0)
ERROR: LoadError: Some tests did not pass: 109 passed, 0 failed, 1 errored, 0 broken.
in expression starting at /home/runner/work/StatsFuns.jl/StatsFuns.jl/test/chainrules.jl:6
in expression starting at /home/runner/work/StatsFuns.jl/StatsFuns.jl/test/runtests.jl:3
non-integer x = 1.100000non-integer x = 1.100000convergence failed in 'qtukey
non-integer x = -2.300000non-integer x = 1.200000non-integer x = 5.400000non-integer x = 11.900000non-integer x = -2.300000non-integer x = 1.200000non-integer x = 5.400000non-integer x = 11.900000non-integer x = -2.300000non-integer x = 1.200000non-integer x = 5.400000non-integer x = 11.900000non-integer x = -2.300000non-integer x = 1.200000non-integer x = 5.400000non-integer x = 11.900000non-integer x = -2.500000non-integer x = 1.300000non-integer x = 8.700000non-integer x = -2.500000non-integer x = 1.300000non-integer x = 8.700000non-integer x = -2.500000non-integer x = 1.300000non-integer x = 8.700000non-integer x = -2.500000non-integer x = 1.300000non-integer x = 8.700000Package StatsFuns errored during testing
Error: Process completed with exit code 1.

It seems like one of the tests has some non-negligible probability of failing?

[devmotion]

That seems to indicate that some of the partials/tangents are generated incorrectly (or at least in a numerically problematic way) in the ChainRules tests, or some of the derivative rules have numerically problematic regimens (from the posted snipped it's unclear whether the rule or finite differencing would be to blame).

[stevengj]

(To be clear, the occasional test failures are for other functions, unrelated to owens_t introduced here.)

[devmotion]

Unfortunately, the Boost tests don't seem to cover all of the branches, so I'll need to add more.

Do you plan to add more tests, or is this PR ready for review?

[stevengj]

I've added tests to cover the branches that the Boost tests missed. I think it is ready for review?

[devmotion]

I was curious and checked the accuracy against Arb:

The reference is Arb's verified numerical integration of the integral representation

T(h, a) = 1/(2π) ∫₀ᵃ exp(-h²(1+x²)/2)/(1+x²) dx

Setup. 518 (h, a) points covering all dispatch paths (T1–T6, the a > 1 remapping, a = 1, h = 0, and a few sign-symmetric points), with h up to 20 and a up to 100. References were computed at 512-bit precision; the largest relative enclosure radius was 3.9e-31 (eps(Float64) = 2.2e-16), so the references are reliable.

Float64 results.

• relative error: max 1.49e-14, mean 4.28e-16, median 1.53e-16
• ULP error: max 118, mean 2.8, median 1
• the largest errors occur at h = 20, in the a = 1 and a > 1 branches
• the series methods T1–T6 stay within 9 ULPs

Per dispatch path:

path              count   max relerr  max ULPs
T1                  173     7.78e-16         7
T2                   30     6.42e-16         3
T3                   30     6.67e-16         5
T4                  108     1.51e-15         9
T5                   32     4.83e-16         3
T6                    8     4.02e-16         3
a > 1, h <= 0.67     48     3.99e-16         2
a > 1, h > 0.67      60     1.49e-14       118
h = 0                 2     0.00e+00         0
|a| = 1              27     1.49e-14       118

Float32 / Float16. These call the Float64 path and round, so their relative error against Arb is at the level of eps(Float32) ≈ 1e-7 and eps(Float16) ≈ 1e-3:

   h        a  Float32 relerr  Float16 relerr
0.25      0.5        9.62e-09        4.85e-05
 1.3      0.8        2.41e-08        7.85e-05
   2     0.99        4.82e-09        8.35e-05
 0.5        2        5.01e-08        1.48e-04

Script

# Accuracy check of `StatsFuns.owens_t` against Arb.
#
# Standalone script: not part of runtests.jl, and Arblib is not a test dependency.
# Run it in a temporary environment that has StatsFuns and Arblib:
#
#     julia> import Pkg; Pkg.activate(; temp = true)
#     julia> Pkg.develop(path = dirname(@__DIR__)); Pkg.add("Arblib")
#     julia> include("test/owens_t_error_analysis.jl")
#
# The reference is Arb's verified integration of the integral representation
#     T(h, a) = 1/(2π) ∫₀ᵃ exp(-h²(1+x²)/2)/(1+x²) dx
# The (h, a) grid covers every dispatch path (T1–T6 and the a > 1 remapping).
# Float32/Float16 just call the Float64 path and round, so they are only spot-checked.

using StatsFuns, Arblib, Printf, Statistics, Test

# 512-bit precision and a large evaluation budget keep the enclosure tight even for
# large h, where the integrand is a sharp spike near x = 0 and T(h, a) is as small
# as ~1e-89. Returns the value and the enclosure radius relative to it.
const INT_OPTS = Arblib.calc_integrate_opt_struct(0, 10_000_000, 10_000, 0, 0)

function owens_t_arb(h, a; prec = 512)
    H = Arb(h; prec)
    f = x -> exp(-H^2 * (1 + x^2) / 2) / (1 + x^2)
    I = Arblib.integrate(f, 0, a; prec, rtol = 1e-30, atol = Arblib.Mag(), opts = INT_OPTS)
    T = real(I) / (2 * Arb(π; prec))
    value = Float64(T)
    return value, Float64(Arblib.radius(Arf, T)) / abs(value)
end

# `meth` is local to `owens_t_dispatch` in src/owens_t.jl; mirror it to name which of
# T1–T6 handles a given (h, a) with 0 ≤ a ≤ 1.
const METH = (1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 3, 4, 4, 4, 4, 5, 6)

function regime_label(h, a)
    h == 0 && return "h = 0"
    abs(a) == 1 && return "|a| = 1"
    abs(a) > 1 && return abs(h) <= 0.67 ? "a > 1, h <= 0.67" : "a > 1, h > 0.67"
    return "T$(METH[StatsFuns.owens_t_compute_code(abs(h), abs(a))])"
end

# h and a values straddle the internal breakpoints (owens_t_hrange / owens_t_arange)
# and push past a = 1 into both remapping branches.
const HS = (0.01, 0.02, 0.05, 0.08, 0.1, 0.125, 0.2, 0.26, 0.4, 0.5, 0.6, 0.67,
            0.7, 1.0, 1.6, 1.7, 2.0, 2.33, 2.4, 3.0, 3.36, 3.4, 4.0, 4.8, 6.0, 10.0, 20.0)
const AS = (0.01, 0.025, 0.05, 0.09, 0.12, 0.15, 0.25, 0.36, 0.45, 0.5, 0.7, 0.9,
            0.99, 0.99999, 1.0, 1.2, 2.0, 5.0, 100.0)

grid = [(h, a) for h in HS for a in AS]
append!(grid, [(0.0, 0.5), (0.0, 2.0), (1.3, -0.4), (-1.3, 0.4), (-2.5, -0.8)])  # h = 0, sign symmetry

results = map(grid) do (h, a)
    computed = owens_t(h, a)
    ref, relrad = owens_t_arb(h, a)
    (; h, a, label = regime_label(h, a),
       relerr = abs(computed - ref) / abs(ref), ulps = abs(computed - ref) / eps(ref), relrad)
end

relerrs = [r.relerr for r in results]
ulpss = [r.ulps for r in results]

println("Owen's T Float64 vs Arb — $(length(results)) points")
@printf("relative error: max %.2e  mean %.2e  median %.2e\n", maximum(relerrs), mean(relerrs), median(relerrs))
@printf("ULP error:      max %.0f  mean %.1f  median %.0f\n", maximum(ulpss), mean(ulpss), median(ulpss))
@printf("max Arb enclosure relative radius: %.1e  (eps = %.1e)\n", maximum(r.relrad for r in results), eps(Float64))

println("\nper dispatch path:")
@printf("  %-16s %6s %12s %9s\n", "path", "count", "max relerr", "max ULPs")
for lab in sort(unique(r.label for r in results))
    sub = filter(r -> r.label == lab, results)
    @printf("  %-16s %6d %12.2e %9.0f\n", lab, length(sub), maximum(r.relerr for r in sub), maximum(r.ulps for r in sub))
end

println("\nworst 5 by relative error:")
@printf("  %8s %10s  %-16s %12s %9s\n", "h", "a", "path", "relerr", "ULPs")
for r in sort(results; by = r -> r.relerr, rev = true)[1:5]
    @printf("  %8.4g %10.5g  %-16s %12.2e %9.0f\n", r.h, r.a, r.label, r.relerr, r.ulps)
end

# Float32/Float16 round the inputs first, so take the reference at the rounded inputs
# to isolate the output error of the fallback from the input rounding.
println("\nFloat32 / Float16 spot check:")
@printf("  %8s %8s %14s %14s\n", "h", "a", "Float32 relerr", "Float16 relerr")
for (h, a) in ((0.25, 0.5), (1.3, 0.8), (2.0, 0.99), (0.5, 2.0))
    h32, a32 = Float32(h), Float32(a)
    h16, a16 = Float16(h), Float16(a)
    ref32, _ = owens_t_arb(Float64(h32), Float64(a32))
    ref16, _ = owens_t_arb(Float64(h16), Float64(a16))
    @printf("  %8.4g %8.4g %14.2e %14.2e\n", h, a,
            abs(owens_t(h32, a32) - ref32) / abs(ref32), abs(owens_t(h16, a16) - ref16) / abs(ref16))
end

@testset "owens_t accuracy vs Arb" begin
    @test maximum(relerrs) < 5e-14
    @test maximum(r.relrad for r in results) < 1e-20
end

[devmotion]

One thing I also noticed: owens_t(0, Inf) returns NaN even though Owen's T function is 1/4 for these values. I think it would be good in general to add tests for such non-finite and special inputs (involving Inf, -Inf, NaN).

[stevengj]

I fixed the Inf and NaN cases and added some tests, including for the correct sign of ±0.0.

[ararslan]

Drive-by minor review

[devmotion]

Looks good to me. If you update the version number, I could tag a new release right away when the PR is merged.

[stevengj]

Done: bumped the version to 2.2.0

[devmotion]

Thank you!