Quick answer without playing with the code yet. treatment is still a dataframe/series, not an xarray object, so this should not work. Two main approches come to mind:
Try doing treatment.to_xarray(), pandas and xarray allow to automatically convert data from one to the other.
# specify model
with pm.Model(coords=coords) as main_lm_model:
t = pm.Data("treatment", treatment, dims=("obs_id", "treatment_covariate")
...
# then get treatment as
lm_idata.constant_data["treatment"]
Note that you are currently using * when I think you want a matrix product, not a pointwise product. You should be able to use np.linalg.matmul directly, as it acts on the last dimensions only and batches over any initial dimension as necessary (chain and draw in this case).
Yes, if you get an inferencedata from sample you need to use extend like it’s done in the rugby example.
side comment, you already have alpha, treatment_betas… in the posterior, “sampling” them in the posterior predictive is actually not doing anything.