Indexing constantdata by label

Questions v5
1

Hello!
I’ve been working on getting used to the xarray coords/dims concept in pymc, and I really like it (particularly that labels get passed on to arviz!).

Sometimes one creates a normal distribution with a few entries in one dimension, and then indexes into the distribution variable to extract only one of them. (very minimal example below)

with pm.Model(coords={"animal": ["cat", "dog"]}) as model:
    mu = pm.Normal("mu", mu=0, sigma=1, dims="animal")

    pm.Deterministic("cat_mu", mu[0])  # just an example of mu[0] being used

With an xarray DataArray, I could index by label in the following manner:

import xarray as xr

da = xr.DataArray(
    data=[1,2], # this data would represent the values of the Normal distribution, took me some time to understand how coords are used in xarray after coming in from using them in xarray
    coords={"animal": ["cat", "dog"]},
)

da.loc["cat"] # index like this

Is there any way to use such syntax in pymc? I think it could help catch errors, since one might mix up the order of which entry refers to cat or dog on the mu object.

with model:

    pm.Deterministic("cat_mu", mu[1])  # oh no, index 1 refers to dog, not cat!
    pm.Deterministic("cat_mu_proper", mu.loc["cat"])  # woohoo, this is guaranteed cat!