Skip to content

Implement binomial, Poisson, and negative binomial inverse functions in pure Julia #205

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
andreasnoack wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

Implement binomial, Poisson, and negative binomial inverse functions in pure Julia #205

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 143 additions & 5 deletions src/distrs/binom.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,20 +30,158 @@ for l in ("", "log"), compl in (false, true)
end
end

# Rmath implementations
# Inverse CDF: find smallest k such that binomcdf(n, p, k) >= cprob.
# Based on VBA critbinomial by Ian Smith.
# PMF ratios: PMF(k+1)/PMF(k) = (n-k)*p/((k+1)*(1-p))
# PMF(k-1)/PMF(k) = k*(1-p)/((n-k+1)*p)
function _critbinomial(n::Float64, p::Float64, cprob::Float64)
q = 1.0 - p

# Normal approximation for initial guess
σ = sqrt(n * p * q)
i = clamp(floor(n * p + norminvcdf(min(cprob, 1.0 - 1.0e-15)) * σ + 0.5), 0.0, n)

# Compute CDF at the guess
pr = Float64(binomcdf(n, p, i))

if pr >= cprob
# Search left: find smallest k with CDF(k) >= cprob
while i > 0
tpr = Float64(binompdf(n, p, i))
if pr - tpr < cprob
return i # CDF(i-1) < cprob <= CDF(i)
end
pr -= tpr
i -= 1.0
end
return 0.0
else
# Search right
while i < n
i += 1.0
tpr = Float64(binompdf(n, p, i))
pr += tpr
if pr >= cprob
return i
end
end
return n
end
end

# Inverse CCDF: find smallest k such that binomccdf(n, p, k) <= cprob.
# Based on VBA critcompbinomial by Ian Smith.
function _critcompbinomial(n::Float64, p::Float64, cprob::Float64)
q = 1.0 - p

# Normal approximation
σ = sqrt(n * p * q)
i = clamp(floor(n * p - norminvcdf(min(cprob, 1.0 - 1.0e-15)) * σ + 0.5), 0.0, n)

# Compute CCDF at the guess
pr = Float64(binomccdf(n, p, i))

if pr > cprob
# Search right
while i < n
i += 1.0
tpr = Float64(binompdf(n, p, i))
pr -= tpr
if pr <= cprob
return i
end
end
return n
else
# Search left: find smallest k with CCDF(k) <= cprob
while i > 0
tpr = Float64(binompdf(n, p, i))
if pr + tpr > cprob
return i # CCDF(i) <= cprob < CCDF(i-1)
end
pr += tpr
i -= 1.0
end
return 0.0
end
end

# Wrapper with edge cases and post-correction (VBA crit_binomial)
function _binom_invcdf(n::Float64, p::Float64, q::Float64)
if q < 0 || q > 1 || n < 0 || p < 0 || p > 1 || isnan(q) || isnan(n) || isnan(p)
return NaN
elseif q == 0 || p == 0
return 0.0
elseif p == 1
return n
end

i = _critbinomial(n, p, q)

# Post-correction
pr = Float64(binomcdf(n, p, i))
if pr >= q
while i > 0
pr2 = Float64(binomcdf(n, p, i - 1.0))
if pr2 < q
return i
end
i -= 1.0
pr = pr2
end
return 0.0
else
return i + 1.0
end
end

# Wrapper with edge cases and post-correction (VBA comp_crit_binomial)
function _binom_invccdf(n::Float64, p::Float64, q::Float64)
if q < 0 || q > 1 || n < 0 || p < 0 || p > 1 || isnan(q) || isnan(n) || isnan(p)
return NaN
elseif q == 1 || p == 0
return 0.0
elseif q == 0 || p == 1
return n
end

i = _critcompbinomial(n, p, q)

# Post-correction
pr = Float64(binomccdf(n, p, i))
if pr <= q
while i > 0
pr2 = Float64(binomccdf(n, p, i - 1.0))
if pr2 > q
return i
end
i -= 1.0
pr = pr2
end
return 0.0
else
return i + 1.0
end
end

# Public API

function binominvcdf(n::Real, p::Real, q::Real)
T = float(Base.promote_typeof(n, p, q))
return convert(T, Rmath.qbinom(q, n, p, true, false))
return convert(T, _binom_invcdf(Float64(n), Float64(p), Float64(q)))
end

function binominvccdf(n::Real, p::Real, q::Real)
T = float(Base.promote_typeof(n, p, q))
return convert(T, Rmath.qbinom(q, n, p, false, false))
return convert(T, _binom_invccdf(Float64(n), Float64(p), Float64(q)))
end

function binominvlogcdf(n::Real, p::Real, lq::Real)
T = float(Base.promote_typeof(n, p, lq))
return convert(T, Rmath.qbinom(lq, n, p, true, true))
return convert(T, _binom_invcdf(Float64(n), Float64(p), exp(Float64(lq))))
end

function binominvlogccdf(n::Real, p::Real, lq::Real)
T = float(Base.promote_typeof(n, p, lq))
return convert(T, Rmath.qbinom(lq, n, p, false, true))
return convert(T, _binom_invccdf(Float64(n), Float64(p), exp(Float64(lq))))
end
164 changes: 157 additions & 7 deletions src/distrs/nbinom.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,22 +66,172 @@ function nbinomlogccdf(r::Real, p::Real, k::Real)
end
end

# Rmath implementations
# TODO: implement https://arxiv.org/abs/2001.03953
# for inverting the incomplete beta function wrt the 2nd argument.
# Inverse CDF: find smallest k such that nbinomcdf(r, p, k) >= cprob.
# Based on VBA critnegbinom by Ian Smith.
# PMF ratio: PMF(k+1)/PMF(k) = (k+r)/(k+1) * (1-p)
function _critnbinom(r::Float64, p::Float64, cprob::Float64)
q = 1.0 - p

# Normal approximation: mean = r*q/p, var = r*q/p^2
μ = r * q / p
σ = sqrt(μ / p)
i = max(0.0, floor(μ + norminvcdf(min(cprob, 1.0 - 1.0e-15)) * σ + 0.5))

# Compute CDF at the guess
pr = Float64(nbinomcdf(r, p, i))

if pr >= cprob
# Search left
while i > 0
tpr = Float64(nbinompdf(r, p, i))
if pr - tpr < cprob
return i
end
pr -= tpr
i -= 1.0
end
return 0.0
else
# Search right
for _ in 1:10_000
i += 1.0
tpr = Float64(nbinompdf(r, p, i))
pr += tpr
if pr >= cprob
return i
end
end
return i
end
end

# Inverse CCDF: find smallest k such that nbinomccdf(r, p, k) <= cprob.
# Based on VBA critcompnegbinom by Ian Smith.
function _critcompnbinom(r::Float64, p::Float64, cprob::Float64)
q = 1.0 - p

# Normal approximation
μ = r * q / p
σ = sqrt(μ / p)
i = max(0.0, floor(μ - norminvcdf(min(cprob, 1.0 - 1.0e-15)) * σ + 0.5))

# Compute CCDF at the guess
pr = Float64(nbinomccdf(r, p, i))

if pr > cprob
# Search right
for _ in 1:10_000
i += 1.0
tpr = Float64(nbinompdf(r, p, i))
pr -= tpr
if pr <= cprob
return i
end
end
return i
else
# Search left
while i > 0
tpr = Float64(nbinompdf(r, p, i))
if pr + tpr > cprob
return i
end
pr += tpr
i -= 1.0
end
return 0.0
end
end

# Wrapper with edge cases
function _nbinom_invcdf(r::Float64, p::Float64, q::Float64)
if q < 0 || q > 1 || r <= 0 || p < 0 || p > 1 || isnan(q) || isnan(r) || isnan(p)
return NaN
elseif q == 0 || p == 1
return 0.0
elseif q == 1
return Inf
elseif p == 0
return Inf
end

i = _critnbinom(r, p, q)

# Post-correction
pr = Float64(nbinomcdf(r, p, i))
if pr >= q
while i > 0
pr2 = Float64(nbinomcdf(r, p, i - 1.0))
if pr2 < q
return i
end
i -= 1.0
end
return 0.0
else
return i + 1.0
end
end

function _nbinom_invccdf(r::Float64, p::Float64, q::Float64)
if q < 0 || q > 1 || r <= 0 || p < 0 || p > 1 || isnan(q) || isnan(r) || isnan(p)
return NaN
elseif q == 0 || p == 0
return Inf
elseif q == 1
return 0.0
elseif p == 1
return 0.0
end

i = _critcompnbinom(r, p, q)

# Post-correction
pr = Float64(nbinomccdf(r, p, i))
if pr <= q
while i > 0
pr2 = Float64(nbinomccdf(r, p, i - 1.0))
if pr2 > q
return i
end
i -= 1.0
end
return 0.0
else
return i + 1.0
end
end

# Public API

function nbinominvcdf(r::Real, p::Real, q::Real)
T = float(Base.promote_typeof(r, p, q))
return convert(T, Rmath.qnbinom(q, r, p, true, false))
return convert(T, _nbinom_invcdf(Float64(r), Float64(p), Float64(q)))
end

function nbinominvccdf(r::Real, p::Real, q::Real)
T = float(Base.promote_typeof(r, p, q))
return convert(T, Rmath.qnbinom(q, r, p, false, false))
return convert(T, _nbinom_invccdf(Float64(r), Float64(p), Float64(q)))
end

function nbinominvlogcdf(r::Real, p::Real, lq::Real)
T = float(Base.promote_typeof(r, p, lq))
return convert(T, Rmath.qnbinom(lq, r, p, true, true))
_lq = Float64(lq)
result = if _lq > -1
_nbinom_invccdf(Float64(r), Float64(p), -expm1(_lq))
else
_nbinom_invcdf(Float64(r), Float64(p), exp(_lq))
end
return convert(T, result)
end

function nbinominvlogccdf(r::Real, p::Real, lq::Real)
T = float(Base.promote_typeof(r, p, lq))
return convert(T, Rmath.qnbinom(lq, r, p, false, true))
_lq = Float64(lq)
result = if _lq > -1
_nbinom_invcdf(Float64(r), Float64(p), -expm1(_lq))
else
_nbinom_invccdf(Float64(r), Float64(p), exp(_lq))
end
return convert(T, result)
end
Loading
Loading