Merge pull request #40 from aisopous/ai/mixture

Add support for mixtures

Project.toml

@@ -1,7 +1,7 @@
 name = "KeyedDistributions"
 uuid = "2576fb08-064d-4cab-b15d-8dda7fcb9a6d"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.14"
+version = "0.1.15"
 
 [deps]
 AutoHashEquals = "15f4f7f2-30c1-5605-9d31-71845cf9641f"

src/KeyedDistributions.jl

@@ -10,7 +10,7 @@ using PDMatsExtras: submat
 using Random: AbstractRNG
 
 export KeyedDistribution, KeyedSampleable
-export KeyedMvNormal, KeyedGenericMvTDist, MvNormalLike, MvTLike
+export KeyedMvNormal, KeyedGenericMvTDist, MvNormalLike, MvTLike, KeyedMixtureModel, MixtureModelLike
 export axiskeys, distribution
 
 for T in (:Distribution, :Sampleable)
@@ -110,6 +110,9 @@ const MvNormalLike = Union{MvNormal, KeyedMvNormal}
 const KeyedGenericMvTDist = KeyedDistribution{Multivariate, Continuous, <:GenericMvTDist}
 const MvTLike = Union{GenericMvTDist, KeyedGenericMvTDist}
 
+const KeyedMixtureModel = KeyedDistribution{<:VariateForm, <:ValueSupport, <:AbstractMixtureModel}
+const MixtureModelLike = Union{MixtureModel, KeyedMixtureModel}
+
 # Use submat to preserve the covariance matrix PDMat type
 function Base.getindex(d::KeyedMvNormal, i::Vector)::KeyedMvNormal
     return KeyedDistribution(MvNormal(d.d.μ[i], submat(d.d.Σ, i)), axiskeys(d)[1][i])
@@ -123,6 +126,23 @@ function Base.getindex(d::KeyedGenericMvTDist, i::Vector)::KeyedGenericMvTDist
     return KeyedDistribution(GenericMvTDist(d.d.df, d.d.μ[i], submat(d.d.Σ, i)), axiskeys(d)[1][i])
 end
 
+function _marginalize(d::MvNormal, i::Vector)::MvNormal
+    return MvNormal(d.μ[i], submat(d.Σ, i))
+end
+
+function _marginalize(d::MvTDist, i::Vector)::MvTDist
+    T = typeof(d)
+    return MvTDist(d.df, d.μ[i], submat(d.Σ, i))
+end
+
+function Base.getindex(mm::KeyedMixtureModel, i::Vector)::KeyedMixtureModel
+    margcomps = map(Distributions.components(mm)) do c
+        _marginalize(c, i)
+    end
+    return KeyedDistribution(MixtureModel(margcomps), only(axiskeys(mm))[i])
+end
+
+
 # Access methods
 
 """
@@ -205,6 +225,12 @@ function Distributions._logpdf(d::KeyedDistribution, x::AbstractArray)
     return Distributions._logpdf(unkeyed_dist, x)
 end
 
+function Distributions._logpdf(d::KeyedMixtureModel, x::AbstractArray)
+    # Assume components of KeyedMixtureModel are unkeyed.
+    unkeyed_dist = distribution(d)
+    return Distributions._logpdf(unkeyed_dist, x)
+end
+
 _maybe_parent(x) = x
 _maybe_parent(x::AbstractArray) = parent(x)
 
@@ -276,6 +302,53 @@ function Distributions.insupport(d::KeyedDistribution{<:Univariate}, x::Real)
     return insupport(distribution(d), x)
 end
 
+function Distributions.MixtureModel(
+    cs::Vector{C}, pri::CT
+) where {C<:KeyedDistribution,CT<:Distributions.Categorical}
+    VF = Distributions.variate_form(C)
+    VS = Distributions.value_support(C)
+    k = cs[1].keys
+    all(==(k), (c.keys for c in cs)) ||
+        error("Keys of all mixture components must be the same.")
+    length(cs) == ncategories(pri) ||
+        error("The number of components does not match the length of prior.")
+    return KeyedDistribution(MixtureModel(distribution.(cs), pri), k)
+end
+
+function Distributions.MixtureModel(
+    cs::Vector{C}, pri::CT
+) where {C<:KeyedDistribution,CT<:AbstractVector{<:Real}}
+    return MixtureModel(cs, Categorical(pri))
+end
+
+function KeyedMixtureModel(
+    cs::Vector{<:KeyedDistribution},
+    pri::Union{AbstractVector{<:Real},Distributions.Categorical},
+)
+    return MixtureModel(cs, pri)
+end
+
+function KeyedMixtureModel(mm::MixtureModel, keys::Tuple{Vararg{AbstractVector}})
+    return KeyedDistribution(mm, keys)
+end
+
+# Avoid the double wrap
+function KeyedDistribution(
+    kd::KeyedDistribution, keys::Tuple{Vararg{AbstractVector{T} where T,N} where N}
+)
+    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)))
+        _marginalize(c, inds)
+    end
+    return KeyedDistribution(MixtureModel(margcomps), keys)
+end
+
 # Overload equality comparison between `KeyedT` and underlying `T`
 for T in (:Distribution, :Sampleable)
     KeyedT = Symbol(:Keyed, T)

test/runtests.jl

@@ -46,6 +46,8 @@ using Test
             @test MvNormalLike <: MultivariateDistribution
             @test MvTLike <: MultivariateDistribution
 
+            @test KeyedMixtureModel <: Distribution
+
         end
 
         @testset "1-dimensional constructor" begin
@@ -369,6 +371,49 @@ using Test
 
             @test d([1]) == d[[1]] == KeyedDistribution(GenericMvTDist(3, m[[1]], submat(W, [1])), [1])
         end
+
+        @testset "Mixture of KeyedMvNormal{AbstractArray, PDMats}" begin
+            keys = [:a, :b, :c]
+            d = KeyedDistribution(MvNormal(m, W), keys)
+            pri = [0.5, 0.5]
+            mm = KeyedDistribution(MixtureModel([d.d, d.d], pri), keys)
+            mmk = MixtureModel([d, d], pri)
+            kmm1 = KeyedMixtureModel([d, d], pri)
+            kmm2 = KeyedMixtureModel(mm.d, axiskeys(d))
+            kmm3 = KeyedDistribution(mm, axiskeys(mm))
+            @test mm == mmk
+            @test kmm1 == kmm2
+            @test kmm2 == kmm1
+            @test kmm3 == kmm2
+            @test Distributions.logpdf(mm, zeros(3)) == Distributions.logpdf(d, zeros(3))
+            @test cov(mm) == cov(d)
+            @test cov(mm([:a, :b, :c])) ≈ cov(d)
+            @test cov(mm([:a, :c])) ≈ cov(d[[1, 3]])
+            @test cov(mm([:a])) ≈ cov(d[[1]])
+            @test only(cov(KeyedDistributions._marginalize(distribution(d), [2]))) ≈ cov(d)[2, 2]
+            @test mean(d[[2]]) == mean(mm[[2]])
+            @test cov(d[[2]]) == cov(mm[[2]])
+            @test mm isa MixtureModelLike
+        end
+
+        @testset "Mixture of KeyedMvTDist{AbstractArray, PDMats}" begin
+            keys = [:a, :b, :c]
+            d = KeyedDistribution(MvTDist(2.5, m, Matrix(Symmetric(W))), keys)
+            pri = [0.5, 0.5]
+            mm = KeyedDistribution(MixtureModel([d.d, d.d], pri), keys)
+            mmk = MixtureModel([d, d], pri)
+            kmm1 = KeyedMixtureModel([d, d], pri)
+            kmm2 = KeyedMixtureModel(mm.d, axiskeys(d))
+            @test mm == mmk
+            @test kmm1 == kmm2
+            @test Distributions.logpdf(mm, zeros(3)) ≈ Distributions.logpdf(d, zeros(3))
+            @test cov(mm) == cov(d)
+            @test cov(mm([:a, :b, :c])) ≈ cov(d)
+            @test cov(mm([:a, :c])) ≈ cov(d[[1, 3]])
+            @test cov(mm([:c])) ≈ cov(d[[3]])
+            @test only(cov(KeyedDistributions._marginalize(distribution(d), [2]))) ≈ cov(d)[2, 2]
+            @test mm isa MixtureModelLike
+        end
     end
 
     @testset "NamedDims functions" begin