Add support for multiple BART random variables per model.

[derekpowell]

This PR adds new functionality to support multiple BART RVs defined within a single pymc model, addressing and resolve the limitation discussed in #86.

• added ability to support multiple BART RVs in a single model (each gets its own pgbart sampler with combined sampling handled by pymc's compound sampling functionality)
• tests of new functionality

Example use

This addition allows for some new modeling possibilities. One example is the Bayesian Causal Forests model of Hahn et al.. This model aims to capture heterogenous treatment effects of a binary treatment $Z$ by separately capturing the prognostic predictiveness of those covariates $X$ and their interactions with $z$ (note that different $X$ may also be modeled by $\mu$ and $\tau$).

$$f(x_i, z_i) = \mu(x_i) + \tau(x_i) \cdot z_i$$

This can be implemented as:

import pymc as pm
import pymc_bart as pmb

with pm.Model() as bcf_model:
    # Prognostic function (baseline outcome) - BART1
    mu = pmb.BART("mu", X=X_prog, Y=y, m=50) 
    
    # Treatment effect function - BART2
    tau = pmb.BART("tau", X=X_trt, Y=y, m=50)
    
    # Combine the two BART functions
    # Mean function: mu + Z * tau
    mean_y = mu + Z * tau
    
    # Observation noise
    sigma = pm.HalfNormal("sigma", sigma=1.0)
    
    # Likelihood
    y_obs = pm.Normal("y_obs", mu=mean_y, sigma=sigma, observed=y)

Hope this looks good. I also plan to offer a more detailed demo example for the pymc-examples doc repo if desired!

[aloctavodia]

Thanks, this is great! Adding one example will be fantastic!

[aloctavodia]

Before approving, I tested locally with a couple of examples, and they ran. But the test is failing. Funny that the order of the variables in the model affects the result. This works

        sigma = pm.HalfNormal("sigma", 1)
        mu1 = pmb.BART("mu1", X1, Y1, m=5)
        mu2 = pmb.BART("mu2", X2, Y2, m=5)

but this fails

        mu1 = pmb.BART("mu1", X1, Y1, m=5)
        mu2 = pmb.BART("mu2", X2, Y2, m=5)
        sigma = pm.HalfNormal("sigma", 1)

[aloctavodia]

Is this necessary?

[derekpowell]

Nope! Sorry about that, unnecessary since it's already int

[derekpowell]

Before approving, I tested locally with a couple of examples, and they ran. But the test is failing. Funny that the order of the variables in the model affects the result. This works

        sigma = pm.HalfNormal("sigma", 1)
        mu1 = pmb.BART("mu1", X1, Y1, m=5)
        mu2 = pmb.BART("mu2", X2, Y2, m=5)

but this fails

        mu1 = pmb.BART("mu1", X1, Y1, m=5)
        mu2 = pmb.BART("mu2", X2, Y2, m=5)
        sigma = pm.HalfNormal("sigma", 1)

Hmm that's odd---I can't reproduce this behavior. I tried swapping the order of NUTS and pgbart variables in my tests and they still passed, and in my example notebook I'm working on and it ran. Or are you saying that it runs but the results are just different? To some extent I'd expect that with the sequential nature of the sampling steps but ideally they wouldn't be too different. If you can share reprex I can look into it?

[aloctavodia]

Locally I am getting the same error than here, it does not run unless the bart RVs are defined last. I am away from my laptop most likely until Monday. Are you using the last pymc version?

[derekpowell]

Oh, looks like I am on pymc = 5.19.1 (I used environment-dev.yml)

[derekpowell]

Alright I am able to reproduce this issue with pymc 5.23. Things work if you manually specify the compound step as in tests/test_bart.py::test_multiple_bart_variables_manual_step. Looks like the breaking change / regression (if that's what it is) comes in pymc v5.20.1.

For now, I suppose I could:

1. Remove the failing test
2. Add a version check in PGBART to warn that compound steps must be manually specified for versions >=5.20.1 (I can also figure out what most up-to-date release that works is). This is a bit suboptimal, not sure if this is too hacky.

I'm afraid that diagnosing what has changed on the pymc side beyond that is likely to be beyond me.