Improve GammaIPSampler and GammaMTSampler

src/multivariate/mvtdist.jl

@@ -156,7 +156,7 @@ end
 
 # Sampling (for GenericMvTDist)
 function _rand!(rng::AbstractRNG, d::GenericMvTDist, x::AbstractVector{<:Real})
-    chisqd = Chisq(d.df)
+    chisqd = Chisq{partype(d)}(d.df)
     y = sqrt(rand(rng, chisqd) / d.df)
     unwhiten!(d.Σ, randn!(rng, x))
     x .= x ./ y .+ d.μ
@@ -165,7 +165,7 @@ end
 
 function _rand!(rng::AbstractRNG, d::GenericMvTDist, x::AbstractMatrix{T}) where T<:Real
     cols = size(x,2)
-    chisqd = Chisq(d.df)
+    chisqd = Chisq{partype(d)}(d.df)
     y = Matrix{T}(undef, 1, cols)
     unwhiten!(d.Σ, randn!(rng, x))
     rand!(rng, chisqd, y)

src/samplers/gamma.jl

@@ -162,32 +162,49 @@ end
 # doi:10.1145/358407.358414
 # http://www.cparity.com/projects/AcmClassification/samples/358414.pdf
 
-struct GammaMTSampler <: Sampleable{Univariate,Continuous}
-    d::Float64
-    c::Float64
-    κ::Float64
+# valid for shape >= 1
+struct GammaMTSampler{T<:Real} <: Sampleable{Univariate,Continuous}
+    d::T
+    c::T
+    κ::T
+    r::T
 end
 
 function GammaMTSampler(g::Gamma)
-    d = shape(g) - 1/3
-    c = 1.0 / sqrt(9.0 * d)
+    # Setup (Step 1)
+    d = shape(g) - 1//3
+    c = inv(3 * sqrt(d))
+
+    # Pre-compute scaling factor
     κ = d * scale(g)
-    GammaMTSampler(d, c, κ)
+
+    # We also pre-compute the factor in the squeeze function
+    return GammaMTSampler(promote(d, c, κ, 331//10_000)...)
 end
 
-function rand(rng::AbstractRNG, s::GammaMTSampler)
+function rand(rng::AbstractRNG, s::GammaMTSampler{T}) where {T<:Real}
+    d = s.d
+    c = s.c
+    κ = s.κ
+    r = s.r
+    z = zero(T)
     while true
-        x = randn(rng)
-        v = 1.0 + s.c * x
-        while v <= 0.0
-            x = randn(rng)
-            v = 1.0 + s.c * x
+        # Generate v (Step 2)
+        x = randn(rng, T)
+        cbrt_v = 1 + c * x
+        while cbrt_v <= z # requires x <= -sqrt(9 * shape - 3)
+            x = randn(rng, T)
+            cbrt_v = 1 + c * x
         end
-        v *= (v * v)
-        u = rand(rng)
-        x2 = x * x
-        if u < 1.0 - 0.331 * abs2(x2) || log(u) < 0.5 * x2 + s.d * (1.0 - v + log(v))
-            return v*s.κ
+        v = cbrt_v^3
+
+        # Generate uniform u (Step 3)
+        u = rand(rng, T)
+
+        # Check acceptance (Step 4 and 5)
+        xsq = x^2
+        if u < 1 - r * xsq^2 || log(u) < xsq / 2 + d * logmxp1(v)
+            return v * κ
         end
     end
 end
@@ -199,24 +216,15 @@ struct GammaIPSampler{S<:Sampleable{Univariate,Continuous},T<:Real} <: Sampleabl
     nia::T #-1/scale
 end
 
-function GammaIPSampler(d::Gamma,::Type{S}) where S<:Sampleable
-    GammaIPSampler(Gamma(1.0 + shape(d), scale(d)), -1.0 / shape(d))
+GammaIPSampler(d::Gamma) = GammaIPSampler(d, GammaMTSampler)
+function GammaIPSampler(d::Gamma, ::Type{S}) where {S<:Sampleable}
+    shape_d = shape(d)
+    sampler = S(Gamma{partype(d)}(1 + shape_d, scale(d)))
+    return GammaIPSampler(sampler, -inv(shape_d))
 end
-GammaIPSampler(d::Gamma) = GammaIPSampler(d,GammaMTSampler)
 
 function rand(rng::AbstractRNG, s::GammaIPSampler)
     x = rand(rng, s.s)
     e = randexp(rng)
     x*exp(s.nia*e)
 end
-
-# function sampler(d::Gamma)
-#     if d.shape < 1.0
-#         # TODO: d.shape = 0.5 : use scaled chisq
-#         GammaIPSampler(d)
-#     elseif d.shape == 1.0
-#         Exponential(d.scale)
-#     else
-#         GammaGDSampler(d)
-#     end
-# end

src/univariate/continuous/chisq.jl

@@ -96,7 +96,14 @@ gradlogpdf(d::Chisq{T}, x::Real) where {T<:Real} =  x > 0 ? (d.ν/2 - 1) / x - 1
 
 #### Sampling
 
-rand(rng::AbstractRNG, d::Chisq) =
-    (ν = d.ν; rand(rng, Gamma(ν / 2.0, 2.0one(ν))))
+function rand(rng::AbstractRNG, d::Chisq)
+    α = dof(d) / 2
+    θ = oftype(α, 2)
+    return rand(rng, Gamma{typeof(α)}(α, θ))
+end
 
-sampler(d::Chisq) = (ν = d.ν; sampler(Gamma(ν / 2.0, 2.0one(ν))))
+function sampler(d::Chisq)
+    α = dof(d) / 2
+    θ = oftype(α, 2)
+    return sampler(Gamma{typeof(α)}(α, θ))
+end

src/univariate/continuous/gamma.jl

@@ -105,9 +105,10 @@ function rand(rng::AbstractRNG, d::Gamma)
         # TODO: shape(d) = 0.5 : use scaled chisq
         return rand(rng, GammaIPSampler(d))
     elseif shape(d) == 1.0
-        return rand(rng, Exponential(d.θ))
+        θ = 
+        return rand(rng, Exponential{partype(d)}(scale(d)))
     else
-        return rand(rng, GammaGDSampler(d))
+        return rand(rng, GammaMTSampler(d))
     end
 end
 
@@ -116,9 +117,9 @@ function sampler(d::Gamma)
         # TODO: shape(d) = 0.5 : use scaled chisq
         return GammaIPSampler(d)
     elseif shape(d) == 1.0
-        return sampler(Exponential(d.θ))
+        return sampler(Exponential{partype(d)}(scale(d)))
     else
-        return GammaGDSampler(d)
+        return GammaMTSampler(d)
     end
 end
 

test/fit.jl

@@ -370,7 +370,7 @@ end
         d = fit(dist, func[2](dist(5.0, 3.0), N + 1))
         @test isa(d, dist)
         @test isapprox(location(d), 5.0, atol=0.02)
-        @test isapprox(scale(d)   , 3.0, atol=0.02)
+        @test isapprox(scale(d)   , 3.0, atol=0.03)
     end
 end
 

test/samplers.jl

@@ -101,6 +101,28 @@ import Distributions:
         end
     end
 
+    @testset "GammaIPSampler" begin
+        @testset "d=$d" for d in [Gamma(0.1, 1.0), Gamma(0.9, 1.0)]
+            s = sampler(d)
+            @test s isa GammaIPSampler{<:GammaMTSampler}
+            @test s.s isa GammaMTSampler
+            test_samples(s, d, n_tsamples)
+            test_samples(s, d, n_tsamples, rng=rng)
+
+            s = @inferred(GammaIPSampler(d, GammaMTSampler))
+            @test s isa GammaIPSampler{<:GammaMTSampler}
+            @test s.s isa GammaMTSampler
+            test_samples(s, d, n_tsamples)
+            test_samples(s, d, n_tsamples, rng=rng)
+
+            s = @inferred(GammaIPSampler(d, GammaGDSampler))
+            @test s isa GammaIPSampler{<:GammaGDSampler}
+            @test s.s isa GammaGDSampler
+            test_samples(s, d, n_tsamples)
+            test_samples(s, d, n_tsamples, rng=rng)
+        end
+    end
+
     @testset "Random.Sampler" begin
         for dist in (
             Binomial(5, 0.3),

test/univariate/continuous/tdist.jl

@@ -4,6 +4,9 @@ using ForwardDiff
 using Test
 
 @testset "Type stability of `rand` (#1614)" begin
-    @inferred(rand(TDist(big"1.0")))
+    if VERSION >= v"1.9.0-DEV.348"
+        # randn(::BigFloat) was only added in https://github.com/JuliaLang/julia/pull/44714
+        @inferred(rand(TDist(big"1.0")))
+    end
     @inferred(rand(TDist(ForwardDiff.Dual(1.0))))
 end