Modeling bidirectional relationships / loopy graphs

In an application scenario, I have two variables that I have to assume are mutually influencing each other. Let’s say the model could look like this, with all variables binary and all variables except for B observed:

A ---> (B) <---> D
        |        |
        v        |
        C <----- +

How would I go about modeling the bidirectional edge between B and D in a (num)pyro model?

Ideally, I would like to do something like

...
B = sample("B", dist.Bernoulli(logits=k1 + k2*A + k3*D))
D = sample("D", dist.Bernoulli(logits=k4 + k5*B), obs=...)
...

but this obviously does not work because D is not yet defined in the first line. (The graph is loopy.)

I guess one approach would be to merge B and D into a single 2D binary variable and model their joint distribution instead? That’s a) a little “ugly” to me personally but ok, but more importantly b) how would I then handle the fact that D is observed but B is not?

I could also imagine that I could do something using a factor, but I have never worked with that and am a bit unsure about how I would have to apply that here.

(I do realize that this is at least as much a general modeling question as it is a pyro-specific question…)

Edit: I see now that there’s obs_mask that might work for the partial observations - will try to get that to work tomorrow.

factor would be the most generic way to encode bi-directional factors

Trying to piece things together from other forum posts, do you mean something as simple as

...
B = sample("B", dist.ImproperUniform(dist.constraint.boolean, (), ()))
D = d_obs
factor("B_log_prob", dist.Bernoulli(logits=k1 + k2*A + k3*D).log_prob(B))
factor("D_log_prob", dist.Bernoulli(logits=k4 + k5*B).log_prob(D))
...

or (equivalently, I guess?)

...
B = sample("B", dist.ImproperUniform(dist.constraint.boolean, (), ()))
D = sample("D", dist.Bernoulli(logits=k4 + k5*B), obs=d_obs)
factor("B_log_prob", dist.Bernoulli(logits=k1 + k2*A + k3*D).log_prob(B))
...

should work & do what I want?

Is there any variation to this pattern that would allow me to also sample from the model?

factor(“B_log_prob”, dist.Bernoulli(logits=k1 + k2A + k3D).log_prob(B))

yes that sort of thing

Is there any variation to this pattern that would allow me to also sample from the model?

a model built as a DAG out of primitive distributions (each of which has a sampler associated with it) is exactly that class of distributions that admits easy (ancestral) sampling: once you go beyond that (e.g. a random markov field) you generically need to do some sort of inference.

Thanks again! I am able to get this to work for a continuous example, i.e., the following seems to work:

def model():
    # True and unique solution is A=0, B=1, C=2, D=3. (Those are the means, respectively.)
    A = sample("A", dist.Normal(0, 1))
    B = sample("B", dist.ImproperUniform(dist.constraints.real, (), ()))
    D = sample("D", dist.Normal(1 + 2*B, 1))
    C = sample("C", dist.Normal(D - B, 1))    
    numpyro.factor("B_log_prob", dist.Normal(A + D-2, 1).log_prob(B))

nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_samples=20000, num_chains=8, num_warmup=10000)
mcmc.run(random.PRNGKey(0))
mcmc.print_summary()  # yields means 0, 1, 2, 3, as expected

However, I am struggling to make a discrete version work. I thought the following should work:

def model():
    # True and unique solution is p(A)=0.2, p(B)=0.4, p(C)=0.6, p(D)=0.8.
    A = sample("A", dist.Bernoulli(probs=0.2), infer={"enumerate": "parallel"})
    B = sample("B", dist.ImproperUniform(dist.constraints.boolean, (), ()), infer={"enumerate": "parallel"})
    D = sample("D", dist.Bernoulli(probs=1 - 0.5*B), infer={"enumerate": "parallel"})
    C = sample("C", dist.Bernoulli(probs=0.5*B + 0.5*D), infer={"enumerate": "parallel"})
    numpyro.factor("B_log_prob", dist.Bernoulli(probs=0.1 + 0.5*A + 0.25*D).log_prob(B))

infer_discrete(model, rng_key=random.PRNGKey(0))()

… but this yields a NotImplementedError while trying to sample from the ImproperUniform.

If I try to do

nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_samples=20000, num_chains=8, num_warmup=10000)
mcmc.run(random.PRNGKey(0))

I get a ValueError: Continuous inference cannot handle discrete sample site 'B'. (Same problem if I use kernel = MixedHMC(HMC(model)) instead of NUTS.)

How do you do inference for a model with a dist.ImproperUniform(dist.constraints.boolean)?

i don’t think that will work

you can probably do something like

B = sample("B", dist.Bernoulli(probs=0.5), infer={"enumerate": "parallel"})

or (for good measure)

B = sample("B", dist.Bernoulli(probs=0.5).mask(False), infer={"enumerate": "parallel"})

to explicitly mask out the corresponding log_prob (a flat prior doesn’t effect the posterior so either is ok)