add owens_t function

LICENSE.md

@@ -20,3 +20,29 @@ The StatsFuns.jl package is licensed under the MIT "Expat" License:
 > CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 > TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 > SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+Portions are also under the (essentially equivalent) Boost Software License:
+
+> Boost Software License - Version 1.0 - August 17th, 2003
+>
+> Permission is hereby granted, free of charge, to any person or organization
+> obtaining a copy of the software and accompanying documentation covered by
+> this license (the "Software") to use, reproduce, display, distribute,
+> execute, and transmit the Software, and to prepare derivative works of the
+> Software, and to permit third-parties to whom the Software is furnished to
+> do so, all subject to the following:
+>
+> The copyright notices in the Software and this entire statement, including
+> the above license grant, this restriction and the following disclaimer,
+> must be included in all copies of the Software, in whole or in part, and
+> all derivative works of the Software, unless such copies or derivative
+> works are solely in the form of machine-executable object code generated by
+> a source language processor.
+>
+> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+> FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
+> SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
+> FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
+> ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+> DEALINGS IN THE SOFTWARE.

Project.toml

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

README.md

@@ -61,7 +61,9 @@ softmax!,       # inplace softmax
 
 # misc
 logmvgamma,     # logarithm of multivariate gamma function
-lstirling_asym
+lstirling_asym,
+
+owens_t         # Owen's T function
 ```
 
 ## Distribution-related functions

ext/StatsFunsChainRulesCoreExt.jl

@@ -1,7 +1,8 @@
 module StatsFunsChainRulesCoreExt
 
 using StatsFuns
-using StatsFuns: digamma
+using SpecialFunctions: digamma, erf
+
 import ChainRulesCore
 
 ChainRulesCore.@scalar_rule(
@@ -82,4 +83,12 @@ ChainRulesCore.@scalar_rule(
     ),
 )
 
+ChainRulesCore.@scalar_rule(
+    owens_t(h::Real, a::Real),
+    (
+        normpdf(h) * erf((h * a) * invsqrt2) / -2,
+        inv2π * exp(-h^2 * (1 + a^2) / 2) / (1 + a^2),
+    ),
+)
+
 end # module

src/StatsFuns.jl

@@ -4,7 +4,7 @@ using Base: Math.@horner
 using Reexport: @reexport
 using Rmath: Rmath
 using SpecialFunctions: beta_inc, beta_inc_inv, digamma,
-    erfc, erfcinv, erfcx, gamma_inc, gamma_inc_inv, logbeta, loggamma
+    erf, erfc, erfcinv, erfcx, gamma_inc, gamma_inc_inv, logbeta, loggamma
 
 # reexports
 @reexport using IrrationalConstants:
@@ -260,11 +260,15 @@ export
     logmvbeta, # logarithm of multivariate beta function
     lstirling_asym,
 
+    # owens_t
+    owens_t,
+
     # tvpack
     bvncdf    # bivariate standard normal cdf
 
 ## source files
 include("misc.jl")
+include("owens_t.jl")
 include("tvpack.jl")
 
 include("distrs/beta.jl")

src/owens_t.jl

@@ -0,0 +1,235 @@
+# Owen's T function
+#
+# Ported from Boost owen_t.hpp, which is Copyright Benjamin Sobotta 2012
+#      Use, modification and distribution are subject to the
+#      Boost Software License, Version 1.0. (http://www.boost.org/LICENSE_1_0.txt)
+# Port by Steven G. Johnson, also under Boost Software License, Version 1.0.
+#
+# Algorithm based on Mike Patefield and David Tandy, "Fast and accurate calculation
+# of Owen's T-function," Journal of Statistical Software, 5, pp. 1-25 (2000).
+
+# P(-∞<Z<=x)-0.5 with Z being normally distributed.
+owens_t_znorm1(x::Real) = erf(x * invsqrt2) / 2
+
+# Auxiliary function, it computes an array key that is used to determine
+# the specific computation method for Owen's T and the order thereof
+# used in owens_t_dispatch.  Differs from C++ in using 1-based indices
+const owens_t_hrange = (0.02f0, 0.06f0, 0.09f0, 0.125f0, 0.26f0, 0.4f0, 0.6f0, 1.6f0, 1.7f0, 2.33f0, 2.4f0, 3.36f0, 3.4f0, 4.8f0) # length 14
+const owens_t_arange = (0.025f0, 0.09f0, 0.15f0, 0.36f0, 0.5f0, 0.9f0, 0.99999f0) # length 7
+const owens_t_select = (
+    1, 1, 2, 13, 13, 13, 13, 13, 13, 13, 13, 16, 16, 16, 9,
+    1, 2, 2, 3, 3, 5, 5, 14, 14, 15, 15, 16, 16, 16, 9,
+    2, 2, 3, 3, 3, 5, 5, 15, 15, 15, 15, 16, 16, 16, 10,
+    2, 2, 3, 5, 5, 5, 5, 7, 7, 16, 16, 16, 16, 16, 10,
+    2, 3, 3, 5, 5, 6, 6, 8, 8, 17, 17, 17, 12, 12, 11,
+    2, 3, 5, 5, 5, 6, 6, 8, 8, 17, 17, 17, 12, 12, 12,
+    2, 3, 4, 4, 6, 6, 8, 8, 17, 17, 17, 17, 17, 12, 12,
+    2, 3, 4, 4, 6, 6, 18, 18, 18, 18, 17, 17, 17, 12, 12,
+) # 1-based indices (in 8x15 "matrix")
+function owens_t_compute_code(h::Float64, a::Float64)
+    ihint = something(findfirst(>=(h), owens_t_hrange), length(owens_t_hrange) + 1)
+    iaint = something(findfirst(>=(a), owens_t_arange), length(owens_t_arange) + 1)
+    return owens_t_select[(iaint - 1) * 15 + ihint]
+end
+
+const owens_t_ord = (2, 3, 4, 5, 7, 10, 12, 18, 10, 20, 30, 0, 4, 7, 8, 20, 0, 0) # 53-bit precision table
+function owens_t_get_order(icode::Int, ::Type{Float64})
+    return owens_t_ord[icode]
+end
+
+# compute the value of Owen's T function with method T1 from the reference paper
+function owens_t_T1(h::Float64, a::Float64, m::Int)
+    hs = -h * h / 2
+    as = a * a
+    aj = a * inv2π
+    dj = expm1(hs)
+    gj = hs * exp(hs)
+    val = atan(a) * inv2π
+    j, jj = one(m), oftype(h, 1)
+    while true
+        val += dj * aj / jj
+        m <= j && break
+        j += one(m)
+        jj += 2
+        aj *= as
+        dj = gj - dj
+        gj *= hs / j
+    end
+    return val
+end
+
+# compute the value of Owen's T function with method T2 from the reference paper
+function owens_t_T2(h::Float64, a::Float64, m::Int, ah::Float64)
+    maxii = m + m + one(m)
+    hs = h * h
+    as = -a * a
+    y = inv(hs)
+    ii = one(m)
+    val = zero(h)
+    vi = a * normpdf(ah)
+    z = owens_t_znorm1(ah) / h
+    while true
+        val += z
+        maxii <= ii && break
+        z = y * (vi - ii * z)
+        vi *= as
+        ii += oftype(ii, 2)
+    end
+    return val * normpdf(h)
+end
+
+# compute the value of Owen's T function with method T3 from the reference paper
+const owens_t_c2 = (
+    0.9999999999999998751, -0.99999999999988796462, 0.99999999998290743652,
+    -0.99999999896282500134, 0.99999996660459362918, -0.9999993398627247676,
+    0.99999125611136965852, -0.99991777624463387686, 0.99942835555870132569,
+    -0.99697311720723000295, 0.98751448037275303682, -0.95915857980572882813,
+    0.89246305511006708555, -0.76893425990463999675, 0.5889352846848469325,
+    -0.38380345160440256652, 0.20317601701045299653, -0.82813631607004984866e-1,
+    0.24167984735759576523e-1, -0.44676566663971825242e-2, 0.39141169402373836468e-3,
+)
+function owens_t_T3(h::Float64, a::Float64, ah::Float64)
+    as = a * a
+    hs = h * h
+    y = inv(hs)
+    ii = one(h)
+    vi = a * normpdf(ah)
+    zi = owens_t_znorm1(ah) / h
+    val = zero(h)
+    i = 1
+    while true
+        val += zi * owens_t_c2[i]
+        i == lastindex(owens_t_c2) && break
+        i += 1
+        zi = y * (ii * zi - vi)
+        vi *= as
+        ii += 2
+    end
+    return val * normpdf(h)
+end
+
+# compute the value of Owen's T function with method T4 from the reference paper
+function owens_t_T4(h::Float64, a::Float64, m::Int)
+    maxii = m + m + one(m)
+    hs = h * h
+    as = -a * a
+    ii = 1
+    ai = a * exp(-hs * (1 - as) / 2) * inv2π
+    yi = one(h)
+    val = zero(h)
+    while true
+        val += ai * yi
+        maxii <= ii && break
+        ii += 2
+        yi = (1 - hs * yi) / ii
+        ai *= as
+    end
+    return val
+end
+
+# compute the value of Owen's T function with method T5 from the reference paper
+const owens_t_pts = (
+    0.35082039676451715489e-2, 0.3127904233803075374e-1, 0.8526682628321945109e-1,
+    0.16245071730812277011, 0.25851196049125434828, 0.36807553840697533536,
+    0.48501092905604697475, 0.60277514152618576821, 0.71477884217753226516,
+    0.81475510988760098605, 0.89711029755948965867, 0.95723808085944261843, 0.99178832974629703586,
+)
+const owens_t_wts = (
+    0.18831438115323502887e-1, 0.18567086243977649478e-1, 0.18042093461223385584e-1,
+    0.17263829606398753364e-1, 0.1624321997598985673e-1, 0.14994592034116704829e-1,
+    0.13535474469662088392e-1, 0.11886351605820165233e-1, 0.10070377242777431897e-1,
+    0.81130545742299586629e-2, 0.60419009528470238773e-2, 0.38862217010742057883e-2, 0.16793031084546090448e-2,
+)
+#=
+    NOTICE:
+    - The pts[] array contains the squares (!) of the abscissas, i.e. the roots of the Legendre
+        polynomial P_n(x), instead of the plain roots as required in Gauss-Legendre
+        quadrature, because T5(h,a,m) contains only x^2 terms.
+    - The wts[] array contains the weights for Gauss-Legendre quadrature scaled with a factor
+        of 1/(2*pi) according to T5(h,a,m).
+=#
+function owens_t_T5(h::Float64, a::Float64)
+    as = a * a
+    hs = -h * h / 2
+    r = 1 .+ as .* owens_t_pts
+    return sum(owens_t_wts .* exp.(hs .* r) ./ r) * a
+end
+
+# compute the value of Owen's T function with method T6 from the reference paper
+function owens_t_T6(h::Float64, a::Float64)
+    normh = normccdf(h)
+    y = 1 - a
+    r = atan(y, 1 + a)
+    val = normh * (1 - normh) / 2
+    !iszero(r) && (val -= r * exp(-y * h * h / 2r) * inv2π)
+    return val
+end
+
+# This routine dispatches the call to one of six subroutines, depending on the values of h and a.
+# preconditions: h >= 0, 0<=a<=1, ah=a*h.
+# Simple main case for 64-bit precision or less, this is as per the Patefield-Tandy paper:
+function owens_t_dispatch(h::Float64, a::Float64, ah::Float64)
+    # Handle some special cases first, these are from
+    # page 1077 of Owen's original paper:
+    iszero(h) && return atan(a) * inv2π
+    iszero(a) && return zero(h)
+    a == 1 && return normccdf(-h) * normccdf(h) / 2
+    @assert a <= 1 # when a>1 we call this routine with 1/a:
+
+    icode = owens_t_compute_code(h, a)
+    m = owens_t_get_order(icode, typeof(h))
+    meth = (1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 3, 4, 4, 4, 4, 5, 6) # 18 entries
+    method = meth[icode]
+
+    # determine the appropriate method, T1 ... T6
+    method == 1 && return owens_t_T1(h, a, m)
+    method == 2 && return owens_t_T2(h, a, m, ah)
+    method == 3 && return owens_t_T3(h, a, ah)
+    method == 4 && return owens_t_T4(h, a, m)
+    method == 5 && return owens_t_T5(h, a)
+    return owens_t_T6(h, a)
+end
+
+# compute Owen's T function, T(h,a), for arbitrary values of h and a
+function _owens_t(h::Float64, a::Float64)
+    # exploit that T(-h,a) == T(h,a)
+    h = abs(h)
+
+    # Use equation (2) in the paper to remap the arguments
+    # such that h>=0 and 0<=a<=1 for the call of the actual
+    # computation routine.
+    abs_a = abs(a)
+    abs_ah = abs_a * h
+
+    val = if abs_a <= 1
+        owens_t_dispatch(h, abs_a, abs_ah)
+    elseif !isfinite(abs_a)
+        isnan(abs_a) ? oftype(h, NaN) : normcdf(-h) / 2
+    elseif h <= oftype(h, 0.67)
+        normh = owens_t_znorm1(h)
+        normah = owens_t_znorm1(abs_ah)
+        1 // 4 - normh * normah - owens_t_dispatch(abs_ah, inv(abs_a), h)
+    else
+        normh = normccdf(h)
+        normah = normccdf(abs_ah)
+        (normh + normah) / 2 - normh * normah - owens_t_dispatch(abs_ah, inv(abs_a), h)
+    end
+    return copysign(val, a) # exploit that T(h,-a) == -T(h,a)
+end
+
+_owens_t(h::Float32, a::Float32) = Float32(owens_t(Float64(h), Float64(a)))
+_owens_t(h::Float16, a::Float16) = Float16(owens_t(Float64(h), Float64(a)))
+
+"""
+    owens_t(h::Real, a::Real)
+
+Return [Owen's T function](https://en.wikipedia.org/wiki/Owen%27s_T_function)
+
+```math
+T(h, a) = \\frac{1}{2\\pi} \\int_0^a \\frac{e^{-h^2(1+x^2)/2}}{1 + x^2} \\, dx
+```
+
+This is the probability of ``X > h`` and ``0 < Y < aX``, where ``X`` and ``Y`` are
+i.i.d. standard normal random variables.
+"""
+owens_t(h::Real, a::Real) = _owens_t(map(float, promote(h, a))...)

test/chainrules.jl

@@ -50,4 +50,9 @@ using Random
     _, pb = rrule(poislogpdf, 0.0, 1)
     _, x̄1, _ = pb(1)
     @test x̄1 == Inf
+
+    h = randn()
+    a = randn()
+    test_frule(owens_t, h, a)
+    test_rrule(owens_t, h, a)
 end