Half-edge error in switching model

[ValentinKaisermayer]

I have a measurement signal that measures either quantity A or quantity B. An additional signal tells me which one it is. However, there is a dynamic in the model, so when switching it takes some time.

So my model is two random walks for the original quantities:

$$x_{A,k} \sim \mathcal{N}(x_{A,k-1},1)$$ $$x_{B,k} \sim \mathcal{N}(x_{B,k-1},1)$$

The model for combining the two:

$$x_k \sim \mathcal{N}(a*x_{k-1} + b*(\delta_{k-1}*x_{A,k-1} + (1-\delta_{k-1})*x_{B,k-1}),1)$$

And finally, my noisy observations:

$$y_k \sim \mathcal{N}(x_{k},P)$$

However, I have problems implementing this in RxInfer:

using CairoMakie
using RxInfer

xA = cumsum(vcat(2, randn(99)))
xB = cumsum(vcat(-10, randn(99)))
is_at_A = rand(Bool, 100)
x = is_at_A .* xA .+ (1 .- is_at_A) .* xB
y = x .+ randn(100)

let
    fig = Figure()
    ax1 = Axis(fig[1, 1])
    # lines!(x; label="x")
    scatter!(y; label="Observations", color=:black)
    lines!(xA; label="Hidden Signal A")
    lines!(xB; label="Hidden Signal B")
    axislegend(ax1)
    fig
end

Note that my simulated data is without the dynamics, but that should not matter.

@model function split(x0, x0_A, x0_B, y, is_at_A, A, B)
    P ~ Gamma(; shape=0.001, scale=0.001)
    x_prior ~ Normal(; mean=mean(x0), var=var(x0))
    x_A_prior ~ Normal(; mean=mean(x0_A), var=var(x0_A))
    x_B_prior ~ Normal(; mean=mean(x0_B), var=var(x0_B))

    local x, x_A, x_B
    x_prev = x_prior
    x_A_prev = x_A_prior
    x_B_prev = x_B_prior

    for i in 1:length(y)
        x_A[i] ~ Normal(; mean=x_A_prev, precision=10) # random walk
        x_B[i] ~ Normal(; mean=x_B_prev, precision=10) # random walk

        x[i] ~ Normal(; mean=A * x_prev + B * (is_at_A[i] * x_A_prev + (1 - is_at_A[i]) * x_B_prev), precision=10)
        y[i] ~ Normal(; mean=x[i], precision=P)

        x_prev = x[i]
        x_A_prev = x_A[i]
        x_B_prev = x_B[i]
    end
end

A = 0.9
B = 0.1

constraints = @constraints begin
    q(x_prior, x, x_A_prior, x_A, x_B_prior, x_B, y, P) = q(x_prior, x, x_A_prior, x_A, x_B_prior, x_B)q(P)q(y)
end

x0_prior = NormalMeanVariance(0.0, 1000.0)
x0_A_prior = NormalMeanVariance(0.0, 1000.0)
x0_B_prior = NormalMeanVariance(0.0, 1000.0)

initm = @initialization begin
    q(P) = Gamma(0.001, 0.001)
end

result = infer(;
    model=split(; x0=x0_prior, x0_A=x0_A_prior, x0_B=x0_B_prior, is_at_A=is_at_A, A=A, B=B),
    data=(y=y,),
    constraints=constraints,
    initialization=initm,
    returnvars=(x=KeepLast(), x_A=KeepLast(), x_B=KeepLast()),
    iterations=20,
    options=(limit_stack_depth=100,),
)

I'm getting an error:

ERROR: Half-edge has been found: x_A_2690. To terminate half-edges 'Uninformative' node can be used.

Thx in advance!

[ValentinKaisermayer]

I also tried

if is_at_A[i]
    x[i] ~ Normal(; mean=A * x_prev + B * x_A_prev, precision=10)
else
    x[i] ~ Normal(; mean=A * x_prev + B * x_B_prev, precision=10)
end

which results in the same error.

[ValentinKaisermayer]

I updated my model - still the same error - but I think nicer to read:

@model function split(x0, x0_A, x0_B, y, is_at_A, A, B)
    P ~ Gamma(; shape=0.001, scale=0.001)

    x_A[1] ~ x0_A
    x_B[1] ~ x0_B
    for i in 2:length(y)
        x_A[i] ~ Normal(; mean=x_A[i - 1], precision=10) # random walk
        x_B[i] ~ Normal(; mean=x_B[i - 1], precision=10) # random walk
    end

    x[1] ~ x0
    for i in 2:length(y)
        if is_at_A[i - 1]
            x[i] ~ Normal(; mean=A * x[i - 1] + B * x_A[i - 1], precision=10)
        else
            x[i] ~ Normal(; mean=A * x[i - 1] + B * x_B[i - 1], precision=10)
        end
        y[i] ~ Normal(; mean=x[i], precision=P)
    end
end

A = 0.9
B = 0.1

constraints = @constraints begin
    q(x, x_A, x_B, y, P) = q(x, x_A, x_B)q(P)q(y)
end

x0_prior = NormalMeanVariance(0.0, 1000.0)
x0_A_prior = NormalMeanVariance(0.0, 1000.0)
x0_B_prior = NormalMeanVariance(0.0, 1000.0)

initialization = @initialization begin
    q(P) = Gamma(0.001, 0.001)
end

result = infer(;
    model=split(; x0=x0_prior, x0_A=x0_A_prior, x0_B=x0_B_prior, is_at_A=is_at_A, A=A, B=B),
    data=(y=y,),
    constraints=constraints,
    initialization=initialization,
    returnvars=(x=KeepLast(), x_A=KeepLast(), x_B=KeepLast()),
    iterations=20,
    options=(limit_stack_depth=100,),
)

[ValentinKaisermayer]

Ok, I now have something that diverges:

@model function split(xA0, xB0, x0, y, is_at_A)
    P ~ Gamma(; shape=0.001, scale=0.001)
    xA_prior ~ Normal(; mean=mean(xA0), var=var(xA0))
    xB_prior ~ Normal(; mean=mean(xB0), var=var(xB0))
    x_prior ~ Normal(; mean=mean(x0), var=var(x0))

    local xA, xB
    xA_prev = xA_prior
    xB_prev = xB_prior

    for i in 1:length(y)
        xA[i] ~ Normal(; mean=xA_prev, precision=20)
        xB[i] ~ Normal(; mean=xB_prev, precision=20)
        xA_prev = xA[i]
        xB_prev = xB[i]
    end

    local x
    x_prev = x_prior

    for i in 1:length(y)
        x[i] := 0.999 * x_prev + 0.001 * (is_at_A[i] * xA[i] + (1 - is_at_A[i]) * xB[i]) # TODO: indices are wrong
        y[i] ~ Normal(; mean=x[i], precision=P)
        x_prev = x[i]
    end
end
constraints = @constraints begin
    q(xA_prior, xA, xB_prior, xB, y, x, x_prior, P) = q(xA_prior, xA, xB_prior, xB, x, x_prior)q(P)q(y)
end
xA0_prior = NormalMeanVariance(0.0, 10.0)
xB0_prior = NormalMeanVariance(0.0, 10.0)
x0_prior = NormalMeanVariance(0.0, 10.0)
initm = @initialization begin
    q(P) = Gamma(0.001, 0.001)
    # Initialize the marginal
    q(xA) = vague(NormalMeanPrecision)
    q(xB) = vague(NormalMeanPrecision)
    q(x) = vague(NormalMeanPrecision)
    # Initialize the message
    μ(xA) = vague(NormalMeanPrecision)
    μ(xB) = vague(NormalMeanPrecision)
    μ(x) = vague(NormalMeanPrecision)
end
result = infer(;
    model=split(; xA0=xA0_prior, xB0=xB0_prior, x0=x0_prior),
    data=(y=y, is_at_A=is_at_A),
    constraints=constraints,
    initialization=initm,
    returnvars=(xA=KeepLast(), xB=KeepLast()),
    iterations=10,
    options=(limit_stack_depth=250,),
    # free_energy=true, # ERROR:     Failed to compute variable bound free energy component for `name = anonymous_var_graphppl, index = nothing, linked to (NodeData in context  with properties name = is_at_A, index = 1 with extra:
)

I would appreciate any advice.
At least until now, I would have implemented this easily as a Kalman filter.

[albertpod]

Hi @ValentinKaisermayer.

I will take a look at your model this week.

Cheers

[albertpod]

@ValentinKaisermayer I've had a lazy attempt to fix it. For starter, you can simplify the constraints, we only need to specify factors that needed to be factorised. I thought going with delta node would work just fine, but I am getting NaN errors happening around inference over Gamma. Seems there is some problem with differentiation of your function. Here's the solution I tried (no success).

using RxInfer
using ExponentialFamilyProjection

generated_xA = cumsum(vcat(2, randn(99)))
generated_xB = cumsum(vcat(-10, randn(99)))
generated_is_at_A = float.(rand(Bool, 100))
generated_x = generated_is_at_A .* generated_xA .+ (1 .- generated_is_at_A) .* generated_xB
generated_y = generated_x .+ randn(100)

function compute_term(is_at_A, x_A, x_B)
    return is_at_A * x_A + (1 - is_at_A) * x_B
end

@model function split(y, is_at_A, x0, x0_A, x0_B, A, B)
    P ~ Gamma(shape=1.0, rate=1.0)
    x_prior ~ Normal(mean=mean(x0), var=var(x0))
    x_A_prior ~ Normal(mean=mean(x0_A), var=var(x0_A))
    x_B_prior ~ Normal(mean=mean(x0_B), var=var(x0_B))

    x_prev = x_prior
    x_A_prev = x_A_prior
    x_B_prev = x_B_prior

    for i in 1:length(y)
        x_A[i] ~ Normal(mean=x_A_prev, precision=1.0)
        x_B[i] ~ Normal(mean=x_B_prev, precision=1.0)
        term[i] := compute_term(is_at_A[i], x_A[i], x_B[i])
        x[i] ~ Normal(mean=A * x_prev + B*term[i], precision=1.0)
        y[i] ~ Normal(mean=x[i], precision=P)

        x_prev = x[i]
        x_A_prev = x_A[i]
        x_B_prev = x_B[i]
    end
end

constraints = @constraints begin
    q(x, P) = q(x)q(P)
end

xA0_prior = NormalMeanVariance(0.0, 1.0)
xB0_prior = NormalMeanVariance(0.0, 1.0)
x0_prior = NormalMeanVariance(0.0, 1.0)

initm = @initialization begin
    q(P) = GammaShapeRate(1.0, 1.0)
    # Initialize the message
    μ(x_A) = NormalMeanPrecision(0.0, 1.0)
    μ(x_B) = NormalMeanPrecision(0.0, 1.0)
end

meta = @meta begin 
    compute_term() -> Linearization()
end;

result = infer(
    model = split(x0=x0_prior, x0_A=xA0_prior, x0_B=xB0_prior, A=1.0, B=1.0),
    data = (y =generated_y, is_at_A = generated_is_at_A),
    constraints = constraints,
    initialization = initm,
    meta = meta,
    iterations = 10,
    options = (limit_stack_depth = 250,),
)

Given that you are interested in x_A, x_B then the following trick my work for you, although I am afraid that it won't scale super good.
https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/#inference-undefinedrules-fusedelta

[ValentinKaisermayer]

Thank you for trying. I have not tested the model, however I noticed that using A=1 and B=1 cannot work.

They stem from a discretization of

$$ \tau \dot{x} = -x + u $$

$$ x_{k+1} = x_k + T_d \frac{1}{\tau}(-x_k + u_k) $$ where $T_d$ is the sample time. $$ x_{k+1} = (1- T_d \frac{1}{\tau})x_k + T_d \frac{1}{\tau}u_k $$

Where $(1- T_d \frac{1}{\tau})<1$ for stability. As you can see $A=1-B$.

The model without the dynamic works fine, it is basically the System identification example.

[albertpod]

Indeed, the correct model specification should give you better estimates during the inference stage; however, RxInfer.jl has the tooling to run inference regardless of model misspecification. I've played a bit more with the projections mechanism. This is far too advanced in my opinion, so there's no urgency to understand it (we are working on automating this, so you won't need to bother with constraint and projection specifications). Overall, the way to run the inference is:

using RxInfer
using ExponentialFamilyProjection
using BayesBase
using Plots

generated_xA = cumsum(vcat(2, randn(99)))
generated_xB = cumsum(vcat(-10, randn(99)))
generated_is_at_A = float.(rand(Bool, 100))
generated_x = generated_is_at_A .* generated_xA .+ (1 .- generated_is_at_A) .* generated_xB
generated_y = generated_x .+ randn(100)

function compute_term(is_at_A, x_A, x_B)
    return is_at_A * x_A + (1 - is_at_A) * x_B
end

struct TransformedNormalDistribution{X, H, T, P, A, B} <: ContinuousUnivariateDistribution
    x::X
    h::H
    t::T
    p::P
    A::A
    B::B
end

BayesBase.logpdf(dist::TransformedNormalDistribution, x) = logpdf(NormalMeanPrecision(dist.A*dist.x + dist.B*compute_term(dist.h, dist.t, dist.p), 0.1), x)
BayesBase.insupport(dist::TransformedNormalDistribution, x) = true

@node TransformedNormalDistribution Stochastic [out, x, h, t, p, A, B]

@model function split(y, is_at_A, x0, x0_A, x0_B, A, B)
    P ~ Gamma(shape=0.01, rate=1.0)
    x_prior ~ Normal(mean=mean(x0), var=var(x0))
    
    x_A_prior ~ Normal(mean=mean(x0_A), var=var(x0_A))
    x_B_prior ~ Normal(mean=mean(x0_B), var=var(x0_B))

    x_prev = x_prior
    x_A_prev = x_A_prior
    x_B_prev = x_B_prior

    for i in 1:length(y)
        x_A[i] ~ Normal(mean=x_A_prev, var=1.0)
        x_B[i] ~ Normal(mean=x_B_prev, var=1.0)
        x[i]   ~ TransformedNormalDistribution(x_prev, is_at_A[i], x_A[i], x_B[i], A, B)
        y[i]   ~ Normal(mean=x[i], precision=P)

        x_prev = x[i]
        x_A_prev = x_A[i]
        x_B_prev = x_B[i]
    end
end

constraints = @constraints begin
    q(x_prior, x_A, x_B, x, P, x_A_prior, x_B_prior) = q(x_prior)q(x_A)q(x_B)q(x)q(P)q(x_A_prior)q(x_B_prior)
    q(x) = q(x[begin])..q(x[end])
    q(x_A) = q(x_A[begin])..q(x_A[end])
    q(x_B) = q(x_B[begin])..q(x_B[end])
    q(x_A) :: ProjectedTo(NormalMeanVariance)
    q(x_B) :: ProjectedTo(NormalMeanVariance)
    q(x_prior) :: ProjectedTo(NormalMeanVariance)
    q(x) :: ProjectedTo(NormalMeanVariance)
end

xA0_prior = NormalMeanVariance(0.0, 1.0)
xB0_prior = NormalMeanVariance(0.0, 1.0)
x0_prior = NormalMeanVariance(0.0, 1.0)

initm = @initialization begin
    q(P) = GammaShapeRate(0.01, 1.0)
    q(x_A) = NormalMeanVariance(0.0, 1e2)
    q(x_B) = NormalMeanVariance(0.0, 1e2)
    q(x) = NormalMeanVariance(0.0, 1e2)
    q(x_prior) = NormalMeanVariance(0.0, 1e2)
end

result = infer(
    model = split(x0=x0_prior, x0_A=xA0_prior, x0_B=xB0_prior, A=1.0, B=0.0),
    data = (y =generated_y, is_at_A = generated_is_at_A),
    constraints = constraints,
    initialization = initm,
    iterations = 100,
    showprogress = true,
    options = (
        rulefallback = NodeFunctionRuleFallback(),
        limit_stack_depth = 250,
    )
)


scatter(generated_y, label="data")
plot!(generated_x, label="truth")
plot!(mean.(result.posteriors[:x][end]), ribbon=std.(result.posteriors[:x][end]), label="inferred")


plot(generated_xA, label="truth")
plot!(mean.(result.posteriors[:x_A][end]), ribbon=std.(result.posteriors[:x_A][end]), label="inferred")

plot(generated_xB, label="truth")
plot!(mean.(result.posteriors[:x_B][end]), ribbon=std.(result.posteriors[:x_B][end]), label="inferred")

The recovery of x_A and x_B is very poor.

I think if you're dealing with switching dynamics in your signal, it's better to model the signal via switching dynamic models. Would it be possible to assume that your signal is coming from a mixture model, where means are represented by separate random walks?