Sample back the data from trace

numpyro
#1

Hi, I have a simple regression model with data inputs and sampled coefficients. My question is once I have sampled the regression coefficients, is there a way to sample back the possible data? It’s probably easier to look at the code.

import arviz as az
import jax.numpy as jnp
from jax import ops
import jax.random as random
import numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive, init_to_median
from numpyro import handlers

import matplotlib.pyplot as plt

numpyro.set_host_device_count(4)
numpyro.enable_x64(True)

N = 320
alpha, beta = -2., 0.4
scale = 1.2
rng = np.random.RandomState(32)
a = rng.standard_normal(N)
b = rng.standard_normal(N)
noise = rng.standard_normal(N)*scale
y = alpha*a + beta*b + noise

def model(a=None, b=None, y=None):
    alpha = numpyro.sample('alpha', dist.Normal(0, 0.5))
    beta = numpyro.sample('beta', dist.Normal(0, 0.5))
    scale = numpyro.sample('scale', dist.Exponential(1.))
    loc = alpha*a + beta*b
    obs = numpyro.sample('obs', dist.Normal(loc, scale), obs=y)

rng_sample, rng_post = random.split(random.PRNGKey(32), 2)
hmc = MCMC(
    NUTS(model, init_strategy=init_to_median),
    num_warmup=1000,
    num_samples=1000,
    num_chains=4,
    progress_bar=False,
)
hmc.run(rng_sample, extra_fields=("z", "energy", "diverging"), y=y, a=a, b=b)
ppc = Predictive(model, hmc.get_samples())(rng_post, a=a, b=b)
idata = az.from_numpyro(hmc,  posterior_predictive=ppc, )
print( az.summary(idata) )


a_new = rng.standard_normal(N//2)
b_new = rng.standard_normal(N//2)
noise_new = 0.
y_new = alpha*a_new + beta*b_new + noise_new
ppc_new = Predictive(model, hmc.get_samples())(rng_post, a=a_new, b=b_new)


do1 = handlers.do(model, data={'obs': y_new, 'a':a_new})
p1 = Predictive(do1, hmc.get_samples())(rng_post, a=a_new, b=XXXXX )

As expected ppc_new['obs'] gives values close to y_new.
However is it possible to provide y_new, a_new and get back samples of b_new ?