Posterior samples from multi-level model when choosing the index values

I have a simple multi-level model without any parameter pooling in the model defined in the following way without any pooling of parameters:

coords = {"Country":country_list, "obs_id":np.arange(yearly.shape[0])}
with pm.Model(coords=coords) as unpooled_model:
    country_idx = pm.Data("country_idx", yearly.country_idx, dims="obs_id")
    a = pm.Exponential("a", 1/10, dims="Country")
    counts = pm.Poisson("counts", a[yearly.country_idx], observed = yearly["Gold"], dims="obs_id")

I can obtain posterior predictive samples from this model straightforwardly and visualise it with arviz. However, each country index (country_idx) value appears a different number of times within this dataset, and this results in getting imbalanced posterior predictive samples, coming from each country. What I would really like to do is only get posterior predictive samples for certain country_idx values. Whenever I have tried to do this using pm.set_data({“country_idx”:[1.4]}) for instance, I get errors regarding conflicts with the shape of the obs_id dimension in the original model. Is there a smart/easy way to do this?

I have looked through the tutorials on multi-level models and this isn’t done there either. For instance, in the linked tutorial, imagine I wanted posterior predictive samples, but only for the aitkin basement, and not the others. What is the best way to do this in pymc3?

Thanks in advance :slight_smile:

Does my answer on this thread help? There’s a link to a new forthcoming example notebook.

I don’t think the link worked. Could you re-attach it?

Ah yes - The example has been merged into pymc-examples so I deleted by branch. It’s not yet live on the website, but you can access the notebook from the pymc-examples repo…