ArviZ Quickstart

import arviz as az
import numpy as np

ArviZ style sheets

# ArviZ ships with style sheets!
az.style.use("arviz-darkgrid")

Feel free to check the examples of style sheets here.

Get started with plotting

ArviZ is designed to be used with libraries like PyStan and PyMC3, but works fine with raw NumPy arrays.

az.plot_posterior(np.random.randn(100_000));

Plotting a dictionary of arrays, ArviZ will interpret each key as the name of a different random variable. Each row of an array is treated as an independent series of draws from the variable, called a chain. Below, we have 10 chains of 50 draws, each for four different distributions.

size = (10, 50)
az.plot_forest(
    {
        "normal": np.random.randn(*size),
        "gumbel": np.random.gumbel(size=size),
        "student t": np.random.standard_t(df=6, size=size),
        "exponential": np.random.exponential(size=size),
    }
);

ArviZ rcParams

You may have noticed that for both plot_posterior() and plot_forest(), the Highest Density Interval (HDI) is 94%, which you may find weird at first. This particular value is a friendly reminder of the arbitrary nature of choosing any single value without further justification, including common values like 95%, 50% and even our own default, 94%. ArviZ includes default values for a few parameters, you can access them with az.rcParams. To change the default HDI value to let’s say 90% you can do:

az.rcParams["stats.hdi_prob"] = 0.90

import pymc3 as pm

J = 8
y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
schools = np.array(
    [
        "Choate",
        "Deerfield",
        "Phillips Andover",
        "Phillips Exeter",
        "Hotchkiss",
        "Lawrenceville",
        "St. Paul's",
        "Mt. Hermon",
    ]
)

with pm.Model() as centered_eight:
    mu = pm.Normal("mu", mu=0, sd=5)
    tau = pm.HalfCauchy("tau", beta=5)
    theta = pm.Normal("theta", mu=mu, sd=tau, shape=J)
    obs = pm.Normal("obs", mu=theta, sd=sigma, observed=y)

    # This pattern is useful in PyMC3
    prior = pm.sample_prior_predictive()
    centered_eight_trace = pm.sample(return_inferencedata=False)
    posterior_predictive = pm.sample_posterior_predictive(centered_eight_trace)

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

100.00% [8000/8000 00:03<00:00 Sampling 4 chains, 192 divergences]

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 14 seconds.
There were 145 divergences after tuning. Increase `target_accept` or reparameterize.
The acceptance probability does not match the target. It is 0.44992731488390336, but should be close to 0.8. Try to increase the number of tuning steps.
There were 27 divergences after tuning. Increase `target_accept` or reparameterize.
There were 9 divergences after tuning. Increase `target_accept` or reparameterize.
There were 11 divergences after tuning. Increase `target_accept` or reparameterize.
The acceptance probability does not match the target. It is 0.7074472966124431, but should be close to 0.8. Try to increase the number of tuning steps.
The rhat statistic is larger than 1.05 for some parameters. This indicates slight problems during sampling.
The estimated number of effective samples is smaller than 200 for some parameters.

100.00% [4000/4000 00:01<00:00]

Most ArviZ functions work fine with trace objects from PyMC3 while it is recommended to covert to inferenceData before plotting:

az.plot_autocorr(centered_eight_trace, var_names=["mu", "tau"]);

Got error No model on context stack. trying to find log_likelihood in translation.
/Users/yilinxia/opt/miniconda3/envs/arviz/lib/python3.8/site-packages/arviz/data/io_pymc3_3x.py:98: FutureWarning: Using `from_pymc3` without the model will be deprecated in a future release. Not using the model will return less accurate and less useful results. Make sure you use the model argument or call from_pymc3 within a model context.
  warnings.warn(

Convert to InferenceData

For much more powerful querying, analysis and plotting, we can use built-in ArviZ utilities to convert PyMC3 objects to xarray datasets. Note we are also giving some information about labelling.

ArviZ is built to work with InferenceData. The more groups it has access to, the more powerful analyses it can perform. You can check the InferenceData structure specification here. Given below is a plot of the trace, which is common in PyMC3 workflows. Don’t forget to note the intelligent labels.

data = az.from_pymc3(
    trace=centered_eight_trace,
    prior=prior,
    posterior_predictive=posterior_predictive,
    model=centered_eight,
    coords={"school": schools},
    dims={"theta": ["school"], "obs": ["school"]},
)
data

arviz.InferenceData

• posterior

  <xarray.Dataset>
  Dimensions:  (chain: 4, draw: 1000, school: 8)
  Coordinates:
    * chain    (chain) int64 0 1 2 3
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      mu       (chain, draw) float64 4.333 4.754 -0.6848 ... 1.239 3.973 -0.9172
      theta    (chain, draw, school) float64 6.962 4.182 4.186 ... -9.074 1.109
      tau      (chain, draw) float64 4.318 6.168 5.434 2.579 ... 2.299 6.237 5.914
  Attributes:
      created_at:                 2022-04-15T05:26:48.456950
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5
      sampling_time:              13.786029815673828
      tuning_steps:               1000

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • school: 8
  • Coordinates: (3)
    • 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])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (3)
    • mu
      (chain, draw)
      float64
      4.333 4.754 ... 3.973 -0.9172

      array([[ 4.33284091,  4.75392728, -0.68484227, ...,  4.25061968,
               4.25061968,  4.25061968],
             [ 4.55419162,  8.59823361, -4.387941  , ...,  4.580372  ,
               0.57630971,  0.04382526],
             [-3.27410696,  2.41587419,  5.02931395, ...,  6.09772326,
               4.27353754,  5.00312407],
             [ 6.25374453,  3.40273268,  3.66194173, ...,  1.23948628,
               3.97253149, -0.91724975]])

    • theta
      (chain, draw, school)
      float64
      6.962 4.182 4.186 ... -9.074 1.109

      array([[[  6.96238111,   4.18186119,   4.18566067, ...,   8.289649  ,
                 8.32871188,  10.07969623],
              [ 12.4661021 ,   2.86356901,   8.98155129, ...,  -2.33170297,
                16.77615465,   0.63324276],
              [  0.56366463,   6.08491074,  -1.79871861, ...,   7.97006825,
                -1.70697409,   4.39332426],
              ...,
              [  5.62132732,   3.11126918,   5.11902771, ...,   4.3880613 ,
                 2.90994524,   3.16746358],
              [  5.62132732,   3.11126918,   5.11902771, ...,   4.3880613 ,
                 2.90994524,   3.16746358],
              [  5.62132732,   3.11126918,   5.11902771, ...,   4.3880613 ,
                 2.90994524,   3.16746358]],
      
             [[  2.07476826,   8.79715477,   8.13459329, ...,   2.17820939,
                 5.28910913,   2.10935557],
              [ 23.85923077,   0.72006008,  -5.58503933, ...,   4.01020892,
                16.6311221 ,  10.37932748],
              [ 17.0087838 ,   4.30984648, -18.63321535, ...,   1.71918947,
                12.36160819,  16.20115302],
      ...
              [ 14.57125009,   3.45529525,  -2.83173026, ...,  18.19242754,
                 3.48728489,   7.04768857],
              [ 20.86100999,  -3.04979899,  -5.04810251, ...,  10.8372059 ,
                 7.90558349,   6.43301921],
              [  5.08803394,   6.91800009,  17.3786901 , ...,   0.72541794,
                11.97793624,  -3.15497584]],
      
             [[  4.47048733,   3.82157824,   8.251458  , ...,   1.70788895,
                28.65754526,  10.47193666],
              [  5.81806512,  12.47665327,   5.98703722, ...,  14.79699056,
                 2.62401369,   5.86525691],
              [  9.04597323,  -0.31585424,  -4.05449548, ...,   3.20645818,
                 7.43409874,  -1.15682589],
              ...,
              [  3.30180283,   2.86622619,   4.6223238 , ...,   1.15441344,
                 2.86174765,   0.50871093],
              [ 10.38382193,   0.15218507,   5.03826493, ...,   5.86603744,
                -4.37602329,   2.16415305],
              [  5.43756007,   5.79900276,  -6.32216702, ...,   4.5500347 ,
                -9.07419334,   1.1090621 ]]])

    • tau
      (chain, draw)
      float64
      4.318 6.168 5.434 ... 6.237 5.914

      array([[ 4.31800145,  6.16836983,  5.43358493, ...,  0.82256729,
               0.82256729,  0.82256729],
             [ 2.91977829,  8.3696382 , 13.28651529, ...,  8.11279661,
               5.71585321,  7.38198731],
             [ 5.09386531,  6.94315969,  7.65834365, ...,  8.31296143,
               6.80928935,  3.848819  ],
             [ 7.13191747,  7.97228244,  2.83204967, ...,  2.29897627,
               6.23694218,  5.91435416]])

  • Attributes: (6)

    created_at :

    2022-04-15T05:26:48.456950

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

    sampling_time :

    13.786029815673828

    tuning_steps :

    1000

  
  
• posterior_predictive

  <xarray.Dataset>
  Dimensions:  (chain: 4, draw: 1000, school: 8)
  Coordinates:
    * chain    (chain) int64 0 1 2 3
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      obs      (chain, draw, school) float64 3.892 -3.475 -4.494 ... -25.62 -7.858
  Attributes:
      created_at:                 2022-04-15T05:26:48.654469
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • school: 8
  • Coordinates: (3)
    • 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])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (1)
    • obs
      (chain, draw, school)
      float64
      3.892 -3.475 ... -25.62 -7.858

      array([[[ 3.89197544e+00, -3.47539690e+00, -4.49439069e+00, ...,
               -1.20588914e+01,  1.93744996e+01,  2.87790013e+00],
              [ 1.00573531e+01, -1.23038997e+01, -2.07565844e+01, ...,
               -2.84425180e+00,  3.71474837e+01,  8.15103268e+00],
              [-3.68493961e+01,  7.96737252e+00, -2.86716646e+01, ...,
               -1.82107595e-01,  6.40918750e+00, -4.46226293e+01],
              ...,
              [ 1.52798266e+01,  1.39236441e+00,  1.32598720e+01, ...,
                4.53256953e+00, -1.78589241e+00,  2.12746103e+01],
              [ 2.21557547e+01,  1.68701120e+00, -3.44127957e+00, ...,
                1.49309391e-02,  1.79033145e+01,  8.35258926e+00],
              [-8.39923264e+00,  1.63427030e+01, -7.14703391e+00, ...,
               -9.72705524e+00, -1.51992409e+00, -1.62565517e+01]],
      
             [[-9.12461445e+00, -2.40751772e+00,  7.26594016e+00, ...,
                2.12256874e+01, -4.18134240e-01,  9.90164907e+00],
              [ 2.08523938e+01,  2.61557511e-01,  1.52517310e+01, ...,
                9.50446779e-01,  1.69803282e+01,  8.81512490e-01],
              [ 8.50454833e+00,  1.46470542e+01, -1.31202139e+01, ...,
               -7.40051829e-01,  1.23544006e+01,  3.68706711e+01],
      ...
              [-6.40665545e+00,  3.08796400e+00, -2.75341615e+01, ...,
                3.70263813e+01,  5.29501736e-01, -6.62554108e+00],
              [-6.20354839e+00, -6.56274198e-01, -1.28534169e+01, ...,
                1.39093220e+01,  2.35890006e+01,  2.12122930e+01],
              [ 9.48472462e+00,  4.05560807e-01,  3.19221601e+01, ...,
               -1.13703624e+01,  9.20216829e+00, -2.57663156e+01]],
      
             [[-9.44123975e+00, -3.74859930e+00,  3.65694025e+01, ...,
               -5.87345592e+00,  4.33860562e+01, -2.37112954e+00],
              [ 1.82192272e+01,  2.61586390e+01,  2.02684347e+01, ...,
                1.37116446e+01,  9.95501085e+00,  2.56116130e+01],
              [ 1.42309423e+01, -1.01225583e+01,  1.16372114e+00, ...,
                1.87922211e+01, -2.03304098e+01, -2.06095514e+01],
              ...,
              [ 1.51553671e+01, -6.89522989e+00,  1.10556563e+01, ...,
                4.85337865e+00,  1.89259914e+00,  1.22403246e+01],
              [ 6.47953216e+00, -3.24026999e+00,  3.26419088e+01, ...,
                7.11206103e+00,  9.76000507e+00,  2.46580357e+01],
              [ 3.12896972e+01,  3.05466603e+01, -3.12671294e+01, ...,
                2.99534823e+01, -2.56182864e+01, -7.85823146e+00]]])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:48.654469

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

  
  
• log_likelihood

  <xarray.Dataset>
  Dimensions:  (chain: 4, draw: 1000, school: 8)
  Coordinates:
    * chain    (chain) int64 0 1 2 3
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      obs      (chain, draw, school) float64 -4.611 -3.294 ... -6.887 -3.992
  Attributes:
      created_at:                 2022-04-15T05:26:48.652469
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • school: 8
  • Coordinates: (3)
    • 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])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (1)
    • obs
      (chain, draw, school)
      float64
      -4.611 -3.294 ... -6.887 -3.992

      array([[[-4.61050298, -3.29441455, -3.79237436, ..., -3.53641638,
               -3.6891927 , -3.81500098],
              [-4.16321537, -3.35343824, -3.97191314, ..., -3.36270259,
               -3.22901261, -4.0086979 ],
              [-5.29977206, -3.23986146, -3.69434577, ..., -3.51758526,
               -5.16334776, -3.89860275],
              ...,
              [-4.73988871, -3.34102207, -3.82027454, ..., -3.36426752,
               -4.36007239, -3.9297018 ],
              [-4.73988871, -3.34102207, -3.82027454, ..., -3.36426752,
               -4.36007239, -3.9297018 ],
              [-4.73988871, -3.34102207, -3.82027454, ..., -3.36426752,
               -4.36007239, -3.9297018 ]],
      
             [[-5.12058349, -3.2247009 , -3.93367407, ..., -3.32257008,
               -4.02935736, -3.96027456],
              [-3.66509089, -3.48651125, -3.70457887, ..., -3.35427743,
               -3.23089276, -3.81336365],
              [-3.89544836, -3.28960979, -4.16886597, ..., -3.31897113,
               -3.38048094, -3.83654746],
      ...
              [-4.02772501, -3.32479533, -3.69158256, ..., -4.53823697,
               -4.27461812, -3.84715811],
              [-3.74024469, -3.83201391, -3.69972008, ..., -3.7167124 ,
               -3.73100985, -3.85713633],
              [-4.79356249, -3.22737725, -4.50264251, ..., -3.31714536,
               -3.40284989, -4.16374438]],
      
             [[-4.85729533, -3.30881967, -3.93878371, ..., -3.3189045 ,
               -3.78943998, -3.81291365],
              [-4.72040703, -3.32172575, -3.84927499, ..., -4.10343277,
               -4.4036284 , -3.86738911],
              [-4.42533347, -3.56729078, -3.69369905, ..., -3.3369514 ,
               -3.77971497, -4.07644311],
              ...,
              [-4.98254639, -3.35330179, -3.80500347, ..., -3.31693233,
               -4.36735705, -4.01309073],
              [-4.31661036, -3.52946462, -3.81772589, ..., -3.4146781 ,
               -5.72495572, -3.95860641],
              [-4.75824139, -3.24574557, -3.71308349, ..., -3.36891127,
               -6.88658335, -3.99235432]]])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:48.652469

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

  
  
• sample_stats

  <xarray.Dataset>
  Dimensions:             (chain: 4, draw: 1000)
  Coordinates:
    * chain               (chain) int64 0 1 2 3
    * draw                (draw) int64 0 1 2 3 4 5 6 ... 994 995 996 997 998 999
  Data variables: (12/13)
      diverging           (chain, draw) bool False False False ... False False
      tree_depth          (chain, draw) int64 4 4 4 4 2 4 4 4 ... 3 2 2 3 4 3 3 5
      acceptance_rate     (chain, draw) float64 0.8917 0.6986 ... 0.3739 0.9883
      energy              (chain, draw) float64 64.75 62.72 66.46 ... 62.89 67.11
      n_steps             (chain, draw) float64 15.0 15.0 15.0 ... 7.0 7.0 31.0
      step_size_bar       (chain, draw) float64 0.3722 0.3722 ... 0.3021 0.3021
      ...                  ...
      max_energy_error    (chain, draw) float64 0.224 51.41 ... 1.013 -0.1603
      lp                  (chain, draw) float64 -55.42 -59.67 ... -59.9 -61.02
      process_time_diff   (chain, draw) float64 0.001564 0.001535 ... 0.002764
      perf_counter_diff   (chain, draw) float64 0.001563 0.001534 ... 0.002763
      perf_counter_start  (chain, draw) float64 11.36 11.36 11.36 ... 13.63 13.64
      step_size           (chain, draw) float64 0.4712 0.4712 ... 0.4006 0.4006
  Attributes:
      created_at:                 2022-04-15T05:26:48.462552
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5
      sampling_time:              13.786029815673828
      tuning_steps:               1000

  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])

  • Data variables: (13)
    • diverging
      (chain, draw)
      bool
      False False False ... False False

      array([[False, False, False, ..., False,  True, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False]])

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

      array([[4, 4, 4, ..., 1, 2, 1],
             [3, 4, 5, ..., 4, 4, 4],
             [5, 4, 4, ..., 4, 4, 4],
             [4, 4, 4, ..., 3, 3, 5]])

    • acceptance_rate
      (chain, draw)
      float64
      0.8917 0.6986 ... 0.3739 0.9883

      array([[8.91694429e-001, 6.98563417e-001, 9.26641481e-001, ...,
              2.29957905e-276, 1.22603899e-060, 1.00424735e-147],
             [9.81822853e-001, 6.60284163e-001, 9.11399136e-001, ...,
              9.68904748e-001, 9.89812920e-001, 9.32033826e-001],
             [9.43336994e-001, 8.73353508e-001, 9.97441137e-001, ...,
              7.51902481e-001, 9.99717607e-001, 8.21137628e-001],
             [9.97608193e-001, 9.23969268e-001, 9.59776134e-001, ...,
              1.00000000e+000, 3.73871409e-001, 9.88284699e-001]])

    • energy
      (chain, draw)
      float64
      64.75 62.72 66.46 ... 62.89 67.11

      array([[64.75364368, 62.72020724, 66.45562417, ..., 51.98945828,
              48.51261271, 49.74881529],
             [58.94898856, 62.93417499, 68.40633326, ..., 64.22195429,
              66.22140502, 68.55086236],
             [63.48685234, 66.23090817, 66.35211684, ..., 69.19225609,
              64.82765234, 66.84037177],
             [67.36266346, 68.77338595, 67.22646605, ..., 56.74224406,
              62.88882935, 67.11039219]])

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

      array([[15., 15., 15., ...,  1.,  3.,  1.],
             [ 7., 15., 31., ..., 15., 15., 15.],
             [31., 15., 15., ..., 15., 15., 15.],
             [15., 15., 15., ...,  7.,  7., 31.]])

    • step_size_bar
      (chain, draw)
      float64
      0.3722 0.3722 ... 0.3021 0.3021

      array([[0.37220628, 0.37220628, 0.37220628, ..., 0.37220628, 0.37220628,
              0.37220628],
             [0.27828811, 0.27828811, 0.27828811, ..., 0.27828811, 0.27828811,
              0.27828811],
             [0.30380677, 0.30380677, 0.30380677, ..., 0.30380677, 0.30380677,
              0.30380677],
             [0.3020568 , 0.3020568 , 0.3020568 , ..., 0.3020568 , 0.3020568 ,
              0.3020568 ]])

    • energy_error
      (chain, draw)
      float64
      -0.08259 0.3743 ... 0.9854 -0.03012

      array([[-0.08259485,  0.37433722,  0.04132455, ...,  0.        ,
               0.        ,  0.        ],
             [ 0.0847979 ,  0.44258818,  0.04262487, ...,  0.0482196 ,
              -0.16359819,  0.07108232],
             [ 0.06507037,  0.19417104, -0.07553831, ...,  0.34174769,
              -0.01390547,  0.48718085],
             [-0.00413115,  0.11973169,  0.79717764, ..., -0.8280887 ,
               0.98536643, -0.0301198 ]])

    • max_energy_error
      (chain, draw)
      float64
      0.224 51.41 ... 1.013 -0.1603

      array([[ 2.24036920e-01,  5.14118840e+01,  1.77932859e-01, ...,
               6.34680760e+02,  1.13795217e+03,  3.38475770e+02],
             [-1.55488342e-01,  7.43792632e-01,  2.40398672e+00, ...,
              -2.10362356e-01, -4.06426253e-01,  1.79603010e-01],
             [ 6.00786159e-01,  2.66160868e-01, -9.96876906e-02, ...,
               4.28312255e-01, -1.92922244e-01,  4.87180847e-01],
             [-1.51818016e-01,  5.10961970e-01,  7.97177645e-01, ...,
              -8.28088702e-01,  1.01304119e+00, -1.60305838e-01]])

    • lp
      (chain, draw)
      float64
      -55.42 -59.67 ... -59.9 -61.02

      array([[-55.41863193, -59.66525994, -58.99963096, ..., -46.71350291,
              -46.71350291, -46.71350291],
             [-55.5128382 , -62.19738055, -65.44934107, ..., -60.48624146,
              -59.2929766 , -63.23177663],
             [-59.43723251, -62.92158617, -59.98527815, ..., -61.86236523,
              -61.08019822, -63.36288695],
             [-64.04490726, -60.97693754, -61.44469257, ..., -51.01093052,
              -59.89960925, -61.02333548]])

    • process_time_diff
      (chain, draw)
      float64
      0.001564 0.001535 ... 0.002764

      array([[0.001564, 0.001535, 0.001552, ..., 0.00025 , 0.000418, 0.000243],
             [0.00085 , 0.001559, 0.003108, ..., 0.001609, 0.001615, 0.00164 ],
             [0.004168, 0.001617, 0.001556, ..., 0.001554, 0.001522, 0.001516],
             [0.001635, 0.001659, 0.001589, ..., 0.000806, 0.000784, 0.002764]])

    • perf_counter_diff
      (chain, draw)
      float64
      0.001563 0.001534 ... 0.002763

      array([[0.00156282, 0.00153433, 0.00155187, ..., 0.00024876, 0.00041791,
              0.0002425 ],
             [0.00084959, 0.00155865, 0.00310714, ..., 0.00160806, 0.00161478,
              0.00163952],
             [0.00416877, 0.00162256, 0.00155562, ..., 0.00155387, 0.0015215 ,
              0.00151558],
             [0.00163555, 0.00165882, 0.0015893 , ..., 0.00080551, 0.00078406,
              0.00276349]])

    • perf_counter_start
      (chain, draw)
      float64
      11.36 11.36 11.36 ... 13.63 13.64

      array([[11.35552752, 11.3572312 , 11.35888869, ..., 12.41607785,
              12.4164371 , 12.41697478],
             [11.47275188, 11.4737375 , 11.47542265, ..., 13.0093599 ,
              13.01111528, 13.01287558],
             [11.69751155, 11.70183779, 11.70359012, ..., 13.24330615,
              13.24498009, 13.24661974],
             [12.46543444, 12.46723304, 12.4690387 , ..., 13.63360082,
              13.63451202, 13.63540219]])

    • step_size
      (chain, draw)
      float64
      0.4712 0.4712 ... 0.4006 0.4006

      array([[0.47117558, 0.47117558, 0.47117558, ..., 0.47117558, 0.47117558,
              0.47117558],
             [0.30562691, 0.30562691, 0.30562691, ..., 0.30562691, 0.30562691,
              0.30562691],
             [0.55337819, 0.55337819, 0.55337819, ..., 0.55337819, 0.55337819,
              0.55337819],
             [0.40058279, 0.40058279, 0.40058279, ..., 0.40058279, 0.40058279,
              0.40058279]])

  • Attributes: (6)

    created_at :

    2022-04-15T05:26:48.462552

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

    sampling_time :

    13.786029815673828

    tuning_steps :

    1000

  
  
• prior

  <xarray.Dataset>
  Dimensions:    (chain: 1, draw: 500, school: 8)
  Coordinates:
    * chain      (chain) int64 0
    * draw       (draw) int64 0 1 2 3 4 5 6 7 ... 492 493 494 495 496 497 498 499
    * school     (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      tau_log__  (chain, draw) float64 0.7613 2.303 1.944 ... 0.7547 1.664 3.069
      theta      (chain, draw, school) float64 -5.797 -7.748 ... 26.96 -38.51
      tau        (chain, draw) float64 2.141 10.0 6.988 ... 2.127 5.279 21.52
      mu         (chain, draw) float64 -5.671 -2.992 -5.939 ... 11.19 1.339 1.965
  Attributes:
      created_at:                 2022-04-15T05:26:48.656780
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5

  xarray.Dataset

  • Dimensions:
    • chain: 1
    • draw: 500
    • school: 8
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (4)
    • tau_log__
      (chain, draw)
      float64
      0.7613 2.303 1.944 ... 1.664 3.069

      array([[ 7.61345408e-01,  2.30303944e+00,  1.94414114e+00,
               2.66770568e+00, -1.60938713e+00,  7.50172384e-01,
               3.12262279e+00,  9.19079474e-01,  3.13132523e+00,
               2.36469977e+00,  1.50955070e-03,  3.10785381e+00,
               4.70090086e-01,  1.73433280e-01, -1.93816174e-01,
               3.36467088e+00,  1.61994874e+00,  2.55626413e+00,
               1.67721906e+00,  3.23286285e+00,  3.35966816e+00,
               2.06000066e+00,  1.93176709e+00,  1.64395626e+00,
              -3.05399750e-02,  1.60821827e-01,  1.48205795e+00,
               3.43659791e+00,  1.65573756e+00,  9.19120452e-01,
               1.17278391e+00,  1.90577154e+00,  1.61103600e+00,
               2.71561615e+00,  9.03554130e-01,  2.04231434e+00,
               2.22841888e+00,  2.22282875e+00,  2.35283476e+00,
               1.26852482e+00,  4.34962520e+00,  6.95581373e-01,
               3.09444861e+00, -4.13486542e-01, -4.63757975e-01,
               1.48174728e+00,  2.78322158e+00, -2.94290057e-02,
               6.22867710e-01,  1.99460792e+00,  9.24522141e-01,
               2.96609890e+00,  1.80360271e+00,  1.24978929e+00,
               1.24694020e+00,  6.18511512e-01,  8.47163322e-01,
               1.70421786e+00,  2.06536840e+00,  3.01286127e+00,
      ...
               2.39590452e+00,  5.37298385e-01,  1.25594396e+00,
               6.52248332e-01,  1.39801010e+00,  1.27005926e+00,
               1.07824753e+00,  1.82905813e+00, -7.31420807e-01,
               2.69953205e+00,  2.15660522e+00, -2.72466077e-01,
               2.56562263e+00,  2.04565773e+00, -2.98573361e+00,
               3.75119879e-01,  2.62910111e+00,  1.39509303e+00,
              -2.39154816e+00,  1.17170513e+00,  1.87875408e+00,
               1.33256499e+00,  1.30437565e-01,  1.60610801e+00,
               9.93938832e-01,  1.74140550e+00,  2.19540801e+00,
               2.04222739e+00,  1.72856731e+00,  2.61420625e+00,
              -1.22204797e+00,  6.66300299e-01,  8.12832546e-01,
               1.91503164e+00, -1.04994717e-01,  1.73667681e+00,
               5.82168793e-01,  1.28695908e+00,  2.36110372e+00,
               3.19406051e+00, -6.03649757e-02,  2.55958668e+00,
               1.26016067e+00,  1.38407621e+00,  2.99712804e+00,
               3.09964886e-01,  1.77389988e+00,  5.46767778e-01,
               9.37072462e-01,  1.65367290e+00,  1.91976180e+00,
               1.35052703e+00,  2.62873847e+00,  5.78422802e-01,
              -1.41988443e+00,  1.82989006e+00,  7.54716254e-01,
               1.66378305e+00,  3.06898171e+00]])

    • theta
      (chain, draw, school)
      float64
      -5.797 -7.748 ... 26.96 -38.51

      array([[[ -5.79744075,  -7.74833526,  -7.04798719, ...,  -6.26323207,
                -9.42306823,  -1.33251561],
              [ -1.20127255,  -7.97892117,   6.10211883, ...,  -5.35132756,
                -7.94849495,  -1.59713787],
              [-13.50012033,   5.23169963,  -0.3322567 , ..., -11.86261593,
                -4.65504822, -14.48537168],
              ...,
              [ 10.44162328,  12.97467062,   9.6899943 , ...,  12.86311451,
                 9.98984831,  13.2587891 ],
              [  1.82817192,   2.05339635,   8.08807398, ...,   3.75413424,
                -1.1995294 ,   9.91121507],
              [ -4.33157526,  39.31720694,  -0.90964751, ...,  20.42453677,
                26.95964677, -38.5054179 ]]])

    • tau
      (chain, draw)
      float64
      2.141 10.0 6.988 ... 5.279 21.52

      array([[2.14115501e+00, 1.00045445e+01, 6.98762792e+00, 1.44068772e+01,
              2.00010157e-01, 2.11736498e+00, 2.27058542e+01, 2.50698159e+00,
              2.29043129e+01, 1.06408436e+01, 1.00151069e+00, 2.23729762e+01,
              1.60013834e+00, 1.18938133e+00, 8.23809328e-01, 2.89239761e+01,
              5.05283131e+00, 1.28875810e+01, 5.35065543e+00, 2.53521326e+01,
              2.87796392e+01, 7.84597495e+00, 6.90169538e+00, 5.17560511e+00,
              9.69921659e-01, 1.17447569e+00, 4.40199544e+00, 3.10810376e+01,
              5.23694105e+00, 2.50708432e+00, 3.23097487e+00, 6.72459393e+00,
              5.00799684e+00, 1.51139197e+01, 2.46836041e+00, 7.70842846e+00,
              9.28517351e+00, 9.23341299e+00, 1.05153359e+01, 3.55560352e+00,
              7.74494295e+01, 2.00487432e+00, 2.20750633e+01, 6.61340435e-01,
              6.28915750e-01, 4.40062808e+00, 1.61710334e+01, 9.70999811e-01,
              1.86426656e+00, 7.34932091e+00, 2.52066345e+00, 1.94160278e+01,
              6.07148191e+00, 3.48960759e+00, 3.47967954e+00, 1.85616311e+00,
              2.33301943e+00, 5.49708448e+00, 7.88820336e+00, 2.03455308e+01,
              8.17400306e+00, 1.40208448e+00, 1.02267851e+00, 1.86904002e+01,
              6.27821283e-01, 1.19302769e+00, 4.85563957e+00, 1.33249011e+02,
              1.09740566e+00, 6.21773550e-01, 5.43680628e+00, 1.74774635e+00,
              1.97792031e+00, 1.74807359e+00, 8.59511839e-01, 8.26296707e+00,
              3.88364258e+01, 1.15364696e+01, 6.15072592e+00, 4.48440911e+00,
      ...
              1.50974505e+01, 2.78748447e+01, 3.58169826e+00, 9.18701373e+00,
              1.84092942e+01, 1.32201898e+00, 2.03791370e+00, 1.05925705e+01,
              3.83068632e+00, 1.83120923e+01, 2.07437091e+00, 5.31475717e+00,
              2.69689321e+01, 6.33301385e+00, 1.42471109e+01, 6.88664828e+00,
              1.72377104e+01, 7.47483711e+01, 1.02102637e+00, 3.18239250e+00,
              5.68874968e+00, 1.09781234e+01, 1.71137713e+00, 3.51115119e+00,
              1.91985245e+00, 4.04713856e+00, 3.56106357e+00, 2.93952361e+00,
              6.22801794e+00, 4.81224777e-01, 1.48727704e+01, 8.64175100e+00,
              7.61499261e-01, 1.30087555e+01, 7.73424388e+00, 5.05024407e-02,
              1.45516585e+00, 1.38613046e+01, 4.03534997e+00, 9.14879364e-02,
              3.22749125e+00, 6.54534481e+00, 3.79075419e+00, 1.13932680e+00,
              4.98337819e+00, 2.70185570e+00, 5.70535669e+00, 8.98366575e+00,
              7.70775825e+00, 5.63257838e+00, 1.36563723e+01, 2.94626162e-01,
              1.94702058e+00, 2.25428432e+00, 6.78715356e+00, 9.00329279e-01,
              5.67844148e+00, 1.78991618e+00, 3.62175633e+00, 1.06026473e+01,
              2.43872512e+01, 9.41420875e-01, 1.29304717e+01, 3.52598797e+00,
              3.99113724e+00, 2.00279349e+01, 1.36337724e+00, 5.89379367e+00,
              1.72765980e+00, 2.55249793e+00, 5.22613972e+00, 6.81933392e+00,
              3.85945904e+00, 1.38562787e+01, 1.78322371e+00, 2.41741952e-01,
              6.23320135e+00, 2.12700791e+00, 5.27924474e+00, 2.15199779e+01]])

    • mu
      (chain, draw)
      float64
      -5.671 -2.992 ... 1.339 1.965

      array([[-5.67113287e+00, -2.99242806e+00, -5.93858448e+00,
               1.42047047e+00,  3.95105426e+00, -4.34958322e+00,
              -3.22267386e-01, -1.43473569e+00,  2.04508977e+00,
              -1.02254563e+01,  4.21061986e+00,  8.43590650e-01,
               1.53478955e-01,  8.40624713e+00, -1.04416722e+00,
              -2.93739480e+00,  3.89099238e-01,  6.67180136e+00,
              -1.25101205e-01, -1.03573084e+01,  6.97369960e-01,
              -4.99911257e+00,  5.58463985e+00,  2.92369435e+00,
              -3.46904943e+00,  8.43630435e+00, -4.26970123e-01,
               1.45273285e+00, -1.88519887e+00, -3.09860524e-01,
              -3.12144265e+00, -1.98340853e-01,  3.25383347e+00,
              -8.88900844e+00, -4.29146708e+00,  1.45581881e+00,
              -6.22790288e+00, -2.63759991e+00, -4.06604261e-01,
               4.42019683e+00, -1.39107147e+00, -1.21343669e+00,
              -7.86382641e+00, -6.89368919e+00, -2.10712340e+00,
              -7.44223786e+00, -7.14275069e-01,  4.76769654e+00,
               4.68836842e+00, -1.70560189e+00,  1.68077670e+01,
               2.10660113e+00,  1.15313197e+01,  4.33736941e+00,
               2.05403651e+00,  7.48890149e+00, -6.00211224e+00,
              -9.25950249e+00,  2.32700346e+00,  1.63231802e+01,
      ...
               2.59085297e+00, -1.45760469e+00,  9.26320645e+00,
               3.09945700e+00, -3.73395652e+00,  2.01401710e+00,
               6.16549609e-01,  2.48878144e-01,  3.12042022e+00,
              -4.78785448e+00,  2.24772894e+00,  7.88018356e+00,
              -3.74304689e+00,  3.10107991e+00, -4.02719759e+00,
               1.41128651e+00, -3.44207063e+00,  3.22812803e+00,
               8.28678523e-01,  3.36536992e+00,  4.88497660e+00,
               6.51326927e+00, -4.77430967e+00, -9.29992229e+00,
               1.55141264e+00, -3.60498083e+00,  2.14378204e-01,
              -4.45223066e-01, -8.08226281e+00, -4.48475930e+00,
              -2.91912539e+00, -8.82090418e-01,  4.37347386e+00,
               7.60085348e+00, -3.31532215e+00,  9.23907235e-01,
              -2.71291235e-01,  3.64372478e-01, -3.06748755e-01,
              -3.61452124e+00,  4.98464986e+00, -5.38117754e+00,
              -5.45838650e+00,  2.11363774e+00, -3.36093025e+00,
              -5.77843229e+00,  6.14100511e+00, -7.58549890e+00,
              -8.51202706e-01,  6.01767513e+00,  3.20978070e+00,
              -3.25066013e+00,  4.32675174e+00, -9.39309354e+00,
               3.07580515e+00,  1.69654241e+00,  1.11911807e+01,
               1.33908553e+00,  1.96494493e+00]])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:48.656780

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

  
  
• prior_predictive

  <xarray.Dataset>
  Dimensions:  (chain: 1, draw: 500, school: 8)
  Coordinates:
    * chain    (chain) int64 0
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 492 493 494 495 496 497 498 499
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      obs      (chain, draw, school) float64 -12.02 -21.44 -9.397 ... 39.53 -42.67
  Attributes:
      created_at:                 2022-04-15T05:26:48.658653
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5

  xarray.Dataset

  • Dimensions:
    • chain: 1
    • draw: 500
    • school: 8
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (1)
    • obs
      (chain, draw, school)
      float64
      -12.02 -21.44 ... 39.53 -42.67

      array([[[-12.0177098 , -21.4428818 ,  -9.39706428, ...,   0.0607126 ,
               -11.69881524,   1.6263736 ],
              [ 26.2677061 , -19.61406527, -30.20010464, ..., -28.46708162,
               -12.62995224,   6.75100157],
              [-20.74160183,   7.2537621 , -17.18885728, ..., -21.36994126,
                -2.91668599,  -1.41552866],
              ...,
              [  0.34387747,   5.66070778,  13.69680529, ...,   9.93334957,
                16.56861679, -32.28827645],
              [ -8.16935722,  -3.4544725 ,  23.35732596, ...,  26.3265578 ,
                 1.59141583,  13.60046927],
              [-23.36866302,  38.92331314,  -1.94383522, ...,  10.28784457,
                39.52708582, -42.66937293]]])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:48.658653

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

  
  
• observed_data

  <xarray.Dataset>
  Dimensions:  (school: 8)
  Coordinates:
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      obs      (school) float64 28.0 8.0 -3.0 7.0 -1.0 1.0 18.0 12.0
  Attributes:
      created_at:                 2022-04-15T05:26:48.659375
      arviz_version:              0.12.0
      inference_library:          pymc3
      inference_library_version:  3.11.5

  xarray.Dataset

  • Dimensions:
    • school: 8
  • Coordinates: (1)
    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (1)
    • obs
      (school)
      float64
      28.0 8.0 -3.0 7.0 ... 1.0 18.0 12.0

      array([28.,  8., -3.,  7., -1.,  1., 18., 12.])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:48.659375

    arviz_version :

    0.12.0

    inference_library :

    pymc3

    inference_library_version :

    3.11.5

  
  

az.plot_trace(data, compact=False);

Plotting with PyStan objects

ArviZ is built with first class support for PyStan objects, and can plot raw fit objects in a reasonable manner. Here is the same centered eight schools model:

import nest_asyncio

nest_asyncio.apply()

import stan  # pystan version 3.4.0


schools_code = """
data {
  int<lower=0> J;
  array[J] real y;
  array[J] real<lower=0> sigma;
}

parameters {
  real mu;
  real<lower=0> tau;
  array[J] real theta;
}

model {
  mu ~ normal(0, 5);
  tau ~ cauchy(0, 5);
  theta ~ normal(mu, tau);
  y ~ normal(theta, sigma);
}
generated quantities {
    vector[J] log_lik;
    vector[J] y_hat;
    for (j in 1:J) {
        log_lik[j] = normal_lpdf(y[j] | theta[j], sigma[j]);
        y_hat[j] = normal_rng(theta[j], sigma[j]);
    }
}
"""

schools_dat = {
    "J": 8,
    "y": [28, 8, -3, 7, -1, 1, 18, 12],
    "sigma": [15, 10, 16, 11, 9, 11, 10, 18],
}

posterior = stan.build(schools_code, data=schools_dat, random_seed=1)
fit = posterior.sample(num_chains=4, num_samples=1000)

Building: found in cache, done.
Sampling:   0%
Sampling:  25% (2000/8000)
Sampling:  50% (4000/8000)
Sampling:  75% (6000/8000)
Sampling: 100% (8000/8000)
Sampling: 100% (8000/8000), done.
Messages received during sampling:
  Gradient evaluation took 2.7e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.27 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 3.3e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.33 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 3.8e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.38 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 3.8e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.38 seconds.
  Adjust your expectations accordingly!

az.plot_density(fit, var_names=["mu", "tau"]);

Again, converting to InferenceData (a netCDF datastore that loads data into xarray datasets), we can get much richer labelling and mixing of data. Here is a plot showing where the Hamiltonian sampler had divergences:

data = az.from_pystan(
    posterior=fit,
    posterior_predictive="y_hat",
    observed_data=["y"],
    log_likelihood={"y": "log_lik"},
    coords={"school": schools},
    dims={
        "theta": ["school"],
        "y": ["school"],
        "log_lik": ["school"],
        "y_hat": ["school"],
        "theta_tilde": ["school"],
    },
)
data

arviz.InferenceData

• posterior

  <xarray.Dataset>
  Dimensions:  (chain: 4, draw: 1000, school: 8)
  Coordinates:
    * chain    (chain) int64 0 1 2 3
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      mu       (chain, draw) float64 0.599 0.2557 0.7311 ... 2.197 2.982 1.079
      tau      (chain, draw) float64 2.982 1.938 1.214 1.148 ... 3.032 4.728 2.477
      theta    (chain, draw, school) float64 -1.166 -0.1457 ... 0.7122 1.992
  Attributes:
      created_at:                 2022-04-15T05:26:56.712787
      arviz_version:              0.12.0
      inference_library:          stan
      inference_library_version:  3.4.0
      num_chains:                 4
      num_samples:                1000
      num_thin:                   1
      num_warmup:                 1000
      save_warmup:                0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • school: 8
  • Coordinates: (3)
    • 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])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (3)
    • mu
      (chain, draw)
      float64
      0.599 0.2557 0.7311 ... 2.982 1.079

      array([[ 0.59896547,  0.25571438,  0.73108497, ..., -0.15059322,
               0.67918435, -0.87010202],
             [ 7.05745234,  6.35567999,  7.35843152, ...,  0.87526844,
               4.17743423,  5.30717601],
             [ 2.08297008, -5.89477919,  1.8490305 , ...,  4.23959213,
               0.2636217 ,  0.21021436],
             [ 4.95426303,  7.98776794,  6.56063267, ...,  2.19740965,
               2.9822251 ,  1.07896348]])

    • tau
      (chain, draw)
      float64
      2.982 1.938 1.214 ... 4.728 2.477

      array([[2.9816211 , 1.93844434, 1.21435272, ..., 1.41578689, 2.62687235,
              1.36908555],
             [2.88396261, 1.86794129, 5.55274767, ..., 5.8105006 , 7.01408489,
              5.0303556 ],
             [4.60747686, 6.38604863, 4.23411186, ..., 7.39001531, 5.93210864,
              6.98137627],
             [1.75089626, 2.07194106, 3.05517736, ..., 3.03236931, 4.7281319 ,
              2.47659762]])

    • theta
      (chain, draw, school)
      float64
      -1.166 -0.1457 ... 0.7122 1.992

      array([[[-1.16626572e+00, -1.45653964e-01,  5.36719641e-02, ...,
               -8.19323205e-03,  3.66556523e+00, -2.98744144e+00],
              [ 9.02401788e-01,  1.36596548e+00,  1.18376387e+00, ...,
                6.56968959e-01,  3.88075676e+00, -1.34145977e+00],
              [ 1.55345651e+00,  1.83347355e+00,  8.03374999e-01, ...,
                1.49358604e+00, -7.20698607e-01,  1.75715224e+00],
              ...,
              [ 2.65333426e+00,  3.27547749e+00,  1.54189616e+00, ...,
               -1.23874301e+00,  1.63146674e+00, -3.46618555e-01],
              [-5.01299858e-01, -1.54935053e+00, -1.74932478e-01, ...,
                1.30496146e+00,  2.03053848e+00, -1.79240549e+00],
              [ 8.10687828e-01,  1.60926183e+00, -1.75480928e+00, ...,
               -1.16738703e+00,  1.94338114e+00, -8.25410107e-01]],
      
             [[ 1.13654372e+01,  3.74800218e+00,  5.31715703e+00, ...,
                6.87376020e+00,  9.29336225e+00,  1.07602323e+01],
              [ 6.38217639e+00,  1.03675892e+01,  8.59112203e+00, ...,
                7.55048747e+00,  5.20186434e+00,  2.97345056e+00],
              [ 1.19281893e+01,  2.97759629e+00,  8.08652561e+00, ...,
                5.86551348e+00,  1.36732286e+01,  1.18571960e+01],
      ...
              [ 8.87343117e-01,  5.66400653e+00,  1.09515540e+00, ...,
                2.13114484e-01,  1.80844877e+01, -3.23828833e+00],
              [ 8.81650004e+00, -1.48652535e+00, -2.39061102e+00, ...,
                8.30800373e+00,  5.48373020e+00, -2.21364309e+00],
              [-5.05247433e+00,  2.13887788e+00,  1.07187205e-01, ...,
               -1.04699509e+01,  2.33497483e-01,  1.88955095e+00]],
      
             [[ 6.49816078e+00,  6.12882787e+00,  5.81606607e+00, ...,
                7.58290107e+00,  6.29847639e+00,  7.57222241e+00],
              [ 4.76129335e+00,  4.04002492e+00,  6.42549308e+00, ...,
                7.63963046e+00,  9.70464478e+00,  5.21559739e+00],
              [ 6.62117099e+00,  7.96960577e+00,  5.31348985e+00, ...,
                3.48760422e+00,  2.27810590e+00,  6.22812093e+00],
              ...,
              [ 8.17518250e-01,  1.12940688e+01,  2.72908095e+00, ...,
                4.41756888e+00, -1.38292923e+00,  6.40808597e+00],
              [ 2.38621050e+00,  7.31991929e+00,  1.34191158e+00, ...,
               -3.36833961e+00,  3.84815474e+00,  3.69253428e+00],
              [-2.51098234e+00, -3.79332185e+00,  1.98985001e+00, ...,
                4.01277654e+00,  7.12161371e-01,  1.99214790e+00]]])

  • Attributes: (9)

    created_at :

    2022-04-15T05:26:56.712787

    arviz_version :

    0.12.0

    inference_library :

    stan

    inference_library_version :

    3.4.0

    num_chains :

    4

    num_samples :

    1000

    num_thin :

    1

    num_warmup :

    1000

    save_warmup :

    0

  
  
• posterior_predictive

  <xarray.Dataset>
  Dimensions:  (chain: 4, draw: 1000, school: 8)
  Coordinates:
    * chain    (chain) int64 0 1 2 3
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      y_hat    (chain, draw, school) float64 -17.26 12.4 7.823 ... -14.9 -27.98
  Attributes:
      created_at:                 2022-04-15T05:26:56.789583
      arviz_version:              0.12.0
      inference_library:          stan
      inference_library_version:  3.4.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • school: 8
  • Coordinates: (3)
    • 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])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (1)
    • y_hat
      (chain, draw, school)
      float64
      -17.26 12.4 7.823 ... -14.9 -27.98

      array([[[-17.25779147,  12.40262498,   7.82271588, ...,  -8.7384816 ,
                22.55553611,  -8.49645821],
              [ -3.35601734,  -1.64784505, -18.94799876, ...,   3.52310244,
                -3.01266829,  19.71513976],
              [ 19.00592347, -15.02656352,  -7.14328144, ...,  10.13313095,
                 4.79168981,  -8.20354128],
              ...,
              [ 11.5919755 ,  16.87498854, -15.80881059, ...,   7.03069036,
                 9.15721981, -16.28970477],
              [-10.83346334, -17.82006606,  -0.30552796, ...,   7.70822444,
                18.08774284, -11.23452107],
              [  9.51268645,  11.75165911,  -7.19074301, ...,  11.81762837,
                -5.04037691,  -0.59108975]],
      
             [[  8.44821191,  21.70563558,  27.94278917, ...,   8.88200212,
                15.44813424,   1.90527655],
              [  6.41071938,   7.30605185, -19.26445036, ...,  11.23903461,
                 8.04245814,   8.12076971],
              [-11.66599151,  -0.15876771,  16.69014978, ...,   4.69035113,
                21.03281934,  28.16095658],
      ...
              [  9.89522202,  23.34230484,  30.22587304, ..., -14.87450666,
                23.88778741,  -9.34680157],
              [-14.8415018 ,  -0.82004036,  -6.44316385, ...,  22.78084694,
                 8.35295632, -10.28915826],
              [-16.20831709,   0.42251925,  17.35707566, ...,   1.89175057,
                -1.60735718,  10.76328548]],
      
             [[ -9.31034286,   0.19480025,  25.70674457, ...,   8.23806359,
                 0.86479985,  11.4823678 ],
              [  6.78361137,  16.02670767,  11.49252648, ...,  16.58991871,
                14.5423763 , -21.37315463],
              [ 13.21287397,  25.19899688, -20.53116722, ...,  20.10147804,
                 6.92336566,  15.0655601 ],
              ...,
              [  9.75676243,  -3.83665446,  -4.06313075, ...,  15.01126342,
                20.43427473,  -6.12147598],
              [ 28.40339445,  12.13328205, -38.21887962, ..., -21.5554171 ,
                17.16160043,  21.40088457],
              [ 30.97501199,  -3.91992746,  -2.26244281, ...,  11.71323189,
               -14.89651212, -27.97748326]]])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:56.789583

    arviz_version :

    0.12.0

    inference_library :

    stan

    inference_library_version :

    3.4.0

  
  
• log_likelihood

  <xarray.Dataset>
  Dimensions:  (chain: 4, draw: 1000, school: 8)
  Coordinates:
    * chain    (chain) int64 0 1 2 3
    * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
    * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  Data variables:
      y        (chain, draw, school) float64 -5.517 -3.553 -3.71 ... -4.716 -3.964
  Attributes:
      created_at:                 2022-04-15T05:26:56.767405
      arviz_version:              0.12.0
      inference_library:          stan
      inference_library_version:  3.4.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • school: 8
  • Coordinates: (3)
    • 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])

    • school
      (school)
      <U16
      'Choate' ... 'Mt. Hermon'

      array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
             'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')

  • Data variables: (1)
    • y
      (chain, draw, school)
      float64
      -5.517 -3.553 ... -4.716 -3.964

      array([[[-5.51736886, -3.55328202, -3.70973998, ..., -3.32103403,
               -4.24890373, -4.15595134],
              [-5.25872169, -3.4415757 , -3.72571452, ..., -3.31732005,
               -4.21828877, -4.08399324],
              [-5.18125465, -3.41165387, -3.7197805 , ..., -3.31784053,
               -4.97384641, -3.97121759],
              ...,
              [-5.0546631 , -3.33312919, -3.73181792, ..., -3.33754443,
               -4.56116803, -4.04455565],
              [-5.43215339, -3.6774741 , -3.70711516, ..., -3.31721811,
               -4.49664213, -4.1028758 ],
              [-5.26978584, -3.4257313 , -3.69455558, ..., -3.33624524,
               -4.51059867, -4.06315465]],
      
             [[-4.24189691, -3.31192105, -3.82663487, ..., -3.45940017,
               -3.60055133, -3.81168224],
              [-4.66550051, -3.24955102, -3.95393763, ..., -3.49414325,
               -4.04048501, -3.93504886],
              [-4.20099562, -3.34764632, -3.9315879 , ..., -3.41465704,
               -3.31512838, -3.80934176],
      ...
              [-5.26053576, -3.24880795, -3.72428174, ..., -3.31939244,
               -3.22155932, -4.16765201],
              [-4.44478134, -3.67149444, -3.69225256, ..., -3.53752355,
               -4.00480867, -4.12108136],
              [-6.05469109, -3.39328739, -3.71038392, ..., -3.86046923,
               -4.79976668, -3.96705903]],
      
             [[-4.65438671, -3.23903005, -3.84333003, ..., -3.49590235,
               -3.9061519 , -3.83956525],
              [-4.82707204, -3.29993064, -3.86504272, ..., -3.49900196,
               -3.56558822, -3.88034134],
              [-4.64266502, -3.22152825, -3.82651576, ..., -3.34240478,
               -4.4574134 , -3.86072169],
              ...,
              [-5.26896054, -3.27577807, -3.75563344, ..., -3.36509735,
               -5.10001335, -3.8575657 ],
              [-5.08491365, -3.22383618, -3.72834795, ..., -3.39568666,
               -4.22289725, -3.91581336],
              [-5.69569994, -3.91693583, -3.74015734, ..., -3.35434134,
               -4.71587045, -3.96387372]]])

  • Attributes: (4)

    created_at :

    2022-04-15T05:26:56.767405

    arviz_version :

    0.12.0

    inference_library :

    stan

    inference_library_version :

    3.4.0

  
  
• sample_stats

  <xarray.Dataset>
  Dimensions:          (chain: 4, draw: 1000)
  Coordinates:
    * chain            (chain) int64 0 1 2 3
    * draw             (draw) int64 0 1 2 3 4 5 6 ... 993 994 995 996 997 998 999
  Data variables:
      acceptance_rate  (chain, draw) float64 0.9271 0.7419 1.0 ... 1.0 0.9957
      step_size        (chain, draw) float64 0.1575 0.1575 ... 0.1294 0.1294
      tree_depth       (chain, draw) int64 3 3 3 3 3 3 3 2 3 ... 5 5 6 5 5 5 5 4 4
      n_steps          (chain, draw) int64 7 15 7 7 7 7 7 ... 63 31 31 31 31 15 15
      diverging        (chain, draw) bool False True False ... False False False
      energy           (chain, draw) float64 15.08 14.79 11.82 ... 22.14 20.47
      lp               (chain, draw) float64 -13.65 -10.93 -7.736 ... -17.2 -17.62
  Attributes:
      created_at:                 2022-04-15T05:26:56.743010
      arviz_version:              0.12.0
      inference_library:          stan
      inference_library_version:  3.4.0
      num_chains:                 4
      num_samples:                1000
      num_thin:                   1
      num_warmup:                 1000
      save_warmup:                0

  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])

  • Data variables: (7)
    • acceptance_rate
      (chain, draw)
      float64
      0.9271 0.7419 1.0 ... 1.0 0.9957

      array([[0.9270755 , 0.74190801, 1.        , ..., 0.96679083, 0.97750843,
              0.98695611],
             [0.7884867 , 0.93179185, 0.56461423, ..., 0.8560158 , 0.95561959,
              0.96021533],
             [0.99653394, 0.99770883, 0.99361147, ..., 0.98532261, 0.99298359,
              0.95042382],
             [0.83013179, 0.84937863, 0.97038593, ..., 0.99089004, 1.        ,
              0.99574856]])

    • step_size
      (chain, draw)
      float64
      0.1575 0.1575 ... 0.1294 0.1294

      array([[0.15752514, 0.15752514, 0.15752514, ..., 0.15752514, 0.15752514,
              0.15752514],
             [0.25791386, 0.25791386, 0.25791386, ..., 0.25791386, 0.25791386,
              0.25791386],
             [0.18260411, 0.18260411, 0.18260411, ..., 0.18260411, 0.18260411,
              0.18260411],
             [0.12944726, 0.12944726, 0.12944726, ..., 0.12944726, 0.12944726,
              0.12944726]])

    • tree_depth
      (chain, draw)
      int64
      3 3 3 3 3 3 3 2 ... 5 6 5 5 5 5 4 4

      array([[3, 3, 3, ..., 3, 3, 3],
             [3, 3, 3, ..., 4, 4, 4],
             [5, 4, 4, ..., 4, 4, 5],
             [3, 3, 4, ..., 5, 4, 4]])

    • n_steps
      (chain, draw)
      int64
      7 15 7 7 7 7 ... 31 31 31 31 15 15

      array([[ 7, 15,  7, ...,  7,  7, 15],
             [15,  7,  7, ..., 15, 15, 15],
             [31, 15, 15, ..., 15, 31, 31],
             [ 7,  7, 15, ..., 31, 15, 15]])

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

      array([[False,  True, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False]])

    • energy
      (chain, draw)
      float64
      15.08 14.79 11.82 ... 22.14 20.47

      array([[15.08338244, 14.78696216, 11.81963712, ..., 18.42674544,
              16.23723736, 17.01622924],
             [17.66858589, 16.48327986, 22.40837464, ..., 23.67550019,
              26.07885909, 25.3743376 ],
             [27.47461117, 27.19358915, 26.22077474, ..., 23.55112003,
              24.9480821 , 27.9200462 ],
             [13.44162735, 17.72234442, 19.4001536 , ..., 24.28306676,
              22.14120273, 20.46973503]])

    • lp
      (chain, draw)
      float64
      -13.65 -10.93 ... -17.2 -17.62

      array([[-13.64892316, -10.93191042,  -7.73553408, ..., -13.42255768,
              -12.41458354, -13.38114468],
             [-13.50843047, -13.52754004, -17.05447355, ..., -21.346963  ,
              -20.76992669, -20.68155521],
             [-21.67696795, -24.09496483, -19.65497824, ..., -20.72555624,
              -19.60015713, -21.7556259 ],
             [-10.25840142, -13.99905927, -14.22672354, ..., -18.78044287,
              -17.20034121, -17.61868656]])

  • Attributes: (9)

    created_at :

    2022-04-15T05:26:56.743010

    arviz_version :

    0.12.0

    inference_library :

    stan

    inference_library_version :

    3.4.0

    num_chains :

    4

    num_samples :

    1000

    num_thin :

    1

    num_warmup :

    1000

    save_warmup :

    0

  
  

az.plot_pair(
    data,
    coords={"school": ["Choate", "Deerfield", "Phillips Andover"]},
    divergences=True,
);

See also

working_with_InferenceData