Allow <:AbstractArray{<:Real} for DirichletMultinomial.

src/multivariate/dirichletmultinomial.jl

@@ -2,7 +2,7 @@
     DirichletMultinomial
 
 The [Dirichlet-multinomial distribution](https://en.wikipedia.org/wiki/Dirichlet-multinomial_distribution)
-is the distribution of a draw from a multinomial distribution where each sample has a 
+is the distribution of a draw from a multinomial distribution where each sample has a
 slightly different probability vector, drawn from a common Dirichlet distribution.
 
 This contrasts with the multinomial distribution, which assumes that all observations arise
@@ -29,26 +29,27 @@ vector α.
 DirichletMultinomial(n, k) # Dirichlet-multinomial distribution with n trials and parameter
 vector of length k of ones.
 ```
+
+`α` can be any `AbstractVector{<:Real}`.
 """
-struct DirichletMultinomial{T <: Real} <: DiscreteMultivariateDistribution
+struct DirichletMultinomial{T <: Real,V <: AbstractVector{T}} <: DiscreteMultivariateDistribution
     n::Int
-    α::Vector{T}
+    α::V
     α0::T
 
-    function DirichletMultinomial{T}(n::Integer, α::Vector{T}) where T
+    function DirichletMultinomial{T}(n::Integer, α::V) where {T <: Real, V <: AbstractVector{T}}
         α0 = sum(abs, α)
         sum(α) == α0 || throw(ArgumentError("alpha must be a positive vector."))
         n > 0 || throw(ArgumentError("n must be a positive integer."))
-        new{T}(Int(n), α, α0)
+        new{T,V}(Int(n), α, α0)
     end
 end
-DirichletMultinomial(n::Integer, α::Vector{T}) where {T <: Real} = DirichletMultinomial{T}(n, α)
-DirichletMultinomial(n::Integer, α::Vector{T}) where {T <: Integer} = DirichletMultinomial(n, float(α))
+DirichletMultinomial(n::Integer, α::AbstractVector{T}) where {T <: Real} = DirichletMultinomial{T}(n, α)
+DirichletMultinomial(n::Integer, α::AbstractVector{T}) where {T <: Integer} = DirichletMultinomial(n, float(α))
 DirichletMultinomial(n::Integer, k::Integer) = DirichletMultinomial(n, ones(k))
 
 Base.show(io::IO, d::DirichletMultinomial) = show(io, d, (:n, :α,))
 
-
 # Parameters
 ncategories(d::DirichletMultinomial) = length(d.α)
 length(d::DirichletMultinomial) = ncategories(d)

test/multivariate/dirichletmultinomial.jl

@@ -2,7 +2,7 @@
 
 
 using Distributions
-using Test, Random, SpecialFunctions
+using Test, Random, SpecialFunctions, StaticArrays
 
 Random.seed!(123)
 
@@ -44,18 +44,23 @@ d = fit(DirichletMultinomial, x)
 @test isapprox(cov(d) , cov(x, dims=2)      , atol=.5)
 
 # test Evaluation
-d = DirichletMultinomial(10, 5)
-@test typeof(d) == DirichletMultinomial{Float64}
-@test !insupport(d, func[1](5))
-@test insupport(d, [2, 2, 2, 2, 2])
-@test insupport(d, 2.0 * ones(5))
-@test !insupport(d, 3.0 * ones(5))
-
-for x in (2 * ones(5), [1, 2, 3, 4, 0], [3.0, 0.0, 3.0, 0.0, 4.0], [0, 0, 0, 0, 10])
-    @test pdf(d, x) ≈
-        factorial(d.n) * gamma(d.α0) / gamma(d.n + d.α0) * prod(gamma.(d.α .+ x) ./ (gamma.(x .+ 1) .* gamma.(d.α)))
-    @test logpdf(d, x) ≈
-        logfactorial(d.n) + loggamma(d.α0) - loggamma(d.n + d.α0) + sum(loggamma, d.α + x) - sum(loggamma, d.α) - sum(loggamma.(x .+ 1))
+@testset "Dirichlet evaluation and non-Vector arguments" begin
+    αS = SVector(0.1, 0.7, 9.2, 3.7, 0.4)
+    for (d, T) in (DirichletMultinomial(10, 5) => DirichletMultinomial{Float64,Vector{Float64}},
+                   DirichletMultinomial(10, αS) => DirichletMultinomial{Float64,typeof(αS)})
+        @test typeof(d) == T
+        @test !insupport(d, func[1](5))
+        @test insupport(d, [2, 2, 2, 2, 2])
+        @test insupport(d, 2.0 * ones(5))
+        @test !insupport(d, 3.0 * ones(5))
+
+        for x in (2 * ones(5), [1, 2, 3, 4, 0], [3.0, 0.0, 3.0, 0.0, 4.0], [0, 0, 0, 0, 10])
+            @test pdf(d, x) ≈
+                factorial(d.n) * gamma(d.α0) / gamma(d.n + d.α0) * prod(gamma.(d.α .+ x) ./ (gamma.(x .+ 1) .* gamma.(d.α)))
+            @test logpdf(d, x) ≈
+                logfactorial(d.n) + loggamma(d.α0) - loggamma(d.n + d.α0) + sum(loggamma, d.α + x) - sum(loggamma, d.α) - sum(loggamma.(x .+ 1))
+        end
+    end
 end
 
 # test Sampling