Coordinates in PyMC & InferenceData Objects

PyMC v4 is here and one of the big changes is that the inference routines (e.g., pm.sample()) will return an ArviZ InferenceData object by default (recent releases of PyMV3 have made this optional). This post is intended to explore a bit about how the InferenceData object works, how it intersects with PyMC parameter arrays and how the labeling of dimensions and coordinates can make life easier. If you haven't already checked out Oriol Abril's fantastic post about PyMC, coords, and dims, I would strongly encourage you to do so.

Oriol's post focuses on the use of labeled dimensions in PyMC models and how Arviz acknowledges and utilizes the labels when processing the prior/posterior (e.g., plotting). Here, I want to focus on how named dimensions/cooridnates can be used to create parameter arrays, contrasting it with the way things were done in PyMC3. I will also illustrate how to these coordinates structure the InferenceData objects in which our posterior is stored.

Table of Contents¶

• xarray
• Example Data
• Unlabeled Approach (Old)
• Labeled Approach (New)
• Using pm.Data
• Wrap-Up

In [1]:

import arviz as az
import numpy as np
import pandas as pd
import pymc as pm
import xarray as xr

rng = np.random.default_rng(101010)

xarray ¶

The new(ish) ArviZ inferenceData object is basically an xarray Dataset object. xarray describes itself as a generalization of both Numpy arrays and Pandas dataframes. Its design also borrows heavily from the netCDF file specification. But let's see what xarray objects actually look like. Here is a standard numy array.

In [2]:

npdata = rng.standard_normal(size=(2, 3))
npdata

Out[2]:

array([[-0.91412107, -0.32402107,  1.29165606],
       [ 0.52327975,  0.95708885,  0.52139111]])

Here we will use this numpy array to create an xarray DataArray.

In [3]:

data = xr.DataArray(npdata,
                    dims=("user", "day"),
                    coords={"user": ['Alice', 'Bob'], 
                            "day": ["yesterday", "today", "tomorrow"]
                           }
                   )

data

Out[3]:

<xarray.DataArray (user: 2, day: 3)>
array([[-0.91412107, -0.32402107,  1.29165606],
       [ 0.52327975,  0.95708885,  0.52139111]])
Coordinates:
  * user     (user) <U5 'Alice' 'Bob'
  * day      (day) <U9 'yesterday' 'today' 'tomorrow'

xarray.DataArray

• user: 2
• day: 3

• -0.9141 -0.324 1.292 0.5233 0.9571 0.5214

  array([[-0.91412107, -0.32402107,  1.29165606],
         [ 0.52327975,  0.95708885,  0.52139111]])

• Coordinates: (2)
  • user
    (user)
    <U5
    'Alice' 'Bob'

    array(['Alice', 'Bob'], dtype='<U5')

  • day
    (day)
    <U9
    'yesterday' 'today' 'tomorrow'

    array(['yesterday', 'today', 'tomorrow'], dtype='<U9')

• Attributes: (0)

As you can see, the data hasn't changed at all, but it now comes with all sorts of extra information. Let's get a bit of terminology out of the way. Our xarray object is an DataArray. The array has 2 dimensions: one labeled "user" and the other labeled "day".

In [4]:

data.dims

Out[4]:

('user', 'day')

Each dimension has several coordinates. The "user" dimension includes coordinates "Alice" or "Bob". The "day" dimension includes coordinates "yesterday", "today", and "tomorrow".

In [5]:

data.coords

Out[5]:

Coordinates:
  * user     (user) <U5 'Alice' 'Bob'
  * day      (day) <U9 'yesterday' 'today' 'tomorrow'

These coordinates can be used instead of standard numerical indicies. Here is how we select the value corresponding to "Alice" and "yesterday" from our numpy array.

In [6]:

npdata[0,0]

Out[6]:

-0.9141210682943073

And here is how we select that same value from our xarray DataArray.

In [7]:

data.sel(user="Alice", day="yesterday")

Out[7]:

<xarray.DataArray ()>
array(-0.91412107)
Coordinates:
    user     <U5 'Alice'
    day      <U9 'yesterday'

xarray.DataArray

• -0.9141

  array(-0.91412107)

• Coordinates: (2)
  • user
    ()
    <U5
    'Alice'

    array('Alice', dtype='<U5')

  • day
    ()
    <U9
    'yesterday'

    array('yesterday', dtype='<U9')

• Attributes: (0)

So one benefit of xarray objects is that there are a nice, semantically labeled coordinates. You no longer have to remember that npdata[:,0] corresponds to "yesterday" and that npdata[1,:] corresponds to "Bob". But it's even better than that. You don't even need to remember which axis is which. The order in which the dimensions are stored is not relevant because you are just using the dimensions' labels (though you can use numerical indexing if you wish).

In [8]:

data.sel(day="yesterday", user="Alice")

Out[8]:

<xarray.DataArray ()>
array(-0.91412107)
Coordinates:
    user     <U5 'Alice'
    day      <U9 'yesterday'

xarray.DataArray

• -0.9141

  array(-0.91412107)

• Coordinates: (2)
  • user
    ()
    <U5
    'Alice'

    array('Alice', dtype='<U5')

  • day
    ()
    <U9
    'yesterday'

    array('yesterday', dtype='<U9')

• Attributes: (0)

In [9]:

data.sel(user="Alice", day="yesterday")

Out[9]:

<xarray.DataArray ()>
array(-0.91412107)
Coordinates:
    user     <U5 'Alice'
    day      <U9 'yesterday'

xarray.DataArray

• -0.9141

  array(-0.91412107)

• Coordinates: (2)
  • user
    ()
    <U5
    'Alice'

    array('Alice', dtype='<U5')

  • day
    ()
    <U9
    'yesterday'

    array('yesterday', dtype='<U9')

• Attributes: (0)

Finally, xarray also permits the creation of objects that include multiple DataArray objects. xarray calls such an object a Dataset. This is essentially what we will be dealing with when we get an ArviZ inferenceData object back after sampling.

In [10]:

var2 = xr.DataArray(rng.standard_normal(size=(2, 2)),
                    dims=("x", "y"),
                    coords={"x": [0,1], 
                            "y": [11,42]}
                   )
var3 = xr.DataArray(rng.standard_normal(size=(2, 2)),
                    dims=("a", "b"),
                    coords={"a": [4.2, 11.8], 
                            "b": ['Geneva','London']}
                   )
ds = xr.Dataset(dict(orig=data, 
                     v2=var2, 
                     v3=var3)
               )
ds

Out[10]:

<xarray.Dataset>
Dimensions:  (user: 2, day: 3, x: 2, y: 2, a: 2, b: 2)
Coordinates:
  * user     (user) <U5 'Alice' 'Bob'
  * day      (day) <U9 'yesterday' 'today' 'tomorrow'
  * x        (x) int64 0 1
  * y        (y) int64 11 42
  * a        (a) float64 4.2 11.8
  * b        (b) <U6 'Geneva' 'London'
Data variables:
    orig     (user, day) float64 -0.9141 -0.324 1.292 0.5233 0.9571 0.5214
    v2       (x, y) float64 -0.2851 -1.626 -0.08508 0.9506
    v3       (a, b) float64 -0.4798 1.463 0.07777 0.09208

xarray.Dataset

• Dimensions:
  • user: 2
  • day: 3
  • x: 2
  • y: 2
  • a: 2
  • b: 2
• Coordinates: (6)
  • user
    (user)
    <U5
    'Alice' 'Bob'

    array(['Alice', 'Bob'], dtype='<U5')

  • day
    (day)
    <U9
    'yesterday' 'today' 'tomorrow'

    array(['yesterday', 'today', 'tomorrow'], dtype='<U9')

  • x
    (x)
    int64
    0 1

    array([0, 1])

  • y
    (y)
    int64
    11 42

    array([11, 42])

  • a
    (a)
    float64
    4.2 11.8

    array([ 4.2, 11.8])

  • b
    (b)
    <U6
    'Geneva' 'London'

    array(['Geneva', 'London'], dtype='<U6')

• Data variables: (3)
  • orig
    (user, day)
    float64
    -0.9141 -0.324 ... 0.9571 0.5214

    array([[-0.91412107, -0.32402107,  1.29165606],
           [ 0.52327975,  0.95708885,  0.52139111]])

  • v2
    (x, y)
    float64
    -0.2851 -1.626 -0.08508 0.9506

    array([[-0.28513891, -1.62619068],
           [-0.08508324,  0.95058622]])

  • v3
    (a, b)
    float64
    -0.4798 1.463 0.07777 0.09208

    array([[-0.47979087,  1.46327566],
           [ 0.07776975,  0.09207604]])

• Attributes: (0)

Note that we have 6 dimensions (user, day, x, y, a, and b), but that each variable in our Dataset only uses 2 of them.

To get access to each individual variable out as a DataArray, we can use dictionary-style indexing:

In [11]:

ds['orig']

Out[11]:

<xarray.DataArray 'orig' (user: 2, day: 3)>
array([[-0.91412107, -0.32402107,  1.29165606],
       [ 0.52327975,  0.95708885,  0.52139111]])
Coordinates:
  * user     (user) <U5 'Alice' 'Bob'
  * day      (day) <U9 'yesterday' 'today' 'tomorrow'

xarray.DataArray

'orig'

• user: 2
• day: 3

• -0.9141 -0.324 1.292 0.5233 0.9571 0.5214

  array([[-0.91412107, -0.32402107,  1.29165606],
         [ 0.52327975,  0.95708885,  0.52139111]])

• Coordinates: (2)
  • user
    (user)
    <U5
    'Alice' 'Bob'

    array(['Alice', 'Bob'], dtype='<U5')

  • day
    (day)
    <U9
    'yesterday' 'today' 'tomorrow'

    array(['yesterday', 'today', 'tomorrow'], dtype='<U9')

• Attributes: (0)

xarray provides loads of indexing options, dimension/coordinate-aware computation, split-apply-combine operations, and much more. Go read up on xarray to discover all of its functionality.

Example Data ¶

In the examples below, we'll be working with some synthetic data involving two groups: one group receiving an experimental medication and one group receiving a placebo. For each patient, we have their age, sex, systolic blood pressure, and a flag indicating which treatment they have received. Each patient also has a unique ID. Let's generate it.

In [12]:

group_size = 100
systolic_f = rng.normal(loc=123, scale=5, size=group_size)
systolic_m = rng.normal(loc=127, scale=5, size=group_size)
systolic = np.hstack((systolic_f, systolic_m))
sex = (['female'] * group_size) + (['male'] * group_size)
age = rng.normal(loc=50, scale=7, size=len(systolic))
treat_group = rng.integers(low=1, high=3, size=len(systolic))
patient_id = rng.choice(1000, size=len(systolic), replace=False)
systolic = systolic + (.05 * age)
df = pd.DataFrame({'bp':systolic,
                   'sex':sex,
                   'age':age,
                   'treat_group':treat_group})
df.index = patient_id
df.index.name = "patientID"
df.head()

Out[12]:

	bp	sex	age	treat_group
patientID
979	128.750951	female	53.901474	2
924	123.935881	female	50.596643	1
522	121.683066	female	57.149881	1
912	135.144287	female	50.263664	2
787	121.420322	female	47.802665	1

Old Method ¶

Let's construct a model in which we estimate the mean of each group. But let's do it the way we might have in PyMC3, using an unlabeled parameter array. Because we need to index into the parameter array numerically, the first thing we have to do is to use panda's categorical data type to extract category codes (e.g., 0,1,2...) that numerically indicate group membership.

In [13]:

sex_idx,sex_codes = pd.factorize(df["sex"])
n_sex = len(sex_codes)

In [14]:

with pm.Model() as unlabeled_model:
    # hyper prior
    mu_hyper = pm.Normal('mu_hyper', mu=120, sigma=15)
    # per-group prior
    mu = pm.Normal('mu', mu=mu_hyper, sigma=15, shape=n_sex)
    # likelihood
    likelihood = pm.Normal('likelihood', 
                           mu=mu[sex_idx], 
                           sigma=15, 
                           observed=df['bp']
                          )

In [15]:

with unlabeled_model:
    unlabeled_idata = pm.sample()

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [mu_hyper, mu]

100.00% [4000/4000 00:03<00:00 Sampling 2 chains, 0 divergences]

Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 5 seconds.

In [16]:

unlabeled_idata.posterior

Out[16]:

<xarray.Dataset>
Dimensions:   (chain: 2, draw: 1000, mu_dim_0: 2)
Coordinates:
  * chain     (chain) int64 0 1
  * draw      (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
  * mu_dim_0  (mu_dim_0) int64 0 1
Data variables:
    mu_hyper  (chain, draw) float64 123.8 128.3 126.2 ... 128.1 132.3 116.9
    mu        (chain, draw, mu_dim_0) float64 124.1 128.6 125.5 ... 123.5 130.5
Attributes:
    created_at:                 2022-06-13T20:09:04.820492
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.0.0b6
    sampling_time:              4.526270627975464
    tuning_steps:               1000

xarray.Dataset

• Dimensions:
  • chain: 2
  • draw: 1000
  • mu_dim_0: 2
• Coordinates: (3)
  • chain
    (chain)
    int64
    0 1

    array([0, 1])

  • draw
    (draw)
    int64
    0 1 2 3 4 5 ... 995 996 997 998 999

    array([  0,   1,   2, ..., 997, 998, 999])

  • mu_dim_0
    (mu_dim_0)
    int64
    0 1

    array([0, 1])

• Data variables: (2)
  • mu_hyper
    (chain, draw)
    float64
    123.8 128.3 126.2 ... 132.3 116.9

    array([[123.83704342, 128.27390986, 126.19071339, ..., 129.33831544,
            129.85257537, 120.46615604],
           [133.48677709, 113.22021801, 131.0946918 , ..., 128.12171912,
            132.33759155, 116.9052638 ]])

  • mu
    (chain, draw, mu_dim_0)
    float64
    124.1 128.6 125.5 ... 123.5 130.5

    array([[[124.14093606, 128.55854405],
            [125.52497745, 127.65533327],
            [124.77976856, 128.77416628],
            ...,
            [124.64880362, 130.75414181],
            [124.34703766, 129.40272599],
            [126.18479451, 130.39614565]],
    
           [[125.82809488, 130.04700006],
            [123.95336703, 132.90786846],
            [126.00171379, 127.46920495],
            ...,
            [123.53689027, 129.44499793],
            [122.56246325, 129.00550263],
            [123.46022351, 130.45076266]]])

• Attributes: (6)

  created_at :

  2022-06-13T20:09:04.820492

  arviz_version :

  0.12.1

  inference_library :

  pymc

  inference_library_version :

  4.0.0b6

  sampling_time :

  4.526270627975464

  tuning_steps :

  1000

In [17]:

az.summary(unlabeled_idata, kind="stats", round_to=2)

Out[17]:

	mean	sd	hdi_3%	hdi_97%
mu_hyper	124.93	9.19	107.58	141.55
mu[0]	124.95	1.54	122.00	127.76
mu[1]	129.89	1.46	127.34	132.91

We can also investigate the posterior directly. Let's calculate the means of each group.

In [18]:

# average over chains and draws
unlabeled_idata.posterior['mu'].mean(dim=['chain', 'draw'])

Out[18]:

<xarray.DataArray 'mu' (mu_dim_0: 2)>
array([124.94933509, 129.89383768])
Coordinates:
  * mu_dim_0  (mu_dim_0) int64 0 1

xarray.DataArray

'mu'

• mu_dim_0: 2

• 124.9 129.9

  array([124.94933509, 129.89383768])

• Coordinates: (1)
  • mu_dim_0
    (mu_dim_0)
    int64
    0 1

    array([0, 1])

• Attributes: (0)

More practically, we can calculate the probability that the blood pressures of the males and females differ from one another.

In [19]:

( unlabeled_idata.posterior['mu'].sel(mu_dim_0=0) < 
  unlabeled_idata.posterior['mu'].sel(mu_dim_0=1) ).mean()

Out[19]:

<xarray.DataArray 'mu' ()>
array(0.991)

xarray.DataArray

'mu'

• 0.991

  array(0.991)

• Coordinates: (0)
• Attributes: (0)

Ok. We have our posterior and everything seems to be working. But dealing with the posterior is cumbersome. In particular, we have this annoying mu_dim_0. What is that? In our model, mu is a parameter array of size (2,), so it has one dimension. Because we have not labeled this dimension explicitly, it is automatically labeled for us in the inferenceData object: mu_dim_0. The two coordinates on this dimension are mu_dim_0=0 and mu_dim_0=1. Which of these two refers to the female mean and which refers to the male mean? Good question! It's not immediately clear. We would have to go back and see how we have handled our data. Looking back, we can check the categorical codes that we asked pandas to generate for our data['female'] variable.

In [20]:

sex_codes

Out[20]:

Index(['female', 'male'], dtype='object')

In [21]:

# first patient's index into the sex_codes map
sex_idx[0]

Out[21]:

0

So pandas has assigned the first patient an index of 0 and sex_codes[0]=='female', suggesting that the first patient is female. We can verify this by checking that the first patient's entry in the 'sex' column of our original data.

In [22]:

# first patient's value in the 'sex' column
df['sex'].iloc[0]

Out[22]:

'female'

So mu_dim_0=0 correponds to the females. Thus, mu_dim_0=1 must corresponds to the males. Got all that? No? It's not particularly easy. Let's see how labeled dimensions and coordinates can help.

New Method ¶

Let's begin by defining sex as a dimension and labeling it as such. We will simultaneously label the coordinates, or values, on this dimension.

In [23]:

coords = {'sex':['female','male']}

Easy enough. What do we do with this? We pass it along to our PyMC model, which will allow us to use it when defining parameter arrays (among other things). Let's see.

In [24]:

with pm.Model(coords=coords) as labeled_model:
    # hyper prior
    mu_hyper = pm.Normal('mu_hyper', mu=120, sigma=15)
    # per-group prior
    mu = pm.Normal('mu', mu=mu_hyper, sigma=15, dims='sex')
    # likelihood
    likelihood = pm.Normal('likelihood', 
                           mu=mu[sex_idx],
                           sigma=15, 
                           observed=df['bp'])

Notice that we are again asking for an array of parameters, mu. But instead of asking for shape=n_grps, we now ask for dims='sex'. This tells PyMC that we want an array of a normally-distributed parameters, one for each coordinate on our 'sex' dimension.

Just like in the previous example, we will index into this array using the sex_idx that pandas generated for us earlier. But how do the values in sex_idx correspond to the coordinates?

In [25]:

# a few patients' coordinate values
patient_ids = [0,75,150]
for idx in sex_idx[patient_ids]:
    print(coords['sex'][idx])

female
female
male

So we have verified that indicies in sex_idx map into our coordinates in the desired manner. Let's sample and see how our labeled dimensions and coordinates help us to unpack our posterior.

In [26]:

with labeled_model:
    labeled_idata = pm.sample()

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [mu_hyper, mu]

100.00% [4000/4000 00:03<00:00 Sampling 2 chains, 0 divergences]

Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 4 seconds.

In [27]:

az.summary(labeled_idata, kind="stats", round_to=2)

Out[27]:

	mean	sd	hdi_3%	hdi_97%
mu_hyper	124.82	8.72	108.21	140.74
mu[female]	125.02	1.48	122.23	127.67
mu[male]	129.88	1.51	127.25	132.76

We can already see that the multiple entries in our parameter array mu are now indexed by the meaningful labels, 'female' and 'male', so we don't have to wonder about which summary corresponds to which group. These labels are also useful when interrogating the InferenceData object itself.

In [28]:

# average over chains and draws
labeled_idata.posterior['mu'].mean(dim=['chain', 'draw'])

Out[28]:

<xarray.DataArray 'mu' (sex: 2)>
array([125.01960648, 129.87574546])
Coordinates:
  * sex      (sex) <U6 'female' 'male'

xarray.DataArray

'mu'

• sex: 2

• 125.0 129.9

  array([125.01960648, 129.87574546])

• Coordinates: (1)
  • sex
    (sex)
    <U6
    'female' 'male'

    array(['female', 'male'], dtype='<U6')

• Attributes: (0)

Alternatively, we can pull out the two parameters manually, referring to sex and each of the two coordinates ('female' and 'male') intead of the unlabeled dimension mu_dim_0:

In [29]:

labeled_idata.posterior['mu'].sel(sex='female').mean()

Out[29]:

<xarray.DataArray 'mu' ()>
array(125.01960648)
Coordinates:
    sex      <U6 'female'

xarray.DataArray

'mu'

• 125.0

  array(125.01960648)

• Coordinates: (1)
  • sex
    ()
    <U6
    'female'

    array('female', dtype='<U6')

• Attributes: (0)

In [30]:

labeled_idata.posterior['mu'].sel(sex='male').mean()

Out[30]:

<xarray.DataArray 'mu' ()>
array(129.87574546)
Coordinates:
    sex      <U6 'male'

xarray.DataArray

'mu'

• 129.9

  array(129.87574546)

• Coordinates: (1)
  • sex
    ()
    <U6
    'male'

    array('male', dtype='<U6')

• Attributes: (0)

And here we can compare the two posteriors:

In [31]:

( labeled_idata.posterior['mu'].sel(sex='female') <
  labeled_idata.posterior['mu'].sel(sex='male') ).mean()

Out[31]:

<xarray.DataArray 'mu' ()>
array(0.9885)

xarray.DataArray

'mu'

• 0.9885

  array(0.9885)

• Coordinates: (0)
• Attributes: (0)

So more than 99% of the samples are consistent with the conclusion that the females' blood pressures were less than the males'.

pm.Data ¶

Dimensions/coordinates can also be used when defining pm.Data objects, the standard method of making PyMC models aware of the data you are using (e.g., the data passed to any observed variables). First, let's look at how we might have incorporated data into a model in PyMC3.

In [32]:

with pm.Model() as model:

    # define our one parameter
    coef = pm.Normal('coef', mu=0, sigma=15)

    # our data is used here...
    mean = coef * df['age'].to_numpy()

    # ...and also here, as our observed data
    # likelihood
    likelihood = pm.Normal('likelihood', 
                           mu=mean,
                           sigma=15, 
                           observed=df['bp'].to_numpy())
    idata = pm.sample()

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [coef]

100.00% [4000/4000 00:04<00:00 Sampling 2 chains, 0 divergences]

Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 4 seconds.

Let's see what all in that resulting InferenceData object

In [33]:

idata.groups

Out[33]:

<bound method InferenceData.groups of Inference data with groups:
	> posterior
	> log_likelihood
	> sample_stats
	> observed_data>

Now let's see how we can use pm.Data and incorporate a data-relevant dimension.

In [34]:

coords = {'patient_id':df.index}
with pm.Model(coords=coords) as model:
    
    # here we define pm.Data objects so that PyMC is aware
    # of our model-relevant data
    age = pm.MutableData('age', df['age'], dims='patient_id')
    bloodpressure = pm.MutableData('bloodpressure', df['bp'], dims='patient_id')

    # define our one parameter
    coef = pm.Normal('coef', mu=0, sigma=15)

    # our data is still used here...
    mean = coef * age

    # ...and here, as our observed data
    # likelihood
    likelihood = pm.Normal('likelihood', 
                           mu=mean,
                           sigma=15, 
                           observed=bloodpressure)
    idata = pm.sample()

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [coef]

100.00% [4000/4000 00:02<00:00 Sampling 2 chains, 0 divergences]

Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 3 seconds.

So what's the big deal? Nothing seems to have changed. But when we look inside the InferenceData object, we can now see that the data we used to define the model is carried around with the posterior (and everything else).

In [35]:

idata.groups

Out[35]:

<bound method InferenceData.groups of Inference data with groups:
	> posterior
	> log_likelihood
	> sample_stats
	> observed_data
	> constant_data>

We can see that a new group, constant_data, is included in our InferenceData object. This group contains the data we used when defining the model. Let's take a look.

In [36]:

idata.constant_data

Out[36]:

<xarray.Dataset>
Dimensions:        (patient_id: 200)
Coordinates:
  * patient_id     (patient_id) int64 979 924 522 912 787 ... 783 523 896 968 98
Data variables:
    age            (patient_id) float64 53.9 50.6 57.15 ... 52.35 35.73 50.27
    bloodpressure  (patient_id) float64 128.8 123.9 121.7 ... 128.4 138.2 129.2
Attributes:
    created_at:                 2022-06-13T20:09:32.929479
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.0.0b6

xarray.Dataset

• Dimensions:
  • patient_id: 200
• Coordinates: (1)
  • patient_id
    (patient_id)
    int64
    979 924 522 912 ... 523 896 968 98

    array([979, 924, 522, 912, 787, 673, 348, 691, 888, 419, 738, 477, 265, 164,
           267,   8,  75, 331, 382, 518, 996, 439, 938, 204, 260, 837, 883, 622,
           278, 193,  34, 476, 576, 963,  87, 361, 825, 515, 406, 663,  62, 580,
           434, 886, 336, 385, 231, 662, 790, 615, 614, 283, 919, 618, 803, 981,
           112,  51, 232, 487, 574, 946, 683,  63, 524, 743, 855, 103, 933, 225,
           648, 934, 697, 185, 566, 906,  71, 653, 329, 667, 186, 944, 957, 728,
           150, 671, 442, 495, 334, 861, 664, 195, 737, 913, 864, 121, 755, 651,
           646, 786, 932, 609,   1, 814, 291, 467, 264,   0, 356, 116, 519, 858,
            96, 536, 224,  21, 333, 207, 926, 144, 998, 475,  90, 665, 548, 483,
           296, 474,  37, 301,  49, 725,  99, 340, 950, 857, 134, 596, 106, 339,
            91, 534, 937, 713, 604, 789, 159, 308, 453,   3, 377, 163, 623, 516,
           161, 802, 869, 853, 435, 530, 203, 970, 964, 826, 413, 958, 627, 172,
           464,  84, 879, 119, 723, 777, 372, 702, 369, 191, 292, 552, 849, 494,
           317, 657, 258, 894, 309,  73, 328, 599, 157, 412, 620, 804, 147, 783,
           523, 896, 968,  98])

• Data variables: (2)
  • age
    (patient_id)
    float64
    53.9 50.6 57.15 ... 35.73 50.27

    array([53.90147389, 50.5966426 , 57.14988141, 50.2636637 , 47.80266515,
           51.83350794, 31.24672234, 54.98711296, 56.79480919, 39.61600411,
           48.81551034, 64.1351038 , 51.6376774 , 39.35594828, 59.97290682,
           44.85177214, 50.07865286, 44.72934491, 53.61617882, 42.2609305 ,
           49.75215627, 52.70837572, 36.8520691 , 39.5804395 , 63.51155525,
           40.24011461, 36.80039478, 49.98453089, 39.38630756, 52.6962435 ,
           45.71167688, 56.0473454 , 56.68660792, 34.57290701, 63.97777459,
           50.84349487, 50.18964092, 56.23684139, 39.50047102, 59.44880155,
           43.72187322, 39.17526294, 49.64121212, 42.9349161 , 52.64278635,
           52.98075841, 57.57050989, 46.27913911, 55.19391263, 41.8245741 ,
           48.96178387, 51.12465199, 50.5185618 , 52.14909777, 45.77552833,
           67.69316582, 47.90974365, 51.91804225, 63.34916891, 51.37447252,
           58.15542603, 47.47390834, 55.8719311 , 58.1878177 , 48.22338614,
           42.71759911, 63.25569328, 44.31776113, 60.16767732, 47.47669244,
           49.58980222, 46.95369539, 53.65291143, 55.83435273, 49.02385696,
           42.64590559, 56.5046152 , 57.59629431, 53.5987158 , 43.63616707,
           47.4709282 , 41.84992167, 54.31387604, 43.94013865, 44.71878856,
           38.24305313, 48.59344889, 54.19160917, 52.95513202, 55.55731878,
           44.50433465, 48.21250651, 51.45017361, 50.22629885, 46.50297183,
           55.94119402, 51.79783903, 54.05728996, 40.71937154, 56.53514968,
           41.95110038, 53.96219511, 52.6921578 , 50.25040305, 53.04373928,
           43.85918536, 56.91038233, 36.24019744, 40.84130492, 50.32480411,
           49.78502035, 43.87729868, 56.84651469, 46.46802197, 62.03573191,
           42.68426188, 59.24131504, 35.90738116, 53.16264909, 46.42559554,
           51.65462117, 32.11666734, 54.73898318, 60.02288001, 53.66384316,
           49.80059937, 53.28568796, 44.4674601 , 42.05278314, 46.93854921,
           44.85836355, 44.00386124, 54.00833893, 51.17948845, 53.99430856,
           64.36697205, 41.44201968, 58.68506271, 48.42198285, 53.39474938,
           49.07427063, 47.66980567, 43.7230174 , 52.47178009, 47.26920475,
           47.21423685, 52.28388339, 44.42629236, 46.09326376, 63.9072557 ,
           55.68508072, 52.29819254, 59.5203019 , 36.79794798, 43.01702051,
           38.77371036, 53.84753156, 58.41005559, 57.02715948, 59.02096707,
           38.46087495, 63.04725782, 43.42904936, 54.19551527, 46.56700152,
           55.57232137, 54.30173552, 40.87469066, 45.64028985, 64.90694935,
           52.5770528 , 47.4319567 , 44.71029876, 45.50520595, 43.95913078,
           51.85076477, 34.98688584, 59.26860079, 42.68924166, 59.14932367,
           46.52091295, 60.86755313, 33.12574033, 56.99262012, 48.25237437,
           43.06701885, 57.26031334, 49.05622224, 55.87658069, 57.43002207,
           52.39629248, 45.46863058, 49.73872476, 52.9040998 , 48.15195917,
           43.88624456, 54.0149081 , 52.34962223, 35.73168522, 50.27173604])

  • bloodpressure
    (patient_id)
    float64
    128.8 123.9 121.7 ... 138.2 129.2

    array([128.75095086, 123.93588089, 121.68306616, 135.14428653,
           121.42032247, 129.54757211, 122.15769177, 130.87735338,
           132.9622286 , 126.75440004, 124.88848143, 123.50016818,
           123.60027639, 124.8534441 , 123.26702176, 119.59431226,
           125.22877493, 117.6709765 , 127.72790058, 123.74665199,
           116.50792255, 125.92949427, 129.434414  , 120.0438151 ,
           124.88822096, 126.18654044, 118.8094666 , 123.84500651,
           128.13447883, 136.2779946 , 129.78626533, 119.29053279,
           134.30357741, 126.22620528, 124.86955453, 121.95076967,
           126.02303483, 128.99933748, 130.56674379, 120.39147074,
           122.03288947, 127.34212932, 126.31282055, 125.94066459,
           134.6070909 , 119.79218296, 117.96745121, 128.93068417,
           128.95681347, 123.73199803, 121.88001495, 130.126373  ,
           117.88200922, 128.68011452, 130.3007806 , 123.99538336,
           131.42471748, 124.42113248, 126.27705445, 122.204213  ,
           118.70966982, 125.37992197, 114.461351  , 115.86766297,
           124.58927554, 119.40708549, 129.56842998, 125.77626421,
           125.22217706, 127.98706148, 128.40111957, 126.95620669,
           127.85543559, 119.20164118, 129.3147997 , 116.29135994,
           124.59007472, 124.98814673, 119.9810809 , 124.93157922,
    ...
           135.06030279, 139.43620358, 128.67722976, 133.0580461 ,
           122.63778413, 127.12214163, 124.03817418, 124.32561017,
           128.40873778, 132.55301991, 127.69117963, 134.04253702,
           130.18307922, 129.63935473, 130.33883827, 136.56473611,
           124.36061914, 129.32990715, 131.63713674, 133.970074  ,
           121.25162355, 130.56467731, 122.17722098, 132.86687551,
           131.17441753, 129.50375566, 131.42508367, 122.64899324,
           125.15076313, 135.58096411, 134.49519076, 124.08479858,
           130.83551936, 125.09129055, 127.5018994 , 133.36806244,
           135.36523989, 127.00015284, 132.18720193, 131.33266312,
           135.01897583, 134.97295621, 127.24913291, 122.81024246,
           124.10618208, 133.1631043 , 128.61997767, 127.19543064,
           127.51103751, 135.44940839, 129.42420603, 139.26585953,
           134.53108799, 122.9995672 , 128.8402919 , 127.86624344,
           126.16422532, 127.84773326, 129.89830775, 140.93947383,
           134.04561645, 134.49589489, 128.64972151, 129.58874389,
           125.04355121, 128.62567666, 130.04500277, 132.45087363,
           131.87192528, 129.20645838, 131.26494059, 122.46486604,
           137.94982293, 127.91380956, 126.11790406, 123.42997753,
           129.09118644, 128.4294642 , 138.23660367, 129.21571582])

• Attributes: (4)

  created_at :

  2022-06-13T20:09:32.929479

  arviz_version :

  0.12.1

  inference_library :

  pymc

  inference_library_version :

  4.0.0b6

The fact that this data was associated with the patient_id dimension when we created the data means that we can now use this dimension to index into the data.

In [37]:

idata.constant_data.sel(patient_id=979)

Out[37]:

<xarray.Dataset>
Dimensions:        ()
Coordinates:
    patient_id     int64 979
Data variables:
    age            float64 53.9
    bloodpressure  float64 128.8
Attributes:
    created_at:                 2022-06-13T20:09:32.929479
    arviz_version:              0.12.1
    inference_library:          pymc
    inference_library_version:  4.0.0b6

xarray.Dataset

• Dimensions:
• Coordinates: (1)
  • patient_id
    ()
    int64
    979

    array(979)

• Data variables: (2)
  • age
    ()
    float64
    53.9

    array(53.90147389)

  • bloodpressure
    ()
    float64
    128.8

    array(128.75095086)

• Attributes: (4)

  created_at :

  2022-06-13T20:09:32.929479

  arviz_version :

  0.12.1

  inference_library :

  pymc

  inference_library_version :

  4.0.0b6

Wrap-Up ¶

Hopefully this has given you a quick look into the advantages of using coordinates and dimensions and illustrated some basic, but illustrative scenarios. Head over to the PyMC Discourse if you have questions or need additional help!

In [38]:

%load_ext watermark
%watermark -p pymc,arviz,numpy,pandas

pymc  : 4.0.0b6
arviz : 0.12.1
numpy : 1.22.4
pandas: 1.4.2

Posted on June 4, 2022