Add multivariate hypergeometric distribution

docs/src/multivariate.md

@@ -59,6 +59,7 @@ MvNormalCanon
 MvLogitNormal
 MvLogNormal
 Dirichlet
+MvHypergeometric
 ```
 
 ## Addition Methods

src/Distributions.jl

@@ -133,6 +133,7 @@ export
     MatrixTDist,
     MixtureModel,
     Multinomial,
+    MvHypergeometric,
     MultivariateNormal,
     MvLogNormal,
     MvNormal,

src/multivariate/mvhypergeometric.jl

@@ -0,0 +1,131 @@
+"""
+    MvHypergeometric(m, n)
+
+The [Multivariate hypergeometric distribution](https://en.wikipedia.org/wiki/Hypergeometric_distribution#Multivariate_hypergeometric_distribution)
+generalizes the *hypergeometric distribution*. Consider ``n`` draws from a finite population containing ``k`` types of elements. Suppose that the population has size ``M`` and there are ``m_i`` elements of type ``i`` for ``i = 1, .., k`` with ``m_1+...m_k = M``. Let ``X = (X_1, ..., X_k)`` where ``X_i`` represents the number of elements of type ``i`` drawn, then the distribution of ``X`` is a multivariate hypergeometric distribution. Each sample of a multivariate hypergeometric distribution is a ``k``-dimensional integer vector that sums to ``n`` and satisfies ``0 \\le X_i \\le m_i``. 
+
+
+The probability mass function is given by
+
+```math
+f(x; m, n) = {{{m_1 \\choose x_1}{m_2 \\choose x_2}\\cdots {m_k \\choose x_k}}\\over {N \\choose n}}, 
+\\quad x_1 + \\cdots + x_k = n, \\quad 0 \\le x_i \\le m_i
+```
+
+```julia
+MvHypergeometric(m, n)   # Multivariate hypergeometric distribution for a population with
+                         # m[i] elements of types i and n draws
+```
+"""
+struct MvHypergeometric <: DiscreteMultivariateDistribution
+    m::Vector{Int}  # number of elements of each type
+    n::Int                  # number of draws
+    function MvHypergeometric(m::Vector{Int}, n::Int; check_args::Bool=true)
+        @check_args(
+            MvHypergeometric,
+            (m, all(x -> x >= 0, m)),
+            zero(n) <= n <= sum(m),
+        )
+        new(m, n)
+    end
+end
+
+
+# Parameters
+
+ncategories(d::MvHypergeometric) = length(d.m)
+length(d::MvHypergeometric) = ncategories(d)
+ntrials(d::MvHypergeometric) = d.n
+
+params(d::MvHypergeometric) = (d.m, d.n)
+partype(::MvHypergeometric) = Int
+
+# Statistics
+
+mean(d::MvHypergeometric) = d.n .* d.m ./ sum(d.m)
+
+function var(d::MvHypergeometric)
+    m = d.m
+    n = ntrials(d)
+    M = sum(m)
+    f = n * (M - n) / (M - 1)
+    v = let f = f
+        map(mi -> f * (mi / M) * ((M - mi) / M), m)
+    end
+    v
+end
+
+function cov(d::MvHypergeometric)
+    m = d.m
+    n = ntrials(d)
+    M = sum(m)
+    p = m / M
+    f = n * (M - n) / (M - 1)
+
+    C = -f * (p * p')
+    C[diagind(C)] .= f .* p .* (1 .- p)
+
+    C
+end
+
+
+# Evaluation
+function insupport(d::MvHypergeometric, x::AbstractVector{<:Real})
+    return length(x) == length(d) && (eltype(x) <: Integer || all(isinteger, x)) && all(((xi, mi),) -> zero(xi) <= xi <= mi, zip(x, d.m)) && sum(x) == ntrials(d)
+end
+
+function _logpdf(d::MvHypergeometric, x::AbstractVector{<:Real})
+    m = d.m
+    M = sum(m)
+    n = ntrials(d)
+
+    s = -logabsbinomial(M, n)[1]
+    for i = 1:length(m)
+        xi = x[i]
+        mi = m[i]
+        ((typeof(xi) <: Integer || isinteger(xi)) && (zero(xi) <= xi <= mi)) || return -Float64(Inf)
+        s += logabsbinomial(mi, xi)[1]
+        n -= xi
+    end
+    (n == 0) || return -Float64(Inf)
+    return s
+end
+
+# Sampling is performed by sequentially sampling each entry from the
+# hypergeometric distribution
+function _rand!(rng::AbstractRNG, d::MvHypergeometric, x::AbstractVector{<:Real})
+    k = length(d)
+    n = ntrials(d)
+    m = d.m
+    length(x) == k || throw(DimensionMismatch("Invalid argument dimension."))
+
+    M = sum(m)
+    i = 0
+    km1 = k - 1
+
+    while i < km1 && n > 0
+        i += 1
+        mi = m[i]
+        # Sample from hypergeometric distribution. Element of type i are 
+        # considered successes. All other elements are considered failures.
+        xi = rand(rng, Hypergeometric(mi, M - mi, n))
+        x[i] = xi
+        # Remove elements of type i from population and group to be sampled.
+        n -= xi
+        M -= mi
+    end
+
+    if i == km1
+        x[k] = n
+    else  # n must have been zero.
+        z = zero(eltype(x))
+        for j = i+1:k
+            x[j] = z
+        end
+    end
+
+    return x
+end
+
+
+

src/multivariates.jl

@@ -126,6 +126,7 @@ for fname in ["dirichlet.jl",
               "mvlognormal.jl",
               "mvtdist.jl",
               "product.jl", # deprecated
-              "vonmisesfisher.jl"]
+              "vonmisesfisher.jl",
+              "mvhypergeometric.jl"]
     include(joinpath("multivariate", fname))
 end

test/multivariate/mvhypergeometric.jl

@@ -0,0 +1,135 @@
+# Tests for Multivariate Hypergeometric
+
+using Distributions
+using Test
+
+@testset "Multivariate Hypergeometric" begin
+    @test_throws DomainError MvHypergeometric([5, 3, -2], 4)
+    @test_throws ArgumentError MvHypergeometric([5, 3], 10)
+
+    m = [5, 3, 2]
+    n = 4
+    d = MvHypergeometric(m, n)
+    @test length(d) == 3
+    @test d.n == n
+    @test d.m == m
+    @test ncategories(d) == length(m)
+    @test params(d) == (m, n)
+    @test partype(d) == Int
+
+    @test mean(d) ≈ [2.0, 1.2, 0.8]
+    @test var(d) ≈ [2 / 3, 56 / 100, 32 / 75]
+
+    covmat = cov(d)
+    @test covmat ≈ (8 / 3) .* [1/4 -3/20 -1/10; -3/20 21/100 -3/50; -1/10 -3/50 4/25]
+
+    @test insupport(d, [2, 1, 1])
+    @test !insupport(d, [3, 2, 1])
+    @test !insupport(d, [0, 0, 4])
+
+
+    # random sampling
+    x = rand(d)
+    @test isa(x, Vector{Int})
+    @test sum(x) == n
+    @test all(x .>= 0)
+    @test all(x .<= m)
+    @test insupport(d, x)
+
+    x = rand(d, 100)
+    @test isa(x, Matrix{Int})
+    @test all(sum(x, dims=1) .== n)
+    @test all(x .>= 0)
+    @test all(x .<= m)
+    @test all(insupport(d, x))
+
+    # random sampling with many catergories
+    m = [20, 2, 2, 2, 1, 1, 1]
+    n = 5
+    d2 = MvHypergeometric(m, n)
+    x = rand(d2)
+    @test isa(x, Vector{Int})
+    @test sum(x) == n
+    @test all(x .>= 0)
+    @test all(x .<= m)
+    @test insupport(d2, x)
+
+    # random sampling with a large category
+    m = [2, 1000]
+    n = 5
+    d3 = MvHypergeometric(m, n)
+    x = rand(d3)
+    @test isa(x, Vector{Int})
+    @test sum(x) == n
+    @test all(x .>= 0)
+    @test all(x .<= m)
+    @test insupport(d3, x)
+
+    # log pdf
+    x = [2, 1, 1]
+    @test pdf(d, x) ≈ 2 / 7
+    @test logpdf(d, x) ≈ log(2 / 7)
+    @test logpdf(d, x) ≈ log(pdf(d, x))
+    @test logpdf(d, [2.5, 0.5, 1]) == -Inf
+
+    x = rand(d, 100)
+    pv = pdf(d, x)
+    lp = logpdf(d, x)
+    for i in 1:size(x, 2)
+        @test pv[i] ≈ pdf(d, x[:, i])
+        @test lp[i] ≈ logpdf(d, x[:, i])
+    end
+
+    # test degenerate cases
+    d1 = MvHypergeometric([1], 1)
+    @test logpdf(d1, [1]) ≈ 0
+    @test logpdf(d1, [0]) == -Inf
+    d2 = MvHypergeometric([2, 0], 1)
+    @test logpdf(d2, [1, 0]) ≈ 0
+    @test logpdf(d2, [0, 1]) == -Inf
+
+    d3 = MvHypergeometric([5, 0, 0, 0], 3)
+    @test logpdf(d3, [3, 0, 0, 0]) ≈ 0
+    @test logpdf(d3, [2, 1, 0, 0]) == -Inf
+    @test logpdf(d3, [2, 0, 0, 0]) == -Inf
+
+    # behavior with n = 0
+    d0 = MvHypergeometric([5, 3, 2], 0)
+    @test logpdf(d0, [0, 0, 0]) ≈ 0
+    @test logpdf(d0, [1, 0, 0]) == -Inf
+
+    @test rand(d0) == [0, 0, 0]
+    @test mean(d0) == [0.0, 0.0, 0.0]
+    @test var(d0) == [0.0, 0.0, 0.0]
+    @test insupport(d0, [0, 0, 0])
+    @test !insupport(d0, [1, 0, 0])
+    @test length(d0) == 3
+
+    # compare with hypergeometric
+    ns = 3
+    nf = 5
+    n = 4
+    dh1 = MvHypergeometric([ns, nf], n)
+    dh2 = Hypergeometric(ns, nf, n)
+
+    x = 2
+    @test pdf(dh1, [x, n - x]) ≈ pdf(dh2, x)
+    x = 3
+    @test pdf(dh1, [x, n - x]) ≈ pdf(dh2, x)
+
+    # comparing marginals to hypergeometric
+    m = [5, 3, 2]
+    n = 4
+    d = MvHypergeometric(m, n)
+    dh = Hypergeometric(m[1], sum(m[2:end]), n)
+    x1 = 2
+    @test pdf(dh, x1) ≈ sum([pdf(d, [x1, x2, n - x1 - x2]) for x2 in 0:m[2]])
+
+    # comparing conditionals to hypergeometric
+    x1 = 2
+    dh = Hypergeometric(m[2], m[3], n - x1)
+    q = sum([pdf(d, [x1, x2, n - x1 - x2]) for x2 in 0:m[2]])
+    for x2 = 0:m[2]
+        @test pdf(dh, x2) ≈ pdf(d, [x1, x2, n - x1 - x2]) / q
+    end
+end

test/runtests.jl

@@ -100,6 +100,7 @@ const tests = [
     "univariate/continuous/triangular",
     "statsapi",
     "univariate/continuous/inversegaussian",
+    "multivariate/mvhypergeometric",
 
     ### missing files compared to /src:
     # "common",