support for geometric distribution

Author: palday

src/MixedModels.jl

@@ -7,7 +7,7 @@ using BSplineKit: interpolate
 using Compat: @compat
 using DataAPI: DataAPI, levels, refpool, refarray, refvalue
 using Distributions: Distributions, Bernoulli, Binomial, Chisq, Distribution, Gamma
-using Distributions: InverseGaussian, Normal, Poisson, ccdf
+using Distributions: Geometric, InverseGaussian, Normal, Poisson, ccdf
 using GLM: GLM, GeneralizedLinearModel, IdentityLink, InverseLink, LinearModel
 using GLM: Link, LogLink, LogitLink, ProbitLink, SqrtLink
 using GLM: canonicallink, glm, linkinv, dispersion, dispersion_parameter
@@ -55,6 +55,7 @@ export @formula,
     Grouping,
     Gamma,
     GeneralizedLinearMixedModel,
+    Geometric,
     HelmertCoding,
     HypothesisCoding,
     IdentityLink,

src/generalizedlinearmixedmodel.jl

@@ -383,6 +383,20 @@ function GeneralizedLinearMixedModel(
     )
 end
 
+_distwarn(::T) where T <: Union{Bernoulli,Binomial,Poisson} = nothing
+function _distwarn(::T) where T <: Union{Gamma,InverseGaussian,Normal}
+    @warn """Results for families with a dispersion parameter are not reliable.
+             It is best to avoid trying to fit such models in MixedModels until
+             the authors gain a better understanding of those cases."""
+    return nothing
+end
+
+function _distwarn(::Geometric)
+    @warn "The results for geometric family models have not been confirmed to be reliable."
+end
+
+_distwarn(::T) where {T <: Distribution} = error("$(nameof(T)) distribution is not supported")
+
 function GeneralizedLinearMixedModel(
     f::FormulaTerm,
     tbl::Tables.ColumnTable,
@@ -400,11 +414,7 @@ function GeneralizedLinearMixedModel(
         ArgumentError("use LinearMixedModel for Normal distribution with IdentityLink")
     )
 
-    if !any(isa(d, dist) for dist in (Bernoulli, Binomial, Poisson))
-        @warn """Results for families with a dispersion parameter are not reliable.
-                 It is best to avoid trying to fit such models in MixedModels until
-                 the authors gain a better understanding of those cases."""
-    end
+    _distwarn(d)
 
     LMM = LinearMixedModel(f, tbl; contrasts, wts, amalgamate)
     y = copy(LMM.y)

src/simulate.jl

@@ -76,14 +76,25 @@ Note that `n` is the `n` parameter for the Binomial distribution,
 random draw (an integer in [0, n]) into a probability (a float in [0,1]).
 """
 function _rand(rng::AbstractRNG, ::Binomial, location, scale=missing, n=1)
-    dist = Binomial(Int(n), location) 
+    dist = Binomial(Int(n), location)
     return rand(rng, dist) / n
 end
 
 function _rand(rng::AbstractRNG, ::T, location, scale=missing, n=1) where T <: Union{Bernoulli,Poisson}
     dist = T(location)
     return rand(rng, dist)
-end    
+end
+
+function _rand(rng::AbstractRNG, ::Geometric, location, scale=missing, n=1)
+    dist = Geometric(location)
+    return rand(rng, dist)
+end
+
+function _rand(rng::AbstractRNG, ::T, location, scale, n=1) where T <: Union{Gamma,InverseGaussian}
+    @warn "Gamma and Inverse Gaussian are not supported in fitting and poorly tested in simulation" maxwarn=1
+    dist = T(location, scale)
+    return rand(rng, dist)
+end
 
 function _rand(rng::AbstractRNG, ::Normal, location, scale=missing, n=1)
     # skip constructing a distribution
@@ -280,7 +291,8 @@ function _simulate!(
     # families with a dispersion parameter
     mul!(η, fullrankx(lm), β, one(T), one(T))
 
-    μ = resp === nothing ? linkinv.(Link(m), η) : GLM.updateμ!(resp, η).mu
+    @info "" resp
+    μ = isnothing(resp) ? GLM.linkfun.(Link(m), η) : GLM.updateμ!(resp, η).mu
 
     # convert to the distribution / add in noise
     @inbounds for (idx, val) in enumerate(μ)

test/bootstrap.jl

@@ -64,6 +64,23 @@ end
         gm2sim = refit!(simulate!(StableRNG(42), deepcopy(gm2)); fast=true, progress=false)
         @test isapprox(gm2.β, gm2sim.β; atol=norm(stderror(gm2)))
     end
+
+    @testset "Geometric" begin
+        n = 1000
+        data = (; y=fill(1, n), g=rand('A':'G', n))
+        data = (; y=rand(StableRNG(42), Geometric(0.5), n),
+                  g=rand('A':'G', n))
+        @test_logs (:warn, r"geometric") MixedModel(@formula(y ~ 1 + (1|g)), data, Geometric())
+        m = @suppress MixedModel(@formula(y ~ 1 + (1|g)), data, Geometric())
+        v = simulate(StableRNG(42), m; β=[exp(0.5)], θ=[0])
+        gmean = mean(rand(StableRNG(42), Geometric(), n))
+        @test mean(v) ≈ gmean atol=0.005
+
+        simulate!(StableRNG(42), m; β=[exp(0.5)], θ=[0])
+        v = response(m)
+        @test mean(v) ≈ gmean atol=0.005
+    end
+
     @testset "_rand with dispersion" begin
         @test_throws ArgumentError MixedModels._rand(StableRNG(42), NegativeBinomial(), 1, 1)