Interpret Advanced Usage

The interpret module provides a set of helper functions to aid the user in more advanced and complex analysis not covered within the comparisons, predictions, and slopes functions. These helper functions are data_grid and select_draws. The data_grid can be used to create a pairwise grid of data points for the user to pass to model.predict. Subsequently, select_draws is used to select the draws from the posterior (or posterior predictive) group of the InferenceData object returned by the predict method that correspond to the data points that “produced” that draw.

With access to the appropriately indexed draws, and data used to generate those draws, it enables for more complex analysis such as cross-comparisons and the choice of which model parameter to compute a quantity of interest for; among others. Additionally, the user has more control over the data passed to model.predict. Below, it will be demonstrated how to use these helper functions. First, to reproduce the results from the standard interpret API, and then to compute cross-comparisons.

import warnings

import arviz as az
import numpy as np
import pandas as pd

import bambi as bmb

from bambi.interpret.helpers import data_grid, select_draws

warnings.simplefilter(action='ignore', category=FutureWarning)

Zero Inflated Poisson

We will adopt the zero inflated Poisson (ZIP) model from the comparisons documentation to demonstrate the helper functions introduced above.

The ZIP model will be used to predict how many fish are caught by visitors at a state park using survey data. Many visitors catch zero fish, either because they did not fish at all, or because they were unlucky. We would like to explicitly model this bimodal behavior (zero versus non-zero) using a Zero Inflated Poisson model, and to compare how different inputs of interest \(w\) and other covariate values \(c\) are associated with the number of fish caught. The dataset contains data on 250 groups that went to a state park to fish. Each group was questioned about how many fish they caught (count), how many children were in the group (child), how many people were in the group (persons), if they used a live bait and whether or not they brought a camper to the park (camper).

fish_data = pd.read_csv("https://stats.idre.ucla.edu/stat/data/fish.csv")
cols = ["count", "livebait", "camper", "persons", "child"]
fish_data = fish_data[cols]
fish_data["child"] = fish_data["child"].astype(np.int8)
fish_data["persons"] = fish_data["persons"].astype(np.int8)
fish_data["livebait"] = pd.Categorical(fish_data["livebait"])
fish_data["camper"] = pd.Categorical(fish_data["camper"])

fish_model = bmb.Model(
    "count ~ livebait + camper + persons + child", 
    fish_data, 
    family='zero_inflated_poisson'
)

fish_idata = fish_model.fit(random_seed=1234)

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [psi, Intercept, livebait, camper, persons, child]

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

Create a grid of data

data_grid allows you to create a pairwise grid, also known as a cross-join or cartesian product, of data using the covariates passed to the conditional and the optional variable parameter. Covariates not passed to conditional, but are terms in the Bambi model, are set to typical values (e.g., mean or mode). If you are coming from R, this function is partially inspired from the data_grid function in the R package modelr.

There are two ways to create a pairwise grid of data:

1. user-provided values are passed as a dictionary to conditional where the keys are the names of the covariates and the values are the values to use in the grid.
2. a list of covariates where the elements are the names of the covariates to use in the grid. As only the names of the covariates were passed, default values are computed to construct the grid.

Any unspecified covariates, i.e., covariates not passed to conditional but are terms in the Bambi model, are set to their “typical” values such as mean or mode depending on the data type of the covariate.

User-provided values

To construct a pairwise grid of data for specific covariate values, pass a dictionary to conditional. The values of the dictionary can be of type int, float, list, or np.ndarray.

conditional = {
    "camper": np.array([0, 1]),
    "persons": np.arange(1, 5, 1),
    "child": np.array([1, 2, 3]),
}
user_passed_grid = data_grid(fish_model, conditional)
user_passed_grid.query("camper == 0")

Default computed for unspecified variable: livebait

	camper	persons	child	livebait
0	0	1	1	1
1	0	1	2	1
2	0	1	3	1
3	0	2	1	1
4	0	2	2	1
5	0	2	3	1
6	0	3	1	1
7	0	3	2	1
8	0	3	3	1
9	0	4	1	1
10	0	4	2	1
11	0	4	3	1

Subsetting by camper = 0, it can be seen that a combination of all possible pairs of values from the dictionary (including the unspecified variable livebait) results in a dataframe containing every possible combination of values from the original sets. livebait has been set to 1 as this is the mode of the unspecified categorical variable.

Default values

Alternatively, a list of covariates can be passed to conditional where the elements are the names of the covariates to use in the grid. By doing this, you are telling interpret to compute default values for these covariates. The psuedocode below outlines the logic and functions used to compute these default values:

if is_numeric_dtype(x) or is_float_dtype(x):
    values = np.linspace(np.min(x), np.max(x), 50)

elif is_integer_dtype(x):
    values = np.quantile(x, np.linspace(0, 1, 5))

elif is_categorical_dtype(x) or is_string_dtype(x) or is_object_dtype(x):
    values = np.unique(x)

conditional = ["camper", "persons", "child"]
default_grid = data_grid(fish_model, conditional)

default_grid.shape, user_passed_grid.shape

Default computed for conditional variable: camper, persons, child
Default computed for unspecified variable: livebait

((32, 4), (24, 4))

Notice how the resulting length is different between the user passed and default grid. This is due to the fact that values for child range from 0 to 3 for the default grid.

default_grid.query("camper == 0")

	camper	persons	child	livebait
0	0	1	0	1
1	0	1	1	1
2	0	1	2	1
3	0	1	3	1
4	0	2	0	1
5	0	2	1	1
6	0	2	2	1
7	0	2	3	1
8	0	3	0	1
9	0	3	1	1
10	0	3	2	1
11	0	3	3	1
12	0	4	0	1
13	0	4	1	1
14	0	4	2	1
15	0	4	3	1

Compute comparisons

To use data_grid to help generate data in computing comparisons or slopes, additional data is passed to the optional variable parameter. The name variable is an abstraction for the comparisons parameter contrast and slopes parameter wrt. If you have used any of the interpret functions, these parameter names should be familiar and the use of data_grid should be analogous to comparisons, predictions, and slopes.

variable can also be passed user-provided data (as a dictionary), or a string indicating the name of the covariate of interest. If the latter, a default value will be computed. Additionally, if an argument is passed for variable, then the effect_type needs to be passed. This is because for comparisons and slopes an epsilon value eps needs to be determined to compute the centered and finite difference, respectively. You can also pass a value for eps as a kwarg.

conditional = {
    "camper": np.array([0, 1]),
    "persons": np.arange(1, 5, 1),
    "child": np.array([1, 2, 3, 4])
}
variable = "livebait"

grid = data_grid(fish_model, conditional, variable, effect_type="comparisons")

Default computed for contrast variable: livebait

idata_grid = fish_model.predict(fish_idata, data=grid, inplace=False)

Select draws conditional on data

The second helper function to aid in more advanced analysis is select_draws. This is a function that selects the posterior or posterior predictive draws from the ArviZ InferenceData object returned by model.predict given a conditional dictionary. The conditional dictionary represents the values that correspond to that draw.

For example, if we wanted to select posterior draws where livebait = [0, 1], then all we need to do is pass a dictionary where the key is the name of the covariate and the value is the value that we want to condition on (or select). The resulting InferenceData object will contain the draws that correspond to the data points where livebait = [0, 1]. Additionally, you must pass the InferenceData object returned by model.predict, the data used to generate the predictions, and the name of the data variable data_var you would like to select from the InferenceData posterior group. If you specified to return the posterior predictive samples by passing model.predict(..., kind="pps"), you can use this group instead of the posterior group by passing group="posterior_predictive".

Below, it is demonstrated how to compute comparisons for count_mean for the contrast livebait = [0, 1] using the posterior draws.

idata_grid

arviz.InferenceData

• posterior

  <xarray.Dataset> Size: 2MB
  Dimensions:       (chain: 4, draw: 1000, livebait_dim: 1, camper_dim: 1,
                     __obs__: 64)
  Coordinates:
    * chain         (chain) int64 32B 0 1 2 3
    * draw          (draw) int64 8kB 0 1 2 3 4 5 6 ... 993 994 995 996 997 998 999
    * livebait_dim  (livebait_dim) <U1 4B '1'
    * camper_dim    (camper_dim) <U1 4B '1'
    * __obs__       (__obs__) int64 512B 0 1 2 3 4 5 6 7 ... 57 58 59 60 61 62 63
  Data variables:
      psi           (chain, draw) float64 32kB 0.4751 0.6612 ... 0.5565 0.5387
      Intercept     (chain, draw) float64 32kB -2.655 -2.48 ... -2.59 -2.517
      livebait      (chain, draw, livebait_dim) float64 32kB 1.936 1.784 ... 1.822
      camper        (chain, draw, camper_dim) float64 32kB 0.6152 ... 0.5638
      persons       (chain, draw) float64 32kB 0.864 0.8756 ... 0.8885 0.8896
      child         (chain, draw) float64 32kB -1.399 -1.272 ... -1.474 -1.467
      mu            (chain, draw, __obs__) float64 2MB 0.04118 0.2855 ... 0.08692
  Attributes:
      created_at:                  2025-09-28T17:53:04.333667+00:00
      arviz_version:               0.22.0
      inference_library:           pymc
      inference_library_version:   3.9.2+2907.g7a3db78e6
      sampling_time:               5.744408369064331
      tuning_steps:                1000
      modeling_interface:          bambi
      modeling_interface_version:  0.14.1.dev57+g7b2bb342c.d20250928

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • livebait_dim: 1
    • camper_dim: 1
    • __obs__: 64
  • Coordinates: (5)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

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

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

    • livebait_dim
      (livebait_dim)
      <U1
      '1'

      array(['1'], dtype='<U1')

    • camper_dim
      (camper_dim)
      <U1
      '1'

      array(['1'], dtype='<U1')

    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 6 ... 58 59 60 61 62 63

      array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
             36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
             54, 55, 56, 57, 58, 59, 60, 61, 62, 63])

  • Data variables: (7)
    • psi
      (chain, draw)
      float64
      0.4751 0.6612 ... 0.5565 0.5387

      array([[0.47505204, 0.66118607, 0.57571122, ..., 0.58190278, 0.62567799,
              0.53047598],
             [0.66252093, 0.5891482 , 0.55873105, ..., 0.57024777, 0.62302212,
              0.6379181 ],
             [0.60942378, 0.62435478, 0.63496683, ..., 0.64415829, 0.64225026,
              0.60127422],
             [0.66630496, 0.46980712, 0.45517287, ..., 0.61960178, 0.55646107,
              0.53872486]], shape=(4, 1000))

    • Intercept
      (chain, draw)
      float64
      -2.655 -2.48 ... -2.59 -2.517

      array([[-2.65453728, -2.48014747, -2.43947212, ..., -2.32744299,
              -2.21885658, -2.40076347],
             [-2.49592977, -2.96189017, -2.45732213, ..., -2.42789096,
              -2.57240175, -2.41380188],
             [-1.94887941, -2.81836131, -2.86993076, ..., -2.54960719,
              -2.28847446, -2.67300306],
             [-2.74252748, -1.90686183, -1.91953106, ..., -2.97620999,
              -2.58952781, -2.51701433]], shape=(4, 1000))

    • livebait
      (chain, draw, livebait_dim)
      float64
      1.936 1.784 1.646 ... 1.772 1.822

      array([[[1.93634598],
              [1.78351511],
              [1.64585371],
              ...,
              [1.68453505],
              [1.52531296],
              [1.65994633]],
      
             [[1.75892723],
              [2.0332964 ],
              [1.78381676],
              ...,
              [1.59273165],
              [1.98520374],
              [1.35598811]],
      
             [[1.38668651],
              [1.77360228],
              [1.90977193],
              ...,
              [1.39062137],
              [1.34452461],
              [1.9063861 ]],
      
             [[2.15021993],
              [1.47721892],
              [1.49108566],
              ...,
              [2.08292783],
              [1.77188371],
              [1.82160551]]], shape=(4, 1000, 1))

    • camper
      (chain, draw, camper_dim)
      float64
      0.6152 0.5502 ... 0.6415 0.5638

      array([[[0.61516609],
              [0.55024909],
              [0.63472053],
              ...,
              [0.59075578],
              [0.68666371],
              [0.66837279]],
      
             [[0.6776875 ],
              [0.65136096],
              [0.64614926],
              ...,
              [0.6957036 ],
              [0.62586061],
              [0.74264992]],
      
             [[0.58652207],
              [0.61359844],
              [0.67168102],
              ...,
              [0.87689133],
              [0.81821063],
              [0.64885935]],
      
             [[0.6753028 ],
              [0.41250836],
              [0.38154009],
              ...,
              [0.6717866 ],
              [0.64152592],
              [0.56377294]]], shape=(4, 1000, 1))

    • persons
      (chain, draw)
      float64
      0.864 0.8756 ... 0.8885 0.8896

      array([[0.86396867, 0.87555097, 0.87076137, ..., 0.82525173, 0.8508994 ,
              0.84094239],
             [0.85150775, 0.89639939, 0.8186814 , ..., 0.86780995, 0.81480082,
              0.9294559 ],
             [0.81012519, 0.93785595, 0.89146359, ..., 0.91683796, 0.8736617 ,
              0.85305218],
             [0.82237691, 0.81521132, 0.83520346, ..., 0.88849766, 0.88849482,
              0.8895511 ]], shape=(4, 1000))

    • child
      (chain, draw)
      float64
      -1.399 -1.272 ... -1.474 -1.467

      array([[-1.39918978, -1.27176458, -1.42188044, ..., -1.28764853,
              -1.38335914, -1.22887547],
             [-1.50774849, -1.43886432, -1.22879453, ..., -1.27121457,
              -1.4074095 , -1.42335677],
             [-1.40907997, -1.2767149 , -1.23191848, ..., -1.40730834,
              -1.34722365, -1.38066414],
             [-1.41255481, -1.26441712, -1.29592335, ..., -1.2514328 ,
              -1.4742868 , -1.46734638]], shape=(4, 1000))

    • mu
      (chain, draw, __obs__)
      float64
      0.04118 0.2855 ... 0.01406 0.08692

      array([[[0.04118182, 0.28552879, 0.01016354, ..., 0.42967771,
               0.0152946 , 0.1060431 ],
              [0.0563394 , 0.33526099, 0.01579399, ..., 0.63161635,
               0.02975516, 0.17706512],
              [0.05025772, 0.26060811, 0.01212517, ..., 0.39004478,
               0.0181474 , 0.09410217],
              ...,
              [0.06143106, 0.33110997, 0.01694999, ..., 0.54113266,
               0.02770134, 0.14930869],
              [0.06384377, 0.29346311, 0.01600786, ..., 0.4707822 ,
               0.02568028, 0.11804152],
              [0.06150133, 0.32343723, 0.01799661, ..., 0.67347714,
               0.03747343, 0.19707386]],
      
             [[0.04275922, 0.24826879, 0.00946723, ..., 0.308347  ,
               0.01175819, 0.06827047],
              [0.03006616, 0.22968219, 0.0071316 , ..., 0.36486487,
               0.01132901, 0.08654485],
              [0.05684453, 0.33836892, 0.01663528, ..., 0.6446717 ,
               0.03169408, 0.18866004],
      ...
              [0.04783118, 0.19215437, 0.01170914, ..., 0.43315152,
               0.02639456, 0.10603606],
              [0.06316301, 0.2423162 , 0.01641992, ..., 0.51029649,
               0.0345789 , 0.13265719],
              [0.04073714, 0.27410914, 0.01024179, ..., 0.42846709,
               0.0160092 , 0.10772156]],
      
             [[0.0356964 , 0.30651596, 0.0086928 , ..., 0.42098354,
               0.0119391 , 0.10251806],
              [0.09479225, 0.41526072, 0.0267697 , ..., 0.57721489,
               0.03721004, 0.16300771],
              [0.09252735, 0.41099867, 0.02531966, ..., 0.55219086,
               0.03401783, 0.15110431],
              ...,
              [0.03546726, 0.28472903, 0.01014699, ..., 0.65586302,
               0.02337323, 0.18763886],
              [0.04178074, 0.24575109, 0.00956536, ..., 0.35169207,
               0.0136889 , 0.08051707],
              [0.04528363, 0.27993325, 0.01043953, ..., 0.37701488,
               0.01405998, 0.08691566]]], shape=(4, 1000, 64))

  • Indexes: (5)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='draw', length=1000))

    • livebait_dim
      PandasIndex

      PandasIndex(Index(['1'], dtype='object', name='livebait_dim'))

    • camper_dim
      PandasIndex

      PandasIndex(Index(['1'], dtype='object', name='camper_dim'))

    • __obs__
      PandasIndex

      PandasIndex(Index([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
             36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
             54, 55, 56, 57, 58, 59, 60, 61, 62, 63],
            dtype='int64', name='__obs__'))

  • Attributes: (8)

    created_at :

    2025-09-28T17:53:04.333667+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    3.9.2+2907.g7a3db78e6

    sampling_time :

    5.744408369064331

    tuning_steps :

    1000

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.1.dev57+g7b2bb342c.d20250928

  
  
• sample_stats

  <xarray.Dataset> Size: 528kB
  Dimensions:                (chain: 4, draw: 1000)
  Coordinates:
    * chain                  (chain) int64 32B 0 1 2 3
    * draw                   (draw) int64 8kB 0 1 2 3 4 5 ... 995 996 997 998 999
  Data variables: (12/18)
      process_time_diff      (chain, draw) float64 32kB 0.00177 ... 0.001431
      index_in_trajectory    (chain, draw) int64 32kB -2 4 3 6 -6 ... 5 -4 5 3 13
      max_energy_error       (chain, draw) float64 32kB 2.184 1.865 ... 1.111
      perf_counter_start     (chain, draw) float64 32kB 2.045e+05 ... 2.045e+05
      reached_max_treedepth  (chain, draw) bool 4kB False False ... False False
      energy_error           (chain, draw) float64 32kB -0.1267 0.02556 ... 1.111
      ...                     ...
      n_steps                (chain, draw) float64 32kB 7.0 7.0 7.0 ... 7.0 15.0
      step_size_bar          (chain, draw) float64 32kB 0.463 0.463 ... 0.4197
      energy                 (chain, draw) float64 32kB 762.8 759.5 ... 755.8
      diverging              (chain, draw) bool 4kB False False ... False False
      perf_counter_diff      (chain, draw) float64 32kB 0.001946 ... 0.001552
      smallest_eigval        (chain, draw) float64 32kB nan nan nan ... nan nan
  Attributes:
      created_at:                  2025-09-28T17:53:04.353984+00:00
      arviz_version:               0.22.0
      inference_library:           pymc
      inference_library_version:   3.9.2+2907.g7a3db78e6
      sampling_time:               5.744408369064331
      tuning_steps:                1000
      modeling_interface:          bambi
      modeling_interface_version:  0.14.1.dev57+g7b2bb342c.d20250928

  xarray.Dataset

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

      array([0, 1, 2, 3])

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

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

  • Data variables: (18)
    • process_time_diff
      (chain, draw)
      float64
      0.00177 0.002007 ... 0.001431

      array([[0.0017698, 0.0020073, 0.0019871, ..., 0.0013771, 0.0010966,
              0.0008205],
             [0.0022362, 0.0010605, 0.002827 , ..., 0.0027154, 0.001546 ,
              0.0015522],
             [0.0019279, 0.0019871, 0.0037933, ..., 0.0009135, 0.0016638,
              0.0007657],
             [0.0019091, 0.0032898, 0.0020698, ..., 0.0008818, 0.0007678,
              0.0014314]], shape=(4, 1000))

    • index_in_trajectory
      (chain, draw)
      int64
      -2 4 3 6 -6 -2 3 ... 4 5 -4 5 3 13

      array([[-2,  4,  3, ...,  5, -2,  6],
             [-9, -3,  4, ..., -4,  5,  6],
             [-3, -3, 12, ..., -4, 15, -5],
             [13, -5, -1, ...,  5,  3, 13]], shape=(4, 1000))

    • max_energy_error
      (chain, draw)
      float64
      2.184 1.865 ... 0.07214 1.111

      array([[ 2.18413835,  1.86501249, -0.34778222, ...,  0.38297837,
              -0.55413146, -0.37703349],
             [-0.80521076,  0.66475682, -0.51290278, ...,  0.70322474,
               0.13051808,  0.11547728],
             [-0.25059241,  0.40179839,  0.51608246, ...,  1.00844166,
               0.99224621,  0.49619473],
             [ 0.0922762 , -0.26885173, -0.28300683, ..., -0.11904773,
               0.072136  ,  1.11066162]], shape=(4, 1000))

    • perf_counter_start
      (chain, draw)
      float64
      2.045e+05 2.045e+05 ... 2.045e+05

      array([[204453.01892517, 204453.02114937, 204453.02358408, ...,
              204454.97468938, 204454.97644201, 204454.97781465],
             [204452.79561392, 204452.79831157, 204452.79963924, ...,
              204454.88499231, 204454.88814699, 204454.89010748],
             [204453.021237  , 204453.02358418, 204453.02607497, ...,
              204455.15433299, 204455.15548246, 204455.15747582],
             [204453.08017075, 204453.0824367 , 204453.09739496, ...,
              204455.21121396, 204455.21231873, 204455.21329086]],
            shape=(4, 1000))

    • reached_max_treedepth
      (chain, draw)
      bool
      False False False ... False False

      array([[False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False]], shape=(4, 1000))

    • energy_error
      (chain, draw)
      float64
      -0.1267 0.02556 ... -0.008317 1.111

      array([[-0.12670574,  0.02555692, -0.34337592, ...,  0.37143005,
               0.06623665, -0.34916326],
             [-0.75130688,  0.6325036 ,  0.00702398, ..., -0.07080415,
               0.11888384,  0.0096263 ],
             [-0.07545402, -0.04947368,  0.43992906, ..., -0.12149239,
              -0.02681514,  0.23165522],
             [-0.03899448, -0.0707252 ,  0.04399099, ..., -0.05644536,
              -0.00831714,  1.11066162]], shape=(4, 1000))

    • lp
      (chain, draw)
      float64
      -755.7 -753.7 ... -751.4 -754.4

      array([[-755.67091538, -753.67155712, -750.38917968, ..., -751.76303739,
              -751.72505089, -752.27790912],
             [-750.97967656, -752.7216474 , -752.68607572, ..., -750.70638629,
              -751.43232804, -751.76796888],
             [-752.6836068 , -752.28115879, -753.22404726, ..., -753.01476202,
              -752.20287966, -750.72905087],
             [-753.62922258, -758.97727593, -760.86075571, ..., -752.10404338,
              -751.41332297, -754.38218058]], shape=(4, 1000))

    • divergences
      (chain, draw)
      int64
      0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0

      array([[0, 0, 0, ..., 0, 0, 0],
             [0, 0, 0, ..., 0, 0, 0],
             [0, 0, 0, ..., 0, 0, 0],
             [0, 0, 0, ..., 0, 0, 0]], shape=(4, 1000))

    • largest_eigval
      (chain, draw)
      float64
      nan nan nan nan ... nan nan nan nan

      array([[nan, nan, nan, ..., nan, nan, nan],
             [nan, nan, nan, ..., nan, nan, nan],
             [nan, nan, nan, ..., nan, nan, nan],
             [nan, nan, nan, ..., nan, nan, nan]], shape=(4, 1000))

    • tree_depth
      (chain, draw)
      int64
      3 3 3 3 3 3 4 3 ... 4 3 3 3 3 3 3 4

      array([[3, 3, 3, ..., 3, 3, 3],
             [4, 3, 4, ..., 4, 3, 3],
             [3, 3, 4, ..., 3, 4, 3],
             [4, 4, 3, ..., 3, 3, 4]], shape=(4, 1000))

    • step_size
      (chain, draw)
      float64
      0.3998 0.3998 ... 0.4525 0.4525

      array([[0.39984063, 0.39984063, 0.39984063, ..., 0.39984063, 0.39984063,
              0.39984063],
             [0.43855996, 0.43855996, 0.43855996, ..., 0.43855996, 0.43855996,
              0.43855996],
             [0.47311063, 0.47311063, 0.47311063, ..., 0.47311063, 0.47311063,
              0.47311063],
             [0.45254452, 0.45254452, 0.45254452, ..., 0.45254452, 0.45254452,
              0.45254452]], shape=(4, 1000))

    • acceptance_rate
      (chain, draw)
      float64
      0.4448 0.4935 1.0 ... 0.9765 0.6551

      array([[0.44484216, 0.49353371, 1.        , ..., 0.84425068, 0.98423452,
              0.97354558],
             [0.98381259, 0.71698876, 0.97393013, ..., 0.74891399, 0.91987007,
              0.96645363],
             [0.99044968, 0.87369173, 0.84641834, ..., 0.68483477, 0.66432019,
              0.79750891],
             [0.97159345, 0.95442976, 0.98573047, ..., 1.        , 0.97652523,
              0.65514232]], shape=(4, 1000))

    • n_steps
      (chain, draw)
      float64
      7.0 7.0 7.0 7.0 ... 7.0 7.0 15.0

      array([[ 7.,  7.,  7., ...,  7.,  7.,  7.],
             [15.,  7., 15., ..., 15.,  7.,  7.],
             [ 7.,  7., 15., ...,  7., 15.,  7.],
             [15., 15.,  7., ...,  7.,  7., 15.]], shape=(4, 1000))

    • step_size_bar
      (chain, draw)
      float64
      0.463 0.463 0.463 ... 0.4197 0.4197

      array([[0.4630148 , 0.4630148 , 0.4630148 , ..., 0.4630148 , 0.4630148 ,
              0.4630148 ],
             [0.45538173, 0.45538173, 0.45538173, ..., 0.45538173, 0.45538173,
              0.45538173],
             [0.44061191, 0.44061191, 0.44061191, ..., 0.44061191, 0.44061191,
              0.44061191],
             [0.41966235, 0.41966235, 0.41966235, ..., 0.41966235, 0.41966235,
              0.41966235]], shape=(4, 1000))

    • energy
      (chain, draw)
      float64
      762.8 759.5 754.2 ... 754.0 755.8

      array([[762.79877504, 759.54625228, 754.17185912, ..., 752.71000354,
              753.45414896, 756.11838657],
             [752.63566584, 753.81393693, 757.09634008, ..., 753.19788792,
              752.58634373, 754.93827441],
             [756.43554461, 756.9626317 , 756.31684022, ..., 756.27732139,
              754.66357716, 753.74312919],
             [758.82125131, 763.33711627, 761.61151457, ..., 754.68409891,
              754.01154025, 755.77150518]], shape=(4, 1000))

    • diverging
      (chain, draw)
      bool
      False False False ... False False

      array([[False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False]], shape=(4, 1000))

    • perf_counter_diff
      (chain, draw)
      float64
      0.001946 0.002206 ... 0.001552

      array([[0.00194558, 0.00220633, 0.00218391, ..., 0.00149219, 0.00118894,
              0.0008902 ],
             [0.0024579 , 0.00116698, 0.00310552, ..., 0.00294279, 0.00167569,
              0.0016828 ],
             [0.00211933, 0.00218255, 0.00417005, ..., 0.00099017, 0.00180081,
              0.00083081],
             [0.00209892, 0.00361377, 0.0022756 , ..., 0.00095534, 0.00083204,
              0.00155204]], shape=(4, 1000))

    • smallest_eigval
      (chain, draw)
      float64
      nan nan nan nan ... nan nan nan nan

      array([[nan, nan, nan, ..., nan, nan, nan],
             [nan, nan, nan, ..., nan, nan, nan],
             [nan, nan, nan, ..., nan, nan, nan],
             [nan, nan, nan, ..., nan, nan, nan]], shape=(4, 1000))

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='draw', length=1000))

  • Attributes: (8)

    created_at :

    2025-09-28T17:53:04.353984+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    3.9.2+2907.g7a3db78e6

    sampling_time :

    5.744408369064331

    tuning_steps :

    1000

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.1.dev57+g7b2bb342c.d20250928

  
  
• observed_data

  <xarray.Dataset> Size: 4kB
  Dimensions:  (__obs__: 250)
  Coordinates:
    * __obs__  (__obs__) int64 2kB 0 1 2 3 4 5 6 7 ... 243 244 245 246 247 248 249
  Data variables:
      count    (__obs__) int64 2kB 0 0 0 0 1 0 0 0 0 1 0 ... 4 1 1 0 1 0 0 0 0 0 0
  Attributes:
      created_at:                  2025-09-28T17:53:04.359378+00:00
      arviz_version:               0.22.0
      inference_library:           pymc
      inference_library_version:   3.9.2+2907.g7a3db78e6
      modeling_interface:          bambi
      modeling_interface_version:  0.14.1.dev57+g7b2bb342c.d20250928

  xarray.Dataset

  • Dimensions:
    • __obs__: 250
  • Coordinates: (1)
    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 245 246 247 248 249

      array([  0,   1,   2, ..., 247, 248, 249], shape=(250,))

  • Data variables: (1)
    • count
      (__obs__)
      int64
      0 0 0 0 1 0 0 0 ... 0 1 0 0 0 0 0 0

      array([  0,   0,   0,   0,   1,   0,   0,   0,   0,   1,   0,   0,   1,
               2,   0,   1,   0,   0,   1,   0,   1,   5,   0,   3,  30,   0,
              13,   0,   0,   0,   0,   0,  11,   5,   0,   1,   1,   7,   0,
              14,   0,  32,   0,   1,   0,   0,   0,   1,   5,   0,   1,   0,
              22,   0,  15,   0,   0,   0,   5,   4,   2,   0,   2,  32,   0,
               0,   1,   0,   0,   0,   7,   0,   0,   0,   0,   0,   0,   0,
               0,   2,   3,   1,   5,   0,   2,   1,   0,   1, 149,   0,   1,
               0,   0,   1,   0,   0,   0,   2,   2,  29,   3,   0,   0,   5,
               0,   0,   0,   0,   0,   1,   7,   1,   0,   2,   0,   2,   0,
               0,   0,   1,   0,   0,   0,   0,   0,   3,   4,   3,   3,   8,
               2,   1,   6,   0,   0,   5,   3,  31,   0,   2,   0,   0,   0,
               0,   0,   0,   6,   9,   0,   0,   0,   0,   0,   2,  15,   1,
               2,   3,   0,  65,   5,   0,   0,   0,   0,   1,   8,   0,   0,
               0,   2,   4,   5,   9,   0,   0,   0,   0,  21,   0,   6,   0,
               0,   0,   0,  16,   0,   0,   4,   2,  10,   0,   0,   0,   2,
               1,   3,   0,   0,  21,   0,   0,   2,   0,   3,   0,  38,   0,
               0,   0,   1,   3,   0,   1,   0,   0,   0,   0,   5,   0,   0,
               2,   0,   0,   0,   1,   4,   0,   0,   2,   3,   0,   0,   0,
               0,   1,   2,   0,   6,   4,   1,   1,   0,   1,   0,   0,   0,
               0,   0,   0])

  • Indexes: (1)
    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             240, 241, 242, 243, 244, 245, 246, 247, 248, 249],
            dtype='int64', name='__obs__', length=250))

  • Attributes: (6)

    created_at :

    2025-09-28T17:53:04.359378+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    3.9.2+2907.g7a3db78e6

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.1.dev57+g7b2bb342c.d20250928

  
  

draw_1 = select_draws(idata_grid, grid, {"livebait": 0}, "mu")

draw_1 = select_draws(idata_grid, grid, {"livebait": 0}, "mu")
draw_2 = select_draws(idata_grid, grid, {"livebait": 1}, "mu")

comparison_mean = (draw_2 - draw_1).mean(("chain", "draw"))
comparison_hdi = az.hdi(draw_2 - draw_1)

comparison_df = pd.DataFrame(
    {
        "mean": comparison_mean.values,
        "hdi_low": comparison_hdi.sel(hdi="lower")["mu"].values,
        "hdi_high": comparison_hdi.sel(hdi="higher")["mu"].values,
    }
)
comparison_df.head(10)

	mean	hdi_low	hdi_high
0	0.213995	0.143976	0.286619
1	0.053701	0.031074	0.079068
2	0.013598	0.006305	0.022216
3	0.003475	0.001268	0.006288
4	0.511683	0.373739	0.667900
5	0.128343	0.078841	0.183933
6	0.032483	0.015540	0.051750
7	0.008296	0.003144	0.014865
8	1.225939	0.921290	1.553605
9	0.307349	0.194784	0.436108

We can compare this comparison with bmb.interpret.comparisons.

summary_df = bmb.interpret.comparisons(
    fish_model,
    fish_idata,
    contrast={"livebait": [0, 1]},
    conditional=conditional
)
summary_df.head(10)

	term	estimate_type	value	camper	persons	child	estimate	lower_3.0%	upper_97.0%
0	livebait	diff	(0, 1)	0	1	1	0.213995	0.143976	0.286619
1	livebait	diff	(0, 1)	0	1	2	0.053701	0.031074	0.079068
2	livebait	diff	(0, 1)	0	1	3	0.013598	0.006305	0.022216
3	livebait	diff	(0, 1)	0	1	4	0.003475	0.001268	0.006288
4	livebait	diff	(0, 1)	0	2	1	0.511683	0.373739	0.667900
5	livebait	diff	(0, 1)	0	2	2	0.128343	0.078841	0.183933
6	livebait	diff	(0, 1)	0	2	3	0.032483	0.015540	0.051750
7	livebait	diff	(0, 1)	0	2	4	0.008296	0.003144	0.014865
8	livebait	diff	(0, 1)	0	3	1	1.225939	0.921290	1.553605
9	livebait	diff	(0, 1)	0	3	2	0.307349	0.194784	0.436108

Albeit the other information in the summary_df, the columns estimate, lower_3.0%, upper_97.0% are identical.

Cross comparisons

Computing a cross-comparison is useful for when we want to compare contrasts when two (or more) predictors change at the same time. Cross-comparisons are currently not supported in the comparisons function, but we can use select_draws to compute them. For example, imagine we are interested in computing the cross-comparison between the two rows below.

summary_df.iloc[:2]

	term	estimate_type	value	camper	persons	child	estimate	lower_3.0%	upper_97.0%
0	livebait	diff	(0, 1)	0	1	1	0.213995	0.143976	0.286619
1	livebait	diff	(0, 1)	0	1	2	0.053701	0.031074	0.079068

The cross-comparison amounts to first computing the comparison for row 0, given below, and can be verified by looking at the estimate in summary_df.

cond_10 = {
    "camper": 0,
    "persons": 1,
    "child": 1,
    "livebait": 0 
}

cond_11 = {
    "camper": 0,
    "persons": 1,
    "child": 1,
    "livebait": 1
}

draws_10 = select_draws(idata_grid, grid, cond_10, "mu")
draws_11 = select_draws(idata_grid, grid, cond_11, "mu")

(draws_11 - draws_10).mean(("chain", "draw")).item()

0.21399545475519965

Next, we need to compute the comparison for row 1.

cond_20 = {
    "camper": 0,
    "persons": 1,
    "child": 2,
    "livebait": 0
}

cond_21 = {
    "camper": 0,
    "persons": 1,
    "child": 2,
    "livebait": 1
}

draws_20 = select_draws(idata_grid, grid, cond_20, "mu")
draws_21 = select_draws(idata_grid, grid, cond_21, "mu")

(draws_21 - draws_20).mean(("chain", "draw")).item()

0.053701205869420204

Next, we compute the “first level” comparisons (diff_1 and diff_2). Subsequently, we compute the difference between these two differences to obtain the cross-comparison.

diff_1 = (draws_11 - draws_10)
diff_2 = (draws_21 - draws_20)

cross_comparison = (diff_2 - diff_1).mean(("chain", "draw")).item()
cross_comparison

-0.16029424888577945

To verify this is correct, we can check by performing the cross-comparison directly on the summary_df.

summary_df.iloc[1]["estimate"] - summary_df.iloc[0]["estimate"]

np.float64(-0.16029424888577945)

Summary

In this notebook, the interpret helper functions data_grid and select_draws were introduced and it was demonstrated how they can be used to compute pairwise grids of data and cross-comparisons. With these functions, it is left to the user to generate their grids of data and quantities of interest allowing for more flexibility and control over the type of data passed to model.predict and the quantities of interest computed.

%load_ext watermark
%watermark -n -u -v -iv -w

Last updated: Sun Sep 28 2025

Python implementation: CPython
Python version       : 3.13.7
IPython version      : 9.4.0

numpy : 2.3.3
arviz : 0.22.0
bambi : 0.14.1.dev57+g7b2bb342c.d20250928
pandas: 2.3.2

Watermark: 2.5.0