truncated_normal_rejection_sampler.jl

# Rejection sampler based on algorithm from Robert (1995)
#
#  - Available at http://arxiv.org/abs/0907.4010
#
# Copied this implementation from Distributions.jl, with few modifications
# to make it generic.

using Random, Statistics

Δ2(x, y) = (x - y) * (x + y)
randnt(lb::Real, ub::Real, tp::Real = NaN) = randnt(Random.GLOBAL_RNG, lb, ub, tp)

function randnt(rng::AbstractRNG, lb::Real, ub::Real, tp::Real = 0)
    lb ≤ ub || throw(ArgumentError("invalid bounds lb=$lb, ub=$ub"))
    T = typeof(middle(lb, ub))
    if 3tp > 1   # has considerable chance of falling in [lb, ub]
        while true
            r = randn(rng, T)
            lb ≤ r ≤ ub && return r
        end
    else
        span = ub - lb
        if lb > 0 && span > 2 / (lb + sqrt(lb^2 + 4)) * exp((lb - sqrt(lb^2 + 4)) * lb / 4)
            a = (sqrt(lb^2 + 4) + lb)/2
            while true
                r = randexp(rng, T) / a + lb
                u = rand(rng, T)
                if u < exp(-(r - a)^2 / 2) && r < ub
                    return r
                end
            end
        elseif ub < 0 && span > 2 / (-ub + sqrt(ub^2 + 4)) * exp((ub^2 + ub * sqrt(ub^2 + 4)) / 4)
            a = (sqrt(ub^2 + 4) - ub) / 2
            while true
                r = randexp(rng, T) / a - ub
                u = rand(rng, T)
                if u < exp(-(r - a)^2 / 2) && r < -lb
                    return -r
                end
            end
        else
            while true
                r = lb + rand(rng, T) * (ub - lb)
                u = rand(rng, T)
                if lb > 0
                    rho = exp(Δ2(lb, r) / 2)
                elseif ub < 0
                    rho = exp(Δ2(ub, r) / 2)
                else
                    rho = exp(-r^2 / 2)
                end
                if u < rho
                    return r
                end
            end
        end
    end
end