I want to add the xarray version of this too if you don’t mind, which gets the same result as the numpy code above:
Provided you label your dimensions in the model:
with pm.Model() as model:
x = pm.ConstantData("x", xs, dims=("predictor", "obs_id"))
a = pm.Normal("a", 0.0, 10.0)
betas = pm.Normal('betas', mu=0, sigma=5, dims="predictor")
mu = pm.Deterministic("mu", a + pm.math.dot(betas, xs), dims="obs_id")
pm.Normal("obs", mu=mu, sigma=1, observed=obs, dims="obs_id")
idata = pm.sample()
You can then do:
a = az.extract(idata, var_names="a")
b = az.extract(idata, var_names="betas")
mu_pp = a + xr.dot(idata.constant_data["x"], b, dims="predictor")
plt.scatter(obs, mu_pp.mean("sample"))
plt.show()
Note the definition of a
and b
would also work if used in the code above, extract
defaults to returning combined variables from the posterior (unless a different group is specified).
Later, instead of indicating what dimension to reduce with the tiling, transposing or using integer axis ids, we do so by name. We reduce the predictor dimension in the dot product, and the sample dimension in the mean, so we are only left with the obs_id dim which is the one we want to plot.