Accesing a prior array coordinate from the dimension string label

Questions v5
1

Hello I have been reading the dimensionality tutorial and I am trying to use dimensions to access the prior array values in a dictionary-like maner.

For example, if I need to define priors for the abundance of chemical elements I would have something like this:

abund_dict = {‘oxygen’: pm.Normal(‘oxygen’, mu=5, sigma=5),
‘nitrogen’: pm.Normal(‘nitrogen’, mu=5, sigma=5)
‘sulfur’: pm.Normal(‘sulfur’, mu=5, sigma=5)
‘argon’: pm.Normal(‘argon’, mu=5, sigma=5)}

In this case, I have a dictionary with 4 scalar priors. However, it would simplify my model declaration considerably if I could use a prior vector to declare the 4 element abundances:

abund_priors = pm.Normal(‘abundance’, mu=5, sigma=5, dim=(‘oxygen’, ‘nitrogen’, ‘sulfur’, ‘argon’)

However, is it possible to access a prior coordinate from this array with the dimension name?

element_abundance = abundance_priors.magic_step(‘oxygen’)

I would be very thankful for any insight on the right way to do this.

2

You need to distinguish between “dims” and “labels”. “dims” are sets, and “labels” are set elements. In your case, the dim is “element” and the labels are (‘oxygen’, ‘nitrogen’, ‘sulfur’, ‘argon’)

So you define your model like:

coords = {'element':[‘oxygen’, ‘nitrogen’, ‘sulfur’, ‘argon’]}
with pm.Model(coords=coords) as m:
    abund_priors = pm.Normal('abundance, mu=5, sigma=5, dims='element')

Note that the dimension names and labels are used by arviz for post-estimation tasks (manipulating idata and making plots). They are not used by pytensor inside the actual PyMC model. So if you wanted to select just the prior for oxygen, for example, you need to resort to numerical indexing:

    oxygen_abundance = abund_priors[coords['element'].index('oxygen')]

Of course, if you knew that oxygen is at position 0 you could just directly do that.

On the other hand, if you’ve already sampled the model and want to see values for oxygen, you can use the idata.sel method, as in oxygen_posterior = idata.posterior.abundance_priors.sel(element = 'oxygen')

2 Likes
3

Thank you @jessegrabowski very much for the confirmation.