PyMC3 posterior prediction example; how to incorporate a matrix of betas?

Questions
15

Thanks so much! I have been digging into the linked notebook and stumbled across another problem. It requires the treatment variable to be inside the InferenceData object.

level_labels = pooled_idata.posterior.Level[pooled_idata.constant_data.floor_idx]

I really would like to keep my treatment beta separate from the covariates and find it convenient to do so; especially in higher dimensional orders. The problem is, I cannot get my treatment coords inside the InferenceData object when it is a single beta parameter (coords appears to be only for multi-dimensional items).

# coords
coords = {
          # dim 1: length of data set
          'obs_id': np.arange(len(target)),
          # dim 2: treatment var
          'treatment': ['treatment'],
          # dim 3: covariates
          'covariates': ['covariate_1', 'covariate_2']
    }

# specify model
with pm.Model(coords=coords) as sk_lm:
    # data
    treatment_data = pm.Data('treatment_data', treatment, dims=('obs_id'))
    covariates_data = pm.Data('covariates_data', covariates, dims=('obs_id', 'covariates'))

    # priors
    alpha = pm.Normal('alpha', mu=0, sigma=1)
    treatment_beta = pm.Normal("treatment_beta", mu=30, sigma=2, dims=('treatment'))  # fails
    covariates_betas = pm.Normal('covariates_betas', mu=[66, 33], sigma=2, dims=('covariates'))

If I include a dims argument (which would solve the issue) and link it to the coords, then:

# prior analysis
with sk_lm:
    prior_pc = pm.sample_prior_predictive()

Throws a value error, which after searching, seems to be associated with the shape of mu and it does not like specifying a shape value size of 1.

ValueError: 0-dimensional argument does not have enough dimensions for all core dimensions ('i_1_0',)

How might I account for a single parameter such that it can be linked to the coords and therefore available under InferenceData.posterior?

I have added a category manually like so in the meantime:

lm_idata.posterior.coords['treatment'] = ['treatment']

Edit

When I get home I hope to try using an idata kwarg, drawing on: PyMC3 with labeled coords and dims | Oriol unraveled