Merge pull request #51 from aisopous/ai/ncomponents-probs

Add missing Distributions methods for KeyedMixtureModel
invenia · Feb 7, 2023 · 063fd33 · 063fd33 · rofinn · Feb 7, 2023
2 parents dc5d9e0 + 5a0bbbe
commit 063fd33
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "KeyedDistributions"
 uuid = "2576fb08-064d-4cab-b15d-8dda7fcb9a6d"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.16"
+version = "0.1.17"
 
 [deps]
 AutoHashEquals = "15f4f7f2-30c1-5605-9d31-71845cf9641f"

diff --git a/src/KeyedDistributions.jl b/src/KeyedDistributions.jl
@@ -214,6 +214,12 @@ Distributions.sampler(d::KeyedDistribution) = sampler(distribution(d))
 
 Distributions.params(d::KeyedDistOrSampleable) = params(distribution(d))
 
+Distributions.probs(d::KeyedMixtureModel) = probs(distribution(d))
+
+Distributions.components(d::KeyedMixtureModel) = components(distribution(d))
+
+Distributions.ncomponents(d::KeyedMixtureModel) = ncomponents(distribution(d))
+
 Base.eltype(d::KeyedDistribution) = eltype(distribution(d))
 
 function Distributions._logpdf(d::KeyedDistribution, x::AbstractArray)
@@ -341,8 +347,6 @@ function KeyedDistribution(
     return KeyedDistribution(kd.d, keys)
 end
 
-Distributions.components(kd::KeyedMixtureModel) = Distributions.components(kd.d)
-
 function (mm::KeyedMixtureModel)(keys...)
     margcomps = map(Distributions.components(mm)) do c
         inds = first(map(AxisKeys.findindex, keys, axiskeys(mm)))

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -404,6 +404,9 @@ using Test
             @test mean(d[[2]]) == mean(mm[[2]])
             @test cov(d[[2]]) == cov(mm[[2]])
             @test mm isa MixtureModelLike
+            @test probs(mm) == pri
+            @test ncomponents(mm) == 2
+            @test components(mm) == [d, d]
         end
 
         @testset "Mixture of KeyedMvTDist{AbstractArray, PDMats}" begin