Merge pull request #16 from ReactiveBayes/general-contingency

Generalize Contingency
ReactiveBayes · Jul 9, 2024 · 0aa873c · 0aa873c
2 parents cc2a479 + 6eda8ad
commit 0aa873c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git 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
@@ -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