[Bug]: Cannot use ProjectedTo or CVIProjection for constrained Wishart transformation

[albertpod]

Description

Following the discussion in #565, I'm extending a structural time series (STS) model, however in that discussion I've provided unconstrained Wishart prior on the full transition precision (which seeks correlations between all components), I need to follow the original STS approach and constrain it as R*Q*R' where Q ~ Wishart(df, S) and R is a known constraint matrix.

The model structure is:

z[t] ~ ContinuousTransition(zprev, F, QR)  # where QR = R*Q*R'

Before implementing custom rules, I wanted to experiment using existing functionality (ProjectedTo or CVIProjection), but both approaches fail with different errors.

Expected Behavior

Either ProjectedTo(Wishart) or CVIProjection() should allow me to experiment with the constrained Wishart transformation without having to immediately implement custom inference rules.

Actual Behavior

Throws errors.

Minimal Reproducible Example or steps to reproduce the issue

Approach 1: Custom node with ProjectedTo

Following the node fusion example, I created a custom TransformedWishart node:

using RxInfer

struct TransformedWishart{W, R} <: ContinuousMatrixDistribution
    w::W
    r::R
end

BayesBase.logpdf(dist::TransformedWishart, x) = logpdf(Wishart(size(dist.R, 1) + dist.W.df, dist.R*dist.W.S*transpose(dist.R)), x)
BayesBase.insupport(dist::TransformedWishart, x) = true

@node TransformedWishart Stochastic [out, w, r]

@model function rxsts(y, R)
    Q    ~ Wishart(2, diageye(2))
    zprev ~ MvNormalMeanPrecision(zeros(2), diageye(2))
    QR    ~ TransformedWishart(Q, R)
    F     ~ MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    for t in eachindex(y)
        z[t] ~ ContinuousTransition(zprev, F, QR)
        y[t] ~ Normal(mean = z[t], precision = 1.0)
        zprev = z[t]
    end
end

@constraints function rxsts_constraints()
    q(z, zprev, F, QR, Q) = q(z, zprev)q(F)q(QR)q(Q)
    q(QR) :: ProjectedTo(Wishart)
    q(Q)  :: ProjectedTo(Wishart)
end

@meta function rxsts_meta()
    ContinuousTransition() -> CTMeta(transition)
end

@initialization function rxsts_init()
    q(F)     = MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    q(Q)     = Wishart(2, diageye(2))
    q(QR)    = Wishart(4, diageye(4))
    μ(zprev) = MvNormalMeanPrecision(zeros(2), diageye(2))
end

R = [
    0.1   0.0;
    0.0   0.0;
    0.0   1.0;
    0.0   0.0;
]

results = infer(
    model = rxsts(R=R),
    data = (y = [randn(4) for _ in 1:100],),
    constraints = rxsts_constraints(),
    meta = rxsts_meta(),
    initialization = rxsts_init(),
    returnvars = KeepLast(),
    iterations = 50,
    showprogress = true,
    options = (
        rulefallback = NodeFunctionRuleFallback(),
        limit_stack_depth = 250,
    )
)

---

Approach 2: Deterministic transformation with CVIProjection

Using a deterministic transformation node:

using RxInfer

TransformWishart(Q, R) = R*Q*R'

transition(F) = reshape(F, 4, 4)

@model function rxsts(y, R)
    Q     ~ Wishart(2, diageye(2))
    zprev ~ MvNormalMeanPrecision(zeros(2), diageye(2))
    F     ~ MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    
    QR := TransformWishart(Q, R)
    
    for t in eachindex(y)
        z[t] ~ ContinuousTransition(zprev, F, QR)
        y[t] ~ Normal(mean = z[t], precision = 1.0)
        zprev = z[t]
    end
end

@constraints function rxsts_constraints()
    q(z, zprev, F, QR, Q) = q(z, zprev)q(F)q(QR)q(Q)
end

@meta function rxsts_meta()
    ContinuousTransition() -> CTMeta(transition)
    TransformWishart() -> CVIProjection()
end

@initialization function rxsts_init()
    q(F)     = MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    q(Q)     = Wishart(2, diageye(2))
    q(QR)    = Wishart(4, diageye(4))
    μ(zprev) = MvNormalMeanPrecision(zeros(2), diageye(2))
end

R = [
    0.1   0.0;
    0.0   0.0;
    0.0   1.0;
    0.0   0.0;
]

results = infer(
    model = rxsts(R = R),
    data = (y = [randn(4) for _ in 1:100],),
    constraints = rxsts_constraints(),
    meta = rxsts_meta(),
    initialization = rxsts_init(),
    returnvars = KeepLast(),
    iterations = 50,
    showprogress = true,
    options = (
        limit_stack_depth = 250,
    )
)

Error Message / Stack Trace

**Approach 1: Custom node with ProjectedTo**

**Error:**

ERROR: MethodError: no method matching get_natural_manifold_base(::Type{Wishart}, ::Tuple{}, ::Nothing)
The function `get_natural_manifold_base` exists, but no method is defined for this combination of argument types.


**Approach 2: Deterministic transformation with CVIProjection**

**Error:**

ERROR: RuleMethodError: no method matching rule for the given arguments
Possible fix, define:
@rule TransformedWishart(:w, Marginalisation) (q_out::Wishart, q_r::PointMass, ) = begin
    return ...
end

Julia Version

1.11

RxInfer Version

Latest stable

Environment Information

(ppplayground) pkg> status ExponentialFamilyManifolds
Status `~/Julia/blabla/Project.toml`
⌃ [5c9727c4] ExponentialFamilyManifolds v3.0.2

Session ID (Optional)

No response

Additional Context

No response

[Nimrais]

However I agree, that user interface should be improved it's not a bug per-se but a feature that the constraint language can not infer your dimensions automatically from the Wishart. I modified your code to make inference start, however you have a bug in the logpdf, please implement proreply.

using RxInfer

using ExponentialFamilyProjection
using ExponentialFamily
using BayesBase

function BayesBase.insupport(::ExponentialFamily.WishartFast, ::Any)
    return true
end

struct TransformedWishart{W, R} <: ContinuousMatrixDistribution
    w::W
    r::R
end

transition(F) = reshape(F, 4, 4)

BayesBase.logpdf(dist::TransformedWishart, x) = logpdf(Wishart(size(dist.r, 1) + dist.w.df, dist.r*dist.w.S*transpose(dist.r)), x)
BayesBase.logpdf(dist::TransformedWishart, x::Matrix) = logpdf(Wishart(size(dist.r, 1) + dist.w.df, dist.R*dist.w.S*transpose(dist.r)), x)
BayesBase.insupport(dist::TransformedWishart, x) = true

@node TransformedWishart Stochastic [out, w, r]

@model function rxsts(y, R)
    Q    ~ Wishart(2, diageye(2))
    zprev ~ MvNormalMeanPrecision(zeros(2), diageye(2))
    QR    ~ TransformedWishart(Q, R)
    F     ~ MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    for t in eachindex(y)
        z[t] ~ ContinuousTransition(zprev, F, QR)
        y[t] ~ Normal(mean = z[t], precision = 1.0)
        zprev = z[t]
    end
end

@constraints function rxsts_constraints()
    q(z, zprev, F, QR, Q) = q(z, zprev)q(F)q(QR)q(Q)
    q(QR) :: ProjectedTo(ExponentialFamily.WishartFast, 4, 4)
    q(Q)  :: ProjectedTo(ExponentialFamily.WishartFast, 2, 2)
end

@meta function rxsts_meta()
    ContinuousTransition() -> CTMeta(transition)
end

@initialization function rxsts_init()
    q(F)     = MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    q(Q)     = Wishart(2, diageye(2))
    q(QR)    = Wishart(4, diageye(4))
    μ(zprev) = MvNormalMeanPrecision(zeros(2), diageye(2))
end

R = [
    0.1   0.0;
    0.0   0.0;
    0.0   1.0;
    0.0   0.0;
]

results = infer(
    model = rxsts(R=R),
    data = (y = [randn(4) for _ in 1:100],),
    constraints = rxsts_constraints(),
    meta = rxsts_meta(),
    initialization = rxsts_init(),
    returnvars = KeepLast(),
    iterations = 50,
    showprogress = true,
    options = (
        rulefallback = NodeFunctionRuleFallback(),
        limit_stack_depth = 250,
    )
)

[albertpod]

Thanks, @Nimrais. We are talking about ProjectedTo approach. I was a bit confused how logpdf will treat TransformedWishart . There were some issues with the model as well. Fixed the model and fixed logpdf, but it's still likely incorrect (I will double check later on). Now I get another error and I can't get what does that mean and what should I do about it?

using RxInfer
using ExponentialFamilyProjection
using ExponentialFamily
using BayesBase

function BayesBase.insupport(::ExponentialFamily.WishartFast, ::Any)
    return true
end

struct TransformedWishart{W, R} <: ContinuousMatrixDistribution
    w::W
    r::R
end

transition(F) = reshape(F, 4, 4)

BayesBase.logpdf(dist::TransformedWishart, x::Matrix) = begin
    return logpdf(ExponentialFamily.WishartFast(size(dist.r, 1) + size(dist.w, 1), dist.r*dist.w*transpose(dist.r)), x)
end

BayesBase.insupport(dist::TransformedWishart, x) = true

@node TransformedWishart Stochastic [out, w, r]

@model function rxsts(y, R)
    Q    ~ Wishart(2, diageye(2))
    zprev ~ MvNormalMeanPrecision(zeros(4), diageye(4))
    QR    ~ TransformedWishart(Q, R)
    F     ~ MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    for t in eachindex(y)
        z[t] ~ ContinuousTransition(zprev, F, QR)
        y[t] ~ MvNormalMeanPrecision(z[t], diageye(4))
        zprev = z[t]
    end
end

@constraints function rxsts_constraints()
    q(z, zprev, F, QR, Q) = q(z, zprev)q(F)q(QR)q(Q)
    q(QR) :: ProjectedTo(ExponentialFamily.WishartFast, 4, 4)
    q(Q)  :: ProjectedTo(ExponentialFamily.WishartFast, 2, 2)
end

@meta function rxsts_meta()
    ContinuousTransition() -> CTMeta(transition)
end

@initialization function rxsts_init()
    q(F)     = MvNormalMeanPrecision(vec(diageye(4)), diageye(16))
    q(Q)     = Wishart(2, diageye(2))
    q(QR)    = Wishart(4, diageye(4))
    μ(zprev) = MvNormalMeanPrecision(zeros(4), diageye(4))
end

R = [
    0.1   0.0;
    0.0   0.0;
    0.0   1.0;
    0.0   0.0;
]

results = infer(
    model = rxsts(R=R),
    data = (y = [randn(4) for _ in 1:100],),
    constraints = rxsts_constraints(),
    meta = rxsts_meta(),
    initialization = rxsts_init(),
    returnvars = KeepLast(),
    iterations = 50,
    showprogress = true,
    options = (
        rulefallback = NodeFunctionRuleFallback(),
        limit_stack_depth = 250,
    )
)

ERROR: The initial point must be on the manifold `ExponentialFamilyManifolds.NaturalParametersManifold{ℝ, ExponentialFamily.WishartFast, Tuple{Int64, Int64}, ManifoldsBase.ProductManifold{ℝ, Tuple{ManifoldsBase.PowerManifold{ℝ, Manifolds.PositiveNumbers, ManifoldsBase.TypeParameter{Tuple{1}}, Manifolds.ArrayPowerRepresentation}, Manifolds.SymmetricPositiveDefinite{ManifoldsBase.TypeParameter{Tuple{2}}}}}, Nothing}((2, 2), ProductManifold(PositiveVectors(1), SymmetricPositiveDefinite(2)), nothing)`, got `RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}}}

[Nimrais]

This is actually looks like a bug, I would check tomorrow on the RxInfer session.

[Nimrais]

I was able to figure out where is the error is. But was not able to fix it yet. When we are starting the optimization in the EFP here we are checking that starting point represents correct natural parameters. However first initial sample is not passing the constraints for the Wishart Distribution.

Code reference https://github.com/ReactiveBayes/ExponentialFamilyProjection.jl/blob/main/src/projected_to.jl#L208

It seems somewhere check for the positive definite matrix get in the way.