How to properly do out-of-sample prediction for hierarchical model

Questions v5
, ,
23

Ok, so “id” stays [0,1], but then new_idx must have to change to [1, 1, 2, 3] if we want to predict on one previous group (1) and two new groups (2, 3) correct?

So for example this code works

# Predict two new groups
new_idx = [1, 1, 2, 3]
new_coords = {
    "id":[0, 1],
    "new_id": [2, 3],
    "obs_idx": list(range(len(new_idx))),
}
with pm.Model(coords=new_coords) as pred_m:
    x = pm.Data("x", [0, 0, 0, 0], dims=("obs_idx",))
    mu_group = pm.Normal("mu_group")
    # Needs a new name, so old b is not used!
    new_b = pm.Normal("new_b", mu_group, dims=("new_id",))
    b = pm.Normal("b", mu_group, dims=("id",))
    mu = pm.math.concatenate([new_b[new_idx], b[idx]]) * x
    obs = pm.Normal("obs", mu, dims=("obs_idx",))

    idata = pm.sample_posterior_predictive(idata, var_names=["obs"], predictions=True, extend_inferencedata=True)