Extending HMC model with custom node at the observation side

[apiedrafTNO]

Hi all,
I'm trying to make a model that extends a batch of hidden markov chains with a custom node after the observation transition but before inserting the observation.

The original model consists of N continguous markov chains of length M. The code looks like this:

@model function batch(M, N)
    h_0 = randomvar(N) # Initial state per location
    h   = randomvar(M, N) # Hidden states per location over time
    w   = datavar(Vector{Float64}, M, N) # Observed states per location over time

    T ~ MatrixDirichlet(A_T) # Transition matrix prior
    for n = 1:N # For each location
        h_0[n] ~ Categorical([1.0, 0.0, 0.0]) # Each location is initially operational

        h_min = h_0[n]
        for m = 1:M # For each timepoint
                h[m, n] ~ Transition(h_min, T) # Transition model under matrix T
            end
            w[m, n] ~ Transition(h[m, n], O) # Observation model

            h_min = h[m, n] # Reset previous state
        end
    end
end

The observation matrix O doesn't get sampled because it's fixed.

What I would like to do now is to add a variable z = randomvar(N) that takes values in {0,1} and a node f of the form:
f(z[n],w[m,n])=z[n] * w[m,n]+(1-z[n])* [1,0] and modify the model to be

@model function batch(M, N)
    h_0 = randomvar(N) # Initial state per location
    h   = randomvar(M, N) # Hidden states per location over time
    w   = randomvar(M, N) # Observed states per location over time
    z    = randomvar(N) # new variable that modifies the reading of location n
    y   = datavar(Vector{Float64}, M, N) # Observed states per location over time

    T ~ MatrixDirichlet(A_T) # Transition matrix prior
    for n = 1:N # For each location
        z[n]~ somedistribution...
        h_0[n] ~ Categorical([1.0, 0.0, 0.0]) # Each location is initially operational
        h_min = h_0[n]
        for m = 1:M # For each timepoint
                h[m, n] ~ Transition(h_min, T) # Transition model under matrix T
            w[m, n] ~ Transition(h[m, n], O) # Observation model
            y ~ f(z[n],w[m,n])
            h_min = h[m, n] # Reset previous state
           
        end
    end
end

So far I'm not managing to make it work, but I don't know if it's because I'm doing something wrong or because it's too complicated a model and what I want to do is simply not possible.

Cheers,
Alvaro

[ismailsenoz]

Hi Alvaro,
Nice to hear that you are playing around with RxInfer. Perhaps a bit more clarification is needed to identify the problem. As far as I understand $z_n$ is a Bernoulli distributed random variable (or a discrete random variable with 2 possible outcomes with a different density function but for simplicity lets assume it is Bernoulli). If $z_n$ is as such, definition of the non-linear function $f(z_n, w_{nm})$ needs clarification. Clarification is needed because the random variable $w_{nm}$ is assumed to evolve according to a 3 by 3 transition matrix, whereas in the function $f$ you are trying to add a 2 dimensional vector to a 3 dimensional vector which would return an error. Would you mind writing explicitly the equation without the indices and in tex notation so there is no code related interpretation? That being said I have written down the model and with the node definition a simple scenario runs on my machine.

using RxInfer

struct NonlinearNode end 

struct NonlinearMeta{F}
    fn       :: F   # Nonlinear function, we assume 2 input - 1 output
end

@node NonlinearNode Deterministic [ out, z , w]

f(z,w) = z * w + (1-z)*[1, 0, 0]

@rule NonlinearNode(:z, Marginalisation) (m_out::PointMass, m_w::DiscreteNonParametric, meta::NonlinearMeta) = begin 
    out = mean(m_out)
    ρ   = probvec(m_w)
    samples = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    p = dot(map((w) -> isequal(out, meta.fn(0,w)), samples)', ρ)
    ρ = ExponentialFamily.softmax([p, 1-p])
    return Bernoulli(last(ρ))
end

@rule NonlinearNode(:w, Marginalisation) (m_out::PointMass, m_z::Bernoulli, meta::NonlinearMeta) = begin 
    out = mean(m_out)
    (ρ, mρ)   = probvec(m_z)
    samples = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    p = map((w) -> ρ*isequal(out, meta.fn(0, w)) + mρ*isequal(out, meta.fn(1, w)) , samples)
    
    return Categorical(ExponentialFamily.softmax(p))
end

@model function batch(M, N, p, A_T, O_T)
    h_0 = randomvar(N) # Initial state per location
    h   = randomvar(M, N) # Hidden states per location over time
    w   = randomvar(M, N) # Observed states per location over time
    z   = randomvar(N) # new variable that modifies the reading of location n
    y   = datavar(Vector{Int64}) # Observed states per location over time

    T ~ MatrixDirichlet(A_T) # Transition matrix prior
    O ~ MatrixDirichlet(O_T) # Transition matrix prior
    for n = 1:N # For each location
        z[n] ~ Bernoulli(p) 
        h_0[n] ~ Categorical([1.0, 0.0, 0.0]) # Each location is initially operational
        h_min = h_0[n]
        for m = 1:M # For each timepoint
            h[m, n] ~ Transition(h_min, T) # Transition model under matrix T
            w[m, n] ~ Transition(h[m, n], O) # Observation model
            y ~ NonlinearNode(z[n], w[m,n])
            h_min = h[m, n] # Reset previous state
           
        end
    end
end

@constraints function batch_constraints()
    q(h_0, h, T, O, z, w) = q(h_0,h)q(T)q(O)q(z)q(w)
end

imodel = batch(3, 1, 0.1, ones(3,3),ones(3,3))
data = [1, 0, 0]
idata = (y = data,)

imarginals = (
    T = vague(MatrixDirichlet, 3, 3), 
    O = vague(MatrixDirichlet, 3, 3), 
    h = vague(Categorical, 3),
    z = vague(Bernoulli),
    w = vague(Categorical, 3)
)

ireturnvars = (
    T = KeepLast(),
    O = KeepLast(),
    h= KeepLast(),
    z = KeepLast(),
    w = KeepLast(),
)

@meta function nmeta(fn)
    NonlinearNode(y, z, w) -> NonlinearMeta(fn)
end

result = infer(
    model         = imodel, 
    data          = idata,
    constraints   = batch_constraints(),
    meta          = nmeta(f),
    initmarginals = imarginals, 
    returnvars    = ireturnvars, 
    iterations    = 20, 
);

I am unsure if this is what you intended to do, but perhaps it is close. Let us know.
Cheers

[apiedrafTNO]

Hi Ismail,
I see now that I should have specified a bit more of the model.
The latent variable h[m,n] is 3-dimensional, but w[m,n] is only 2-dimensional because the observation matrix is 3x2.

The equation of the non-linear node is $f(z,w)=z \cdot w+(1-z)\cdot[1,0]$.
The interpretation is that, if $z=0$, then the output of the node should be [1,0] independent of $w$, whereas if $z=1$, then the output is $w$. It's like a conditional overwrite.
The distribution of $z[n]$ is indeed not a Bernoulli but a $Probit(d-n)$ where $d$ is normally distributed, but I think that assuming it to be a Bernoulli is also fine for now. One step at a time.

Seeing your code I see that what I was missing was the message passing rules. I'll tweak these rules to 2-dimensional vectors and try it.

Thanks for such a complete and quick response!

Cheers,
Alvaro

[ismailsenoz]

In the probit extension the rules will be more complicated. We do have approximation strategies for that scenario in our arsenal of approximations. Nevertheless, this functionality is not public. If you need further clarifications please reach out.