How to keep dimensions in a posterior predictive sample?

Questions v5
6

Hi, I’m still mulling this question, and I have another solution (in case someone like me happens across this in the future)

I believe I was overthinking the model. I actually wanted to be observing two totally different variables, so the dimensions of event_count needed to reflect that. The following does exactly what I wanted.

Instead of keeping the data in a ‘long’ format, i keep it wide, with separate columns for each ‘event type’.

n = 1000
ns = np.arange(0,n)

blue_lam = 1.4
green_lam = 3

blue_dist = poisson(blue_lam)
green_dist = poisson(green_lam)
blues = blue_dist.rvs(n)
greens = green_dist.rvs(n)

dt = pd.DataFrame.from_dict({'blue': blues, 'green':greens})

types = dt.columns

coords = {
    "Event Types": types,
    "obs_id": np.arange(dt.shape[0])
}

mdl = pm.Model(coords=coords)
with mdl:
    mu = pm.Uniform("mu",0,5, dims=("Event Types"))
    event_count = pm.Poisson("event_count", mu=mu, observed=dt, dims = ("obs_id","Event Types"))

The posterior predictive samples now give me samples for the different event types. (note that I need to flatten the posterior predictive plot by the dimension i don’t want to see separate plots for.

with mdl:
    post = pm.sample_posterior_predictive(trace)
az.plot_ppc(post, flatten=["obs_id"])

image
image821×314 19.4 KB