Add Continuous Transition Node for Latent State Transformation

src/ReactiveMP.jl

@@ -123,6 +123,7 @@ include("nodes/bifm.jl")
 include("nodes/bifm_helper.jl")
 include("nodes/probit.jl")
 include("nodes/poisson.jl")
+include("nodes/continuous_transition.jl")
 include("nodes/half_normal.jl")
 
 include("nodes/flow/flow.jl")

src/nodes/continuous_transition.jl

@@ -0,0 +1,105 @@
+export CTransition, ContinuousTransition, CTMeta
+
+import LazyArrays
+import StatsFuns: log2π
+
+@doc raw"""
+The ContinuousTransition node transforms an m-dimensional (dx) vector x into an n-dimensional (dy) vector y via a linear (or nonlinear) transformation with a `n×m`-dimensional matrix `A` that is constructed from a vector `a`.
+
+To construct the matrix A, the elements of `a` are filled into A according to the transformation function provided with meta. `a` must be of MultivariateNormalDistributionsFamily type. If you intend to use univariate Gaussian, use it as a vector of length `1``, e.g. `a ~ MvNormalMeanCovariance([0.0], [1.;])`.
+
+Check CTMeta for more details on how to specify the transformation function that **must** return a matrix.
+
+```julia
+y ~ ContinuousTransition(x, a, W) where {meta = CTMeta(transformation, â)}
+```
+Interfaces:
+1. y - n-dimensional output of the ContinuousTransition node.
+2. x - m-dimensional input of the ContinuousTransition node.
+3. a - any-dimensional vector that casts into the matrix `A`. 
+4. W - `n×n`-dimensional precision matrix used to soften the transition and perform variational message passing.
+
+Note that you can set W to a fixed value or put a prior on it to control the amount of jitter.
+"""
+struct ContinuousTransition end
+
+const CTransition = ContinuousTransition
+
+@node ContinuousTransition Stochastic [y, x, a, W]
+
+@doc raw"""
+`CTMeta` is used as a metadata flag in `ContinuousTransition` to define the transformation function for constructing the matrix `A` from vector `a`.
+
+`CTMeta` requires a transformation function and the length of vector `a`, which acts as an expansion point for approximating the transformation linearly. If transformation appears to be linear, then no approximation is performed.
+
+Constructors:
+- `CTMeta(transformation::Function, â::Vector{<:Real})`: Constructs a `CTMeta` struct with the transformation function and allocated basis vectors.
+
+Fields:
+- `ds`: A tuple indicating the dimensionality of the ContinuousTransition (dy, dx).
+- `f`:  Represents the transformation function that transforms vector `a` into matrix `A`
+- `es`: A Vector of unit vectors used in the transformation process.
+
+The `CTMeta` struct plays a pivotal role in defining how the vector `a` is transformed into the matrix `A`, thus influencing the behavior of the `ContinuousTransition` node.
+"""
+struct CTMeta
+    ds::Tuple # dimensionality of ContinuousTransition (dy, dx)
+    f::Function # transformation function
+    es::Vector{<:AbstractVector} # unit vectors
+
+    # NOTE: this meta is not user-friendly, I don't like a vector
+    # perhaps making mutable struct with empty meta first will be better from user perspective
+    # meta for transformation of a vector to a matrix
+    function CTMeta(transformation::Function, â::Vector{<:Real})
+        dy, dx = size(transformation(â))
+        es = [StandardBasisVector(dy, i, one(eltype(first(â)))) for i in 1:dy]
+        return new((dy, dx), transformation, es)
+    end
+end
+
+getunits(meta::CTMeta)          = meta.es
+getdimensionality(meta::CTMeta) = meta.ds
+gettransformation(meta::CTMeta) = meta.f
+
+getjacobians(ctmeta::CTMeta, a) = process_Fs(gettransformation(ctmeta), a)
+process_Fs(f::Function, a) = [ForwardDiff.jacobian(a -> f(a)[i, :], a) for i in 1:size(f(a), 1)]
+
+default_meta(::Type{CTMeta}) = error("ContinuousTransition node requires meta flag explicitly specified")
+
+default_functional_dependencies_pipeline(::Type{<:ContinuousTransition}) = RequireMarginalFunctionalDependencies((3,), (nothing,))
+
+"""
+    `ctcompanion_matrix` casts a vector `a` into a matrix `A` by means of linearization of the transformation function `f` around the expansion point `a0`.
+"""
+function ctcompanion_matrix(a, epsilon, meta::CTMeta)
+    a0 = a .+ epsilon # expansion point
+    Js, es = getjacobians(meta, a0), getunits(meta)
+    f = gettransformation(meta)
+    dy, _ = getdimensionality(meta)
+    # we approximate each row of A by a linear function and create a matrix A composed of the approximated rows
+    A = sum(es[i] * (f(a0)[i, :] + Js[i] * (a - a0))' for i in 1:dy)
+    return A
+end
+
+@average_energy ContinuousTransition (q_y_x::Any, q_a::Any, q_W::Any, meta::CTMeta) = begin
+    ma, Va   = mean_cov(q_a)
+    myx, Vyx = mean_cov(q_y_x)
+    mW       = mean(q_W)
+
+    dy, dx = getdimensionality(meta)
+    Fs, es = getjacobians(meta, ma), getunits(meta)
+    n      = div(ndims(q_y_x), 2)
+    mA     = ctcompanion_matrix(ma, sqrt.(var(q_a)), meta)
+
+    mx, Vx = @views myx[(dy + 1):end], Vyx[(dy + 1):end, (dy + 1):end]
+    my, Vy = @views myx[1:dy], Vyx[1:dy, 1:dy]
+    Vyx    = @view Vyx[1:dy, (dy + 1):end]
+    g₁     = my' * mW * my + tr(Vy * mW)
+    g₂     = mx' * mA' * mW * my + tr(Vyx * mA' * mW)
+    g₃     = g₂
+    G      = sum(sum(es[i]' * mW * es[j] * Fs[i] * (ma * ma' + Va) * Fs[j]' for i in 1:length(Fs)) for j in 1:length(Fs))
+    g₄     = mx' * G * mx + tr(Vx * G)
+    AE     = n / 2 * log2π - (mean(logdet, q_W) - (g₁ - g₂ - g₃ + g₄)) / 2
+
+    return AE
+end

src/rules/continuous_transition/W.jl

@@ -0,0 +1,27 @@
+function compute_delta(my, Vy, mx, Vx, Vyx, mA, Va, ma, Fs, es)
+    G₁ = (my * my' + Vy)
+    G₂ = ((my * mx' + Vyx) * mA')
+    G₃ = transpose(G₂)
+    Ex_xx = mx * mx' + Vx
+    G₅ = sum(sum(es[i] * ma' * Fs[i]'Ex_xx * Fs[j] * ma * es[j]' for i in 1:length(Fs)) for j in 1:length(Fs))
+    G₆ = sum(sum(es[i] * tr(Fs[i]' * Ex_xx * Fs[j] * Va) * es[j]' for i in 1:length(Fs)) for j in 1:length(Fs))
+    return G₁ - G₂ - G₃ + G₅ + G₆
+end
+
+@rule ContinuousTransition(:W, Marginalisation) (q_y_x::MultivariateNormalDistributionsFamily, q_a::MultivariateNormalDistributionsFamily, meta::CTMeta) = begin
+    dy, dx = getdimensionality(meta)
+
+    ma, Va = mean_cov(q_a)
+    Fs, es = getjacobians(meta, ma), getunits(meta)
+
+    mA = ctcompanion_matrix(ma, sqrt.(var(q_a)), meta)
+    myx, Vyx = mean_cov(q_y_x)
+
+    mx, Vx = @views myx[(dy + 1):end], Vyx[(dy + 1):end, (dy + 1):end]
+    my, Vy = @views myx[1:dy], Vyx[1:dy, 1:dy]
+    Vyx    = @views Vyx[1:dy, (dy + 1):end]
+
+    Δ = compute_delta(my, Vy, mx, Vx, Vyx, mA, Va, ma, Fs, es)
+
+    return WishartFast(dy + 2, Δ)
+end

src/rules/continuous_transition/a.jl

@@ -0,0 +1,19 @@
+@rule ContinuousTransition(:a, Marginalisation) (q_y_x::MultivariateNormalDistributionsFamily, q_a::MultivariateNormalDistributionsFamily, q_W::Any, meta::CTMeta) = begin
+    ma = mean(q_a)
+    mW = mean(q_W)
+    myx, Vyx = mean_cov(q_y_x)
+
+    dy, dx = getdimensionality(meta)
+
+    mx, Vx = @views myx[(dy + 1):end], Vyx[(dy + 1):end, (dy + 1):end]
+    my, Vy = @views myx[1:dy], Vyx[1:dy, 1:dy]
+    Vyx    = @view Vyx[1:dy, (dy + 1):end]
+
+    Fs, es = getjacobians(meta, ma), getunits(meta)
+
+    # rank1update(Vyx, mx, my) equivalent to ξ = (Vyx + mx * my') 
+    D = sum(sum(es[j]' * mW * es[i] * Fs[i]' * rank1update(Vx, mx) * Fs[j] for i in 1:length(Fs)) for j in 1:length(Fs))
+    z = sum(Fs[i]' * rank1update(Vyx', mx, my) * mW * es[i] for i in 1:length(Fs))
+
+    return MvNormalWeightedMeanPrecision(z, D)
+end

src/rules/continuous_transition/marginals.jl

@@ -0,0 +1,36 @@
+
+@marginalrule ContinuousTransition(:y_x) (m_y::MultivariateNormalDistributionsFamily, m_x::MultivariateNormalDistributionsFamily, q_a::Any, q_W::Any, meta::CTMeta) = begin
+    return continuous_tranition_marginal(m_y, m_x, q_a, q_W, meta)
+end
+
+function continuous_tranition_marginal(m_y::MultivariateNormalDistributionsFamily, m_x::MultivariateNormalDistributionsFamily, q_a::Any, q_W::Any, meta::CTMeta)
+    ma, Va = mean_cov(q_a)
+
+    Fs, es = getjacobians(meta, ma), getunits(meta)
+
+    mW = mean(q_W)
+
+    mA = ctcompanion_matrix(ma, sqrt.(var(q_a)), meta)
+
+    b_my, b_Vy = mean_cov(m_y)
+    f_mx, f_Vx = mean_cov(m_x)
+
+    inv_b_Vy = cholinv(b_Vy)
+    inv_f_Vx = cholinv(f_Vx)
+
+    Ξ = inv_f_Vx + sum(sum(es[j]' * mW * es[i] * Fs[j] * Va * Fs[i]' for i in 1:length(Fs)) for j in 1:length(Fs))
+
+    W_11 = inv_b_Vy + mW
+
+    # negate_inplace!(mW * mH)
+    W_12 = -(mW * mA)
+
+    W_21 = -(mA' * mW)
+
+    W_22 = Ξ + mA' * mW * mA
+
+    W = [W_11 W_12; W_21 W_22]
+    ξ = [inv_b_Vy * b_my; inv_f_Vx * f_mx]
+
+    return MvNormalWeightedMeanPrecision(ξ, W)
+end

src/rules/continuous_transition/x.jl

@@ -0,0 +1,19 @@
+@rule ContinuousTransition(:x, Marginalisation) (m_y::MultivariateNormalDistributionsFamily, q_a::MultivariateNormalDistributionsFamily, q_W::Any, meta::CTMeta) = begin
+    ma, Va = mean_cov(q_a)
+    my, Wy = mean_precision(m_y)
+    mW = mean(q_W)
+
+    dy, dx = getdimensionality(meta)
+    Fs, es = getjacobians(meta, ma), getunits(meta)
+
+    mA = ctcompanion_matrix(ma, sqrt.(var(q_a)), meta)
+
+    W = sum(sum(es[j]' * mW * es[i] * Fs[j] * Va * Fs[i]' for i in 1:length(Fs)) for j in 1:length(Fs))
+    # Woodbury identity
+    # inv(inv(Wy) + inv(mW)) = Wy - Wy * inv(Wy + mW) * Wy
+    WymW = Wy - Wy * cholinv(Wy + mW) * Wy
+    z = mA' * WymW * my
+    Ξ = mA' * WymW * mA + W
+
+    return MvNormalWeightedMeanPrecision(z, Ξ)
+end

src/rules/continuous_transition/y.jl

@@ -0,0 +1,16 @@
+@rule ContinuousTransition(:y, Marginalisation) (m_x::MultivariateNormalDistributionsFamily, q_a::MultivariateNormalDistributionsFamily, q_W::Any, meta::CTMeta) = begin
+    ma = mean(q_a)
+    mx, Vx = mean_cov(m_x)
+
+    mW = mean(q_W)
+
+    dy, dx = getdimensionality(meta)
+    Fs, es = getjacobians(meta, ma), getunits(meta)
+
+    mA = ctcompanion_matrix(ma, sqrt.(var(q_a)), meta)
+
+    Vy = mA * Vx * mA' + cholinv(mW)
+    my = mA * mx
+
+    return MvNormalMeanCovariance(my, Vy)
+end

src/rules/prototypes.jl

@@ -110,6 +110,12 @@ include("transition/out.jl")
 include("transition/in.jl")
 include("transition/a.jl")
 
+include("continuous_transition/y.jl")
+include("continuous_transition/x.jl")
+include("continuous_transition/a.jl")
+include("continuous_transition/W.jl")
+include("continuous_transition/marginals.jl")
+
 include("autoregressive/y.jl")
 include("autoregressive/x.jl")
 include("autoregressive/theta.jl")

test/nodes/test_continuous_transition.jl

@@ -0,0 +1,47 @@
+module ContinuousTransitionNodeTest
+
+using Test, ReactiveMP, Random, Distributions, BayesBase, ExponentialFamily
+
+import ReactiveMP: getdimensionality, getjacobians, gettransformation, getunits, ctcompanion_matrix
+
+@testset "ContinuousTransitionNode" begin
+    dy, dx = 2, 3
+    a0 = rand(dx * dy)  # Example vector `a0`
+    meta = CTMeta(a -> reshape(a, dy, dx), a0)
+    @testset "Creation" begin
+        node = make_node(ContinuousTransition, FactorNodeCreationOptions(nothing, meta, nothing))
+
+        @test functionalform(node) === ContinuousTransition
+        @test sdtype(node) === Stochastic()
+        @test name.(interfaces(node)) === (:y, :x, :a, :W)
+        @test factorisation(node) === ((1, 2, 3, 4),)
+        @test getdimensionality(metadata(node)) == (dy, dx)  # Based on the transformation function dimensions
+    end
+
+    @testset "AverageEnergy" begin
+        # This is an example setup, you'll need to adjust the distributions and marginals according to your needs
+        q_y_x = MvNormalMeanCovariance(zeros(5), diageye(5))
+        q_a = MvNormalMeanCovariance(zeros(6), diageye(6))  # Adjust the dimension according to `a`
+        q_W = Wishart(3, diageye(2))
+
+        marginals = (Marginal(q_y_x, false, false, nothing), Marginal(q_a, false, false, nothing), Marginal(q_W, false, false, nothing))
+
+        @test score(AverageEnergy(), ContinuousTransition, Val{(:y_x, :a, :W)}(), marginals, meta) ≈ 13.415092731310878
+        @show getjacobians(meta, a0)
+    end
+
+    @testset "CTransition Functionality" begin
+        A = ctcompanion_matrix(a0, zeros(length(a0)), meta)
+
+        @test size(A) == (dy, dx)
+        @test A == gettransformation(meta)(a0)
+    end
+
+    @testset "Metadata Functionality" begin
+        @test getdimensionality(meta) == (dy, dx)
+        @test length(getjacobians(meta, a0)) == dy  # Based on `dy`
+        @test length(getunits(meta)) == dy  # Based on `dy`
+    end
+end
+
+end

test/nodes/test_transfominator.jl

@@ -0,0 +1,22 @@
+score(
+    AverageEnergy(),
+    ContinuousTransition,
+    Val{(:y_x, :h, :Λ)}(),
+    (
+        Marginal(MvNormalMeanPrecision(zeros(4), diageye(4)), false, false, nothing),
+        Marginal(MvNormalMeanPrecision(zeros(4), diageye(4)), false, false, nothing),
+        Marginal(Wishart(2, diageye(2)), false, false, nothing)
+    ),
+    CTMeta((2, 2))
+)
+score(
+    AverageEnergy(),
+    ContinuousTransition,
+    Val{(:y_x, :h, :Λ)}(),
+    (
+        Marginal(MvNormalMeanPrecision(zeros(5), diageye(5)), false, false, nothing),
+        Marginal(MvNormalMeanPrecision(zeros(6), diageye(6)), false, false, nothing),
+        Marginal(Wishart(2, diageye(2)), false, false, nothing)
+    ),
+    CTMeta((2, 3))
+)

test/rules/continuous_transition/test_W.jl

@@ -0,0 +1,70 @@
+module RulesContinuousTransitionTest
+
+using Test, ReactiveMP, BayesBase, Random, ExponentialFamily, Distributions, LinearAlgebra
+
+import ReactiveMP: @test_rules, ctcompanion_matrix, getjacobians, getunits, WishartFast
+
+@testset "rules:ContinuousTransition:W" begin
+    rng = MersenneTwister(42)
+
+    @testset "Linear transformation" begin
+        # the following rule is used for testing purposes only
+        # It is derived separately by Thijs van de Laar
+        function benchmark_rule(q_y_x, mA, ΣA, UA)
+            myx, Vyx = mean_cov(q_y_x)
+
+            dy  = size(mA, 1)
+            Vx  = Vyx[(dy + 1):end, (dy + 1):end]
+            Vy  = Vyx[1:dy, 1:dy]
+            mx  = myx[(dy + 1):end]
+            my  = myx[1:dy]
+            Vyx = Vyx[1:dy, (dy + 1):end]
+
+            G = tr(Vx * UA) * ΣA + mA * Vx * mA' - mA * Vyx' - Vyx * mA' + Vy + ΣA * mx' * UA * mx + (mA * mx - my) * (mA * mx - my)'
+
+            return WishartFast(dy + 2, G)
+        end
+
+        @testset "Structured: (q_y_x::MultivariateNormalDistributionsFamily, q_a::MultivariateNormalDistributionsFamily, meta::CTMeta)" begin
+            for (dy, dx) in [(1, 3), (2, 3), (3, 2), (2, 2)]
+                dydx = dy * dx
+                transformation = (a) -> reshape(a, dy, dx)
+                mA, ΣA, UA = rand(rng, dy, dx), diageye(dy), diageye(dx)
+
+                a0 = Float32.(vec(mA))
+
+                metal = CTMeta(transformation, a0)
+                Lx, Ly = rand(rng, dx, dx), rand(rng, dy, dy)
+                μx, Σx = rand(rng, dx), Lx * Lx'
+                μy, Σy = rand(rng, dy), Ly * Ly'
+
+                qyx = MvNormalMeanCovariance([μy; μx], [Σy zeros(dy, dx); zeros(dx, dy) Σx])
+                qa = MvNormalMeanCovariance(vec(mA), kron(UA, ΣA))
+
+                @test_rules [check_type_promotion = true, atol = 1e-5] ContinuousTransition(:W, Marginalisation) [(
+                    input = (q_y_x = qyx, q_a = qa, meta = metal), output = benchmark_rule(qyx, mA, ΣA, UA)
+                )]
+            end
+        end
+    end
+
+    @testset "Nonlinear transformation" begin
+        @testset "Structured: (q_y_x::MultivariateNormalDistributionsFamily, q_a::Any, q_W::Any, meta::CTMeta)" begin
+            dy, dx = 2, 2
+            dydx = dy * dy
+            transformation = (a) -> [cos(a[1]) -sin(a[1]); sin(a[1]) cos(a[1])]
+            a0 = zeros(Int, 1)
+            metanl = CTMeta(transformation, a0)
+            μx, Σx = zeros(dx), diageye(dx)
+            μy, Σy = zeros(dy), diageye(dy)
+
+            qyx = MvNormalMeanCovariance([μy; μx], [Σy zeros(dy, dx); zeros(dx, dy) Σx])
+            qa = MvNormalMeanCovariance(a0, diageye(1))
+            @test_rules [check_type_promotion = true] ContinuousTransition(:W, Marginalisation) [(
+                input = (q_y_x = qyx, q_a = qa, meta = metanl), output = WishartFast(dy + 2, dy * diageye(dy))
+            )]
+        end
+    end
+end
+
+end