diff --git a/src/densities/contingency.jl b/src/densities/contingency.jl index 5f5f576..6c579b7 100644 --- a/src/densities/contingency.jl +++ b/src/densities/contingency.jl @@ -22,25 +22,25 @@ A `Contingency` distribution over more than two variables requires higher-order - `renormalize`, optional, supports either `Val(true)` or `Val(false)`, specifies whether matrix `P` must be automatically renormalized. Does not modify the original `P` and allocates a new one for the renormalized version. If set to `false` the contingency matrix `P` **must** be normalized by hand, otherwise the result of related calculations might be wrong """ -struct Contingency{T,P<:AbstractMatrix{T}} +struct Contingency{T,P<:AbstractArray{T}} p::P - Contingency{T,P}(A::AbstractMatrix) where {T,P<:AbstractMatrix{T}} = new(A) + Contingency{T,P}(A::AbstractArray) where {T,P<:AbstractArray{T}} = new(A) end -Contingency(P::AbstractMatrix) = Contingency(P, Val(true)) +Contingency(P::AbstractArray) = Contingency(P, Val(true)) -function Contingency(P::M, ::Val{true}) where {T,M<:AbstractMatrix{T}} +function Contingency(P::M, ::Val{true}) where {T,M<:AbstractArray{T}} return Contingency{T,M}(P ./ sum(P)) end -Contingency(P::M, ::Val{false}) where {T,M<:AbstractMatrix{T}} = Contingency{T,M}(P) +Contingency(P::M, ::Val{false}) where {T,M<:AbstractArray{T}} = Contingency{T,M}(P) BayesBase.components(distribution::Contingency) = distribution.p BayesBase.component(distribution::Contingency, k) = distribution.p[:, k] -function BayesBase.vague(::Type{<:Contingency}, dims::Int) - return Contingency(ones(dims, dims) ./ abs2(dims)) +function BayesBase.vague(::Type{<:Contingency}, dims::Int, nvars::Int=2) + return Contingency(ones([dims for _ in 1:nvars]...)) end BayesBase.paramfloattype(distribution::Contingency) = deep_eltype(components(distribution)) diff --git a/test/densities/contingency_tests.jl b/test/densities/contingency_tests.jl index 77310c7..897c7ab 100644 --- a/test/densities/contingency_tests.jl +++ b/test/densities/contingency_tests.jl @@ -20,6 +20,10 @@ end @test typeof(d2) <: Contingency @test components(d2) ≈ ones(4, 4) ./ 16 + + d3 = vague(Contingency, 5, 3) + @test typeof(d3) <: Contingency + @test components(d3) ≈ ones(5, 5, 5) ./ 125 end @testitem "Contingency: entropy" begin @@ -31,6 +35,17 @@ end @test entropy(Contingency(10.0 * [0.09 0.00; 0.00 0.91])) ≈ 0.30253782309749805 @test !isnan(entropy(Contingency([0.0 1.0; 1.0 0.0]))) @test !isinf(entropy(Contingency([0.0 1.0; 1.0 0.0]))) + + @test entropy( + Contingency( + stack([ + 0.3 * [0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3], + 0.7 * [0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3], + ],), + ), + ) ≈ 2.6390313416381166 + + @test entropy(Contingency(ones(2, 2, 2))) == 2.0794415416798357 end @testitem "Contingency: isapprox" begin