Matrixdirichlet mixture

src/nodes/predefined.jl

@@ -24,6 +24,7 @@ include("predefined/gamma_mixture.jl")
 include("predefined/dot_product.jl")
 include("predefined/softdot.jl")
 include("predefined/transition.jl")
+include("predefined/transition_mixture.jl")
 include("predefined/autoregressive.jl")
 include("predefined/bifm.jl")
 include("predefined/bifm_helper.jl")

src/nodes/predefined/transition_mixture.jl

@@ -0,0 +1,155 @@
+export TransitionMixture, TransitionMixtureNode
+
+# Transition Mixture Functional Form
+struct TransitionMixture{N} end
+
+ReactiveMP.as_node_symbol(::Type{<:TransitionMixture}) = :TransitionMixture
+
+interfaces(::Type{<:TransitionMixture}) = Val((:out, :in, :switch, :matrices))
+alias_interface(::Type{<:TransitionMixture}, ::Int64, name::Symbol) = name
+is_predefined_node(::Type{<:TransitionMixture}) = PredefinedNodeFunctionalForm()
+sdtype(::Type{<:TransitionMixture}) = Stochastic()
+collect_factorisation(::Type{<:TransitionMixture}, factorization) = TransitionMixtureNodeFactorisation()
+
+struct TransitionMixtureNodeFactorisation end
+
+struct TransitionMixtureNode{N} <: AbstractFactorNode
+    out::NodeInterface
+    in::NodeInterface
+    switch::NodeInterface
+    matrices::NTuple{N, IndexedNodeInterface}
+    local_clusters::FactorNodeLocalClusters
+end
+
+functionalform(factornode::TransitionMixtureNode{N}) where {N} = TransitionMixture{N}
+getinterface(factornode::TransitionMixtureNode{N}, i::Int64) where {N} = getinterfaces(factornode)[i]
+getinterfaces(factornode::TransitionMixtureNode) = (factornode.out, factornode.in, factornode.switch, factornode.matrices...)
+sdtype(factornode::TransitionMixtureNode) = Stochastic()
+
+interfaceindices(factornode::TransitionMixtureNode, iname::Symbol)                       = (interfaceindex(factornode, iname),)
+interfaceindices(factornode::TransitionMixtureNode, inames::NTuple{N, Symbol}) where {N} = map(iname -> interfaceindex(factornode, iname), inames)
+
+function interfaceindex(factornode::TransitionMixtureNode, iname::Symbol)
+    if iname === :out
+        return 1
+    elseif iname === :in
+        return 2
+    elseif iname === :switch
+        return 3
+    elseif iname === :matrices
+        return 4
+    else
+        error("Unknown interface ':$(iname)' for the [ $(functionalform(factornode)) ] node")
+    end
+end
+
+function factornode(::Type{<:TransitionMixture}, interfaces, factorization)
+    @show interfaces factorization
+    outinterface = interfaces[findfirst(((name, variable),) -> name == :out, interfaces)]
+    ininterface = interfaces[findfirst(((name, variable),) -> name == :in, interfaces)]
+    switchinterface = interfaces[findfirst(((name, variable),) -> name == :switch, interfaces)]
+    matricesinterface = filter(((name, variable),) -> name == :matrices, interfaces)
+    noutinterface = NodeInterface(outinterface...)
+    nininterface = NodeInterface(ininterface...)
+    nswitchinterface = NodeInterface(switchinterface...)
+
+    N = length(matricesinterface)
+    nmatricesinterface = ntuple(i -> IndexedNodeInterface(i, NodeInterface(matricesinterface[i]...)), N)
+
+    m_matrices = ntuple(i -> FactorNodeLocalMarginal(Symbol(:matrices, :_, i)), N)
+    marginals = (FactorNodeLocalMarginal(:out_in_switch), m_matrices...)
+
+    m_mf = ntuple(i -> (i + 3,), N)
+    factornodelocalclusters = FactorNodeLocalClusters(marginals, ((1, 2, 3), m_mf...))
+
+    if length(matricesinterface) < 2
+        error("The number of matrices in `TransitionMixture` must be at least 2. Got `$(length(matricesinterface))` matrices instead.")
+    end
+
+    return TransitionMixtureNode(noutinterface, nininterface, nswitchinterface, nmatricesinterface, factornodelocalclusters)
+end
+
+struct TransitionMixtureNodeFunctionalDependencies <: FunctionalDependencies end
+
+collect_functional_dependencies(::TransitionMixtureNode, ::Nothing) = TransitionMixtureNodeFunctionalDependencies()
+collect_functional_dependencies(::TransitionMixtureNode, ::TransitionMixtureNodeFunctionalDependencies) = TransitionMixtureNodeFunctionalDependencies()
+collect_functional_dependencies(::TransitionMixtureNode, ::Any) =
+    error("The functional dependencies for TransitionMixtureNode must be either `Nothing` or `TransitionMixtureNodeFunctionalDependencies`")
+
+function activate!(factornode::TransitionMixtureNode, options::FactorNodeActivationOptions)
+    dependecies = collect_functional_dependencies(factornode, getdependecies(options))
+    ReactiveMP.initialize_clusters!(factornode.local_clusters, DefaultFunctionalDependencies(), factornode, options)
+    return activate!(dependecies, factornode, options)
+end
+
+function functional_dependencies(::TransitionMixtureNodeFunctionalDependencies, factornode::TransitionMixtureNode{N}, interface, iindex::Int) where {N}
+    message_dependencies = if iindex === 1
+        # message_cluster = filter(!=(:out), first(filter((c) -> :out ∈ c, clusters)))
+        (factornode.in, factornode.switch)
+    elseif iindex === 2
+        (factornode.out, factornode.switch)
+    elseif iindex === 3
+        (factornode.out, factornode.in)
+    else
+        ()
+    end
+
+    marginal_dependencies = if iindex === 1 || iindex === 2 || iindex === 3
+        (factornode.matrices,)
+    elseif 3 < iindex <= N + 3
+        (first(factornode.local_clusters.marginals),)
+    else
+        error("Bad index in functional_dependencies for TransitionMixtureNode")
+    end
+
+    return message_dependencies, marginal_dependencies
+end
+
+function collect_latest_messages(::TransitionMixtureNodeFunctionalDependencies, factornode::TransitionMixtureNode{N}, message_dependencies::Tuple{}) where {N}
+    return nothing, of(nothing)
+end
+
+function collect_latest_messages(
+    ::TransitionMixtureNodeFunctionalDependencies, factornode::TransitionMixtureNode{N}, message_dependencies::Tuple{NodeInterface, NodeInterface}
+) where {N}
+    firstvarinterface = message_dependencies[1]
+    secondvarinterface = message_dependencies[2]
+
+    message_names = Val{(name(firstvarinterface), name(secondvarinterface))}()
+    messages_observable = combineLatestUpdates((messagein(firstvarinterface), messagein(secondvarinterface)), PushNew())
+
+    return message_names, messages_observable
+end
+
+function collect_latest_marginals(
+    ::TransitionMixtureNodeFunctionalDependencies, factornode::TransitionMixtureNode{N}, marginal_dependencies::Tuple{NTuple{N, IndexedNodeInterface}}
+) where {N}
+    matricesinterfaces = first(marginal_dependencies)
+
+    marginal_names = Val{(name(matricesinterfaces[1]),)}()
+    marginals_observable =
+        combineLatest((combineLatest(map((mat) -> getmarginal(getvariable(mat), IncludeAll()), matricesinterfaces), PushNew()),), PushNew()) |>
+        map_to((ManyOf(map((mat) -> getmarginal(getvariable(mat), IncludeAll()), matricesinterfaces)),))
+
+    return marginal_names, marginals_observable
+end
+
+function collect_latest_marginals(
+    ::TransitionMixtureNodeFunctionalDependencies, factornode::TransitionMixtureNode{N}, marginal_dependencies::Tuple{NodeInterface, NodeInterface, NodeInterface}
+) where {N}
+    outinterface = marginal_dependencies[1]
+    ininterface = marginal_dependencies[2]
+    switchinterface = marginal_dependencies[3]
+
+    marginal_names       = Val{(name(outinterface), name(ininterface), name(switchinterface))}()
+    marginals_observable = combineLatestUpdates((getmarginal(getvariable(outinterface), IncludeAll()), getmarginal(getvariable(ininterface), IncludeAll()), getmarginal(getvariable(switchinterface), IncludeAll())), PushNew())
+
+    return marginal_names, marginals_observable
+end
+
+# FreeEnergy related functions
+
+# @average_energy TransitionMixture(q_out_in_switch::Any, q_matrices::ManyOf{N, Any}) where {N} = begin
+#     U = 0.0
+#     for i in 1:N
+#         U -=*mean(log, q_matrices[i])

src/rules/predefined.jl

@@ -107,6 +107,12 @@ include("transition/out.jl")
 include("transition/in.jl")
 include("transition/a.jl")
 
+include("transition_mixture/marginals.jl")
+include("transition_mixture/out.jl")
+include("transition_mixture/in.jl")
+include("transition_mixture/switch.jl")
+include("transition_mixture/matrices.jl")
+
 include("continuous_transition/y.jl")
 include("continuous_transition/x.jl")
 include("continuous_transition/a.jl")

src/rules/transition_mixture/in.jl

@@ -0,0 +1,19 @@
+@rule TransitionMixture{N}(:in, Marginalisation) (m_out::Categorical, m_switch::Categorical, m_matrices::ManyOf{N, PointMass}) where {N} = begin
+    πs = probvec(m_switch)
+    ndims = length(probvec(m_out))
+    a = tiny * ones(ndims)
+    for i in 1:N
+        a += πs[i] * mean(m_matrices[i])' * probvec(m_out)
+    end
+    return Categorical(a ./ sum(a))
+end
+
+@rule TransitionMixture{N}(:in, Marginalisation) (m_out::Categorical, m_switch::Categorical, q_matrices::ManyOf{N, Union{MatrixDirichlet, PointMass}}) where {N} = begin
+    πs = probvec(m_switch)
+    ndims = length(probvec(m_out))
+    a = tiny * ones(ndims)
+    for i in 1:N
+        a += πs[i] * clamp.(exp.(mean(Base.Broadcast.BroadcastFunction(log), q_matrices[i])), tiny, Inf)' * probvec(m_out)
+    end
+    return Categorical(a ./ sum(a))
+end

src/rules/transition_mixture/marginals.jl

@@ -0,0 +1,31 @@
+@marginalrule TransitionMixture{N}(:out_in_switch) (
+    m_out::Categorical, m_in::Categorical, m_switch::Categorical, q_matrices::ManyOf{N, Union{MatrixDirichlet, PointMass}}
+) where {N} = begin
+    πs = probvec(m_switch)
+    π_m_out = probvec(m_out)
+    π_m_in = probvec(m_in)
+    B = Array{Float64, 3}(undef, N, length(π_m_out), length(π_m_in))
+    for i in 1:N
+        B[i, :, :] = (πs[i] * diagm(π_m_out) * clamp.(exp.(mean(Base.Broadcast.BroadcastFunction(log), q_matrices[i])), tiny, Inf) * diagm(π_m_in))
+    end
+    return Contingency(B)
+end
+
+function marginalrule(
+    ::Type{<:TransitionMixture{N}}, ::Val{:out_in_switch}, ::Val{(:out, :in, :switch)}, messages, qnames::Val{Q}, marginals, meta::Nothing, __node::Nothing
+) where {N, Q}
+    @assert length(Q) == length(marginals) "The length of `qnames` must match the length of `marginals`"
+    @assert N == length(marginals) "The `N` must match the length of `marginals`"
+    m_out = getdata(messages[1])
+    m_in = getdata(messages[2])
+    m_switch = getdata(messages[3])
+    q_matrices = getdata.(marginals)
+    πs = probvec(m_switch)
+    π_m_out = probvec(m_out)
+    π_m_in = probvec(m_in)
+    B = Array{Float64, 3}(undef, N, length(π_m_out), length(π_m_in))
+    for i in 1:N
+        B[i, :, :] = (πs[i] * diagm(π_m_out) * clamp.(exp.(mean(Base.Broadcast.BroadcastFunction(log), q_matrices[i])), tiny, Inf) * diagm(π_m_in))
+    end
+    return Contingency(B)
+end

src/rules/transition_mixture/matrices.jl

@@ -0,0 +1,3 @@
+@rule TransitionMixture((:matrices, k), Marginalisation) (q_out_in_switch::Contingency,) = begin
+    return MatrixDirichlet(components(q_out_in_switch)[k, :, :] .+ 1)
+end

src/rules/transition_mixture/out.jl

@@ -0,0 +1,19 @@
+@rule TransitionMixture{N}(:out, Marginalisation) (m_in::Categorical, m_switch::Categorical, m_matrices::ManyOf{N, PointMass}) where {N} = begin
+    πs = probvec(m_switch)
+    ndims = length(probvec(m_in))
+    a = tiny * ones(ndims)
+    for i in 1:N
+        a += πs[i] * mean(m_matrices[i]) * probvec(m_in)
+    end
+    return Categorical(a ./ sum(a))
+end
+
+@rule TransitionMixture{N}(:out, Marginalisation) (m_in::Categorical, m_switch::Categorical, q_matrices::ManyOf{N, Union{MatrixDirichlet, PointMass}}) where {N} = begin
+    πs = probvec(m_switch)
+    ndims = length(probvec(m_in))
+    a = tiny * ones(ndims)
+    for i in 1:N
+        a += πs[i] * clamp.(exp.(mean(Base.Broadcast.BroadcastFunction(log), q_matrices[i])), tiny, Inf) * probvec(m_in)
+    end
+    return Categorical(a ./ sum(a))
+end

src/rules/transition_mixture/switch.jl

@@ -0,0 +1,15 @@
+@rule TransitionMixture{N}(:switch, Marginalisation) (m_out::Categorical, m_in::Categorical, m_matrices::ManyOf{N, PointMass}) where {N} = begin
+    a = zeros(N)
+    for k in 1:N
+        a[k] = probvec(m_out)' * mean(m_matrices[k]) * probvec(m_in)
+    end
+    return Categorical(a ./ sum(a))
+end
+
+@rule TransitionMixture{N}(:switch, Marginalisation) (m_out::Categorical, m_in::Categorical, q_matrices::ManyOf{N, Union{MatrixDirichlet, PointMass}}) where {N} = begin
+    a = zeros(N)
+    for k in 1:N
+        a[k] = probvec(m_out)' * clamp.(exp.(mean(Base.Broadcast.BroadcastFunction(log), q_matrices[k])), tiny, Inf) * probvec(m_in)
+    end
+    return Categorical(a ./ sum(a))
+end

test/rules/transition_mixture/in_tests.jl

@@ -0,0 +1,28 @@
+@testitem "rules:TransitionMixture:in" begin
+    using ReactiveMP, BayesBase, Random, ExponentialFamily, Distributions
+
+    import ReactiveMP: @test_rules
+    import MatrixCorrectionTools: NoCorrection, AddToDiagonalEntries, ReplaceZeroDiagonalEntries
+
+    @testset "Belief Propagation: (m_out::Categorical, m_switch::Categorical, m_matrices::ManyOf{N, PointMass})" begin
+        @test_rules [check_type_promotion = false] TransitionMixture{2}(:in, Marginalisation) [(
+            input = (m_out = Categorical(0.2, 0.8), m_switch = Categorical([1.0, 0.0]), m_matrices = ManyOf(PointMass([0.1 0.2; 0.9 0.8]), PointMass([0.3 0.4; 0.7 0.6]))),
+            output = Categorical([0.5211267605633803, 0.4788732394366197])
+        )]
+    end
+
+    @testset "Variational Messages: (m_out::Categorical, m_switch::Categorical, q_matrices::ManyOf{N, MatrixDirichlet})" begin
+        @test_rules [check_type_promotion = false] TransitionMixture{2}(:in, Marginalisation) [
+            (
+                input = (m_out = Categorical(0.2, 0.8), m_switch = Categorical([0.3, 0.7]), q_matrices = ManyOf(PointMass([0.1 0.2; 0.9 0.8]), PointMass([0.3 0.4; 0.7 0.6]))),
+                output = Categorical([0.5239616613418148, 0.4760383386581853])
+            ),
+            (
+                input = (
+                    m_out = Categorical(0.1, 0.9), m_switch = Categorical([0.3, 0.7]), q_matrices = ManyOf(MatrixDirichlet([0.1 0.2; 0.9 0.8]), MatrixDirichlet([0.3 0.4; 0.7 0.6]))
+                ),
+                output = Categorical([0.5641249477901458, 0.4358750522098543])
+            )
+        ]
+    end
+end

test/rules/transition_mixture/marginals_tests.jl

@@ -0,0 +1,75 @@
+@testitem "marginalrules:TransitionMixture" begin
+    using ReactiveMP, BayesBase, Random, ExponentialFamily, Distributions
+
+    import ReactiveMP: @test_marginalrules
+    @testset "out_in_switch: (m_out::Categorical, m_in::Categorical, m_switch::Categorical, q_matrices::ManyOf{N, Union{PointMass, MatrixDirichlet}})" begin
+        @test ReactiveMP.marginalrule(
+            TransitionMixture{2},
+            Val(:out_in_switch),
+            Val((:out, :in, :switch)),
+            (
+                ReactiveMP.Message(Categorical(0.2, 0.8), false, true, nothing),
+                ReactiveMP.Message(Categorical(0.1, 0.9), false, true, nothing),
+                ReactiveMP.Message(Categorical(0.3, 0.7), false, true, nothing)
+            ),
+            Val((:matrices_1, :matrices_2)),
+            (ReactiveMP.Marginal(PointMass([0.1 0.2; 0.9 0.8]), false, true, nothing), ReactiveMP.Marginal(PointMass([0.3 0.4; 0.7 0.6]), false, true, nothing)),
+            Nothing(),
+            Nothing()
+        ) ≈ Contingency(
+            permutedims(
+                stack([
+                    [0.0009966777408637875 0.017940199335548173; 0.035880398671096346 0.28704318936877077],
+                    [0.0069767441860465115 0.08372093023255814; 0.06511627906976744 0.5023255813953489]
+                ]),
+                (3, 1, 2)
+            )
+        )
+
+        @test ReactiveMP.marginalrule(
+            TransitionMixture{2},
+            Val(:out_in_switch),
+            Val((:out, :in, :switch)),
+            (
+                ReactiveMP.Message(Categorical(0.2, 0.4, 0.4), false, true, nothing),
+                ReactiveMP.Message(Categorical(0.1, 0.5, 0.4), false, true, nothing),
+                ReactiveMP.Message(Categorical(0.3, 0.7), false, true, nothing)
+            ),
+            Val((:matrices_1, :matrices_2)),
+            (
+                ReactiveMP.Marginal(
+                    MatrixDirichlet(
+                        [
+                            3.759109596106528 7.651848771355311 9.463024296540873
+                            0.05227997457759148 7.63832734893159 4.123039206788142
+                            1.7218244939296623 0.8732644376397225 5.797052273450543
+                        ]
+                    ),
+                    false,
+                    true,
+                    nothing
+                ),
+                ReactiveMP.Marginal(
+                    PointMass(
+                        [
+                            9.712482722213574 9.112005385625451 9.022921164635415
+                            2.158924800523061 9.014480452856436 7.641284931051801
+                            2.3506858907854378 5.462580545036949 5.43795594079331
+                        ]
+                    ),
+                    false,
+                    true,
+                    nothing
+                )
+            ),
+            Nothing(),
+            Nothing()
+        ) ≈ Contingency{Float64, Array{Float64, 3}}(
+            [
+                0.0007817591287544585 1.4372761803507031e-12 0.0005989009423266854; 0.027299775547839533 0.012136580144260207 0.013214581490095035;;;
+                0.0027518963746678685 0.0054934040277940565 0.00034292322803211974; 0.12805979116407934 0.2533783586358652 0.15354181526669014;;;
+                0.0022880577360659344 0.0018546054978035274 0.002707028738712131; 0.10144624381859718 0.17182454330668553 0.12227973495029375
+            ]
+        )
+    end
+end

test/rules/transition_mixture/matrices_tests.jl

@@ -0,0 +1,16 @@
+@testitem "rules:TransitionMixture:matrices" begin
+    using ReactiveMP, BayesBase, Random, ExponentialFamily, Distributions
+
+    import ReactiveMP: @test_rules
+    import MatrixCorrectionTools: NoCorrection, AddToDiagonalEntries, ReplaceZeroDiagonalEntries
+
+    @testset "Variational Messages: (q_out_in_switch::Contingency)" begin
+        @test_rules [check_type_promotion = false] TransitionMixture{2}((:matrices, k = 1), Marginalisation) [(
+            input = (q_out_in_switch = Contingency([0.03 0.06; 0.07 0.14;;; 0.18 0.03; 0.42 0.07]),), output = MatrixDirichlet([1.03 1.18; 1.06 1.03])
+        )]
+
+        @test_rules [check_type_promotion = false] TransitionMixture{2}((:matrices, k = 2), Marginalisation) [(
+            input = (q_out_in_switch = Contingency([0.03 0.06; 0.07 0.14;;; 0.18 0.03; 0.42 0.07]),), output = MatrixDirichlet([1.07 1.42; 1.1400000000000001 1.07])
+        )]
+    end
+end