Indexing for vector variable

Questions v5
1

What’s the syntax to index a vector-valued variable? The following code does what I want, but it’s rather clunky (and will only get more clunky the longer the vector).

sample_idx, sample_dims = pd.factorize(df['sample'])

coords = {'samples': sample_dims}

with pm.Model(coords=coords) as my_model:
   samples_indices = pm.Data('samples_indices', sample_idx)

   intercept_x = pm.Normal('intercept_x', mu=0, sigma=1, dims=('samples'))
   intercept_y = pm.Normal('intercept_y', mu=0, sigma=1, dims=('samples'))

   intercepts = pm.math.stack([intercept_x[samples_indices], intercept_y[samples_indices]])

Surely there’s a better way to do this. Something like the following?

sample_idx, sample_dims = pd.factorize(df['sample'])

coords = {'samples': sample_dims,
          'vars': ('x', 'y')}

with pm.Model(coords=coords) as my_model:
   samples_indices = pm.Data('samples_indices', sample_idx)

   intercept = pm.Normal('intercept', mu=0, sigma=1, dims=('samples', 'vars'))

   intercepts = intercept[samples_indices,:]

That doesn’t work, but I assume there must be some similarly elegant syntax to do this.

2

I’m like 99% sure the 2nd code snippet should work. What is the specific error you get?

3

Ah, you’re right. I was getting an error, but it was because I got the order of the dims wrong, and was trying to add an nxm matrix to an mxn matrix… :man_facepalming: