arviz.extract

arviz.extract(data, group='posterior', combined=True, var_names=None, filter_vars=None, num_samples=None, keep_dataset=False, rng=None)

Extract an InferenceData group or subset of it.

Parameters

idataInferenceData or InferenceData_like

InferenceData from which to extract the data.

groupstr, optional

Which InferenceData data group to extract data from.

combinedbool, optional

Combine chain and draw dimensions into sample. Won’t work if a dimension named sample already exists.

var_namesstr or list of str, optional

Variables to be extracted. Prefix the variables by ~ when you want to exclude them.

filter_vars: {None, “like”, “regex”}, optional

If None (default), interpret var_names as the real variables names. If “like”, interpret var_names as substrings of the real variables names. If “regex”, interpret var_names as regular expressions on the real variables names. A la pandas.filter. Like with plotting, sometimes it’s easier to subset saying what to exclude instead of what to include

num_samplesint, optional

Extract only a subset of the samples. Only valid if combined=True

keep_datasetbool, optional

If true, always return a DataSet. If false (default) return a DataArray when there is a single variable.

rngbool, int, numpy.Generator, optional

Shuffle the samples, only valid if combined=True. By default, samples are shuffled if num_samples is not None, and are left in the same order otherwise. This ensures that subsetting the samples doesn’t return only samples from a single chain and consecutive draws.

Returns

xarray.DataArray or xarray.Dataset

Examples

The default behaviour is to return the posterior group after stacking the chain and draw dimensions.

import arviz as az
idata = az.load_arviz_data("centered_eight")
az.extract(idata)

<xarray.Dataset>
Dimensions:  (sample: 2000, school: 8)
Coordinates:
  * school   (school) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  * sample   (sample) object MultiIndex
  * chain    (sample) int64 0 0 0 0 0 0 0 0 0 0 0 0 ... 3 3 3 3 3 3 3 3 3 3 3 3
  * draw     (sample) int64 0 1 2 3 4 5 6 7 ... 492 493 494 495 496 497 498 499
Data variables:
    mu       (sample) float64 7.872 3.385 9.1 7.304 ... 1.859 1.767 3.486 3.404
    theta    (school, sample) float64 12.32 11.29 5.709 ... -2.623 8.452 1.295
    tau      (sample) float64 4.726 3.909 4.844 1.857 ... 2.741 2.932 4.461
Attributes:
    created_at:                 2022-10-13T14:37:37.315398
    arviz_version:              0.13.0.dev0
    inference_library:          pymc
    inference_library_version:  4.2.2
    sampling_time:              7.480114936828613
    tuning_steps:               1000

xarray.Dataset

• Dimensions:
  • sample: 2000
  • school: 8
• Coordinates: (4)
  • school
    (school)
    object
    'Choate' ... 'Mt. Hermon'

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

  • sample
    (sample)
    object
    MultiIndex

    array([(0, 0), (0, 1), (0, 2), ..., (3, 497), (3, 498), (3, 499)], dtype=object)

  • chain
    (sample)
    int64
    0 0 0 0 0 0 0 0 ... 3 3 3 3 3 3 3 3

    array([0, 0, 0, ..., 3, 3, 3])

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

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

• Data variables: (3)
  • mu
    (sample)
    float64
    7.872 3.385 9.1 ... 3.486 3.404

    array([7.87179637, 3.38455431, 9.10047569, ..., 1.76673325, 3.48611194,
           3.40446391])

  • theta
    (school, sample)
    float64
    12.32 11.29 5.709 ... 8.452 1.295

    array([[12.32068558, 11.28562322,  5.70850575, ...,  3.53251469,
             4.1827505 ,  0.19295578],
           [ 9.90536689,  9.12932369,  5.7579324 , ...,  2.00890089,
             7.55425055,  6.4984278 ],
           [14.9516155 ,  3.13926327, 10.94458539, ...,  0.51080602,
             4.45603444, -0.89442402],
           ...,
           [16.90179529,  2.39308809,  8.14332707, ...,  4.70724881,
             1.5289581 ,  7.93646013],
           [13.19805933, 10.05522282,  7.60475341, ...,  3.07331446,
             1.09609826,  6.76245459],
           [15.06136584,  6.17672421,  8.76764689, ..., -2.62306892,
             8.45228161,  1.2950506 ]])

  • tau
    (sample)
    float64
    4.726 3.909 4.844 ... 2.932 4.461

    array([4.72574006, 3.90899361, 4.8440252 , ..., 2.74060666, 2.93237926,
           4.46124596])

• Indexes: (2)
  • school
    PandasIndex

    PandasIndex(Index(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
           'Hotchkiss', 'Lawrenceville', 'St. Paul's', 'Mt. Hermon'],
          dtype='object', name='school'))

  • sample
    chain
    draw
    PandasMultiIndex

    PandasIndex(MultiIndex([(0,   0),
                (0,   1),
                (0,   2),
                (0,   3),
                (0,   4),
                (0,   5),
                (0,   6),
                (0,   7),
                (0,   8),
                (0,   9),
                ...
                (3, 490),
                (3, 491),
                (3, 492),
                (3, 493),
                (3, 494),
                (3, 495),
                (3, 496),
                (3, 497),
                (3, 498),
                (3, 499)],
               name='sample', length=2000))

• Attributes: (6)

  created_at :

  2022-10-13T14:37:37.315398

  arviz_version :

  0.13.0.dev0

  inference_library :

  pymc

  inference_library_version :

  4.2.2

  sampling_time :

  7.480114936828613

  tuning_steps :

  1000

You can also indicate a subset to be returned, but in variables and in samples:

az.extract(idata, var_names="theta", num_samples=100)

<xarray.DataArray 'theta' (school: 8, sample: 100)>
array([[ 1.16356393e+01,  9.80188050e+00,  6.93102223e+00,
         4.40479805e+00,  3.37453765e+00,  1.66027661e+00,
         6.88803121e+00,  7.68194332e+00,  3.42527056e+00,
         3.24161727e+00,  9.76390798e+00,  2.66587615e+00,
         1.25933432e+01,  1.20017602e+01, -2.56500961e+00,
        -1.19450497e+00,  3.16324840e-01, -4.35165446e-01,
        -1.45575080e+00,  7.57620852e+00,  7.81921367e-01,
         9.80188050e+00,  2.61377333e-01, -5.10846749e+00,
         6.25495846e+00,  9.58014932e+00, -1.32636459e+00,
         8.27358881e-01,  1.92743101e+01,  9.85349670e-01,
         2.67717800e+00,  9.80188050e+00, -8.82173225e-01,
         4.27281294e+00,  9.71751869e+00,  5.21907152e+00,
         3.23315640e+00,  7.25006163e+00,  1.35503511e+01,
         1.12856232e+01,  3.88495610e+00, -6.98493081e-01,
         1.32968861e+01,  9.14079559e+00,  6.90074411e+00,
        -5.37572195e-01,  3.25808233e+00, -6.24775331e-01,
         9.76585447e+00,  9.88063755e+00,  1.04382282e+01,
         1.47359226e+01,  7.34034281e+00,  1.59326273e+00,
         2.37221256e+00,  5.29726546e-01,  1.31022009e+01,
        -1.33919123e-02,  5.36521047e+00, -3.75264208e+00,
...
         2.76188750e+00,  1.13577324e+01, -7.89320888e-01,
         5.26780466e+00,  3.35132225e+00, -3.55332084e+00,
         5.44159038e+00,  5.04095986e+00,  8.86523760e+00,
         1.25813350e+01,  4.83350971e+00,  2.22113452e+00,
         4.58559411e+00,  1.94926476e+00,  1.58236330e+01,
         2.86838389e-01,  6.99697482e+00,  2.27260207e+00,
        -2.26780401e+00,  5.83475486e-01,  6.55331619e+00,
        -3.15352465e+00,  5.49126949e+00,  4.12192432e+00,
         3.35132225e+00,  3.90203385e+00,  4.11797530e+00,
         9.31886541e+00,  9.03110639e+00,  7.63726943e+00,
         1.59962412e+00,  3.58661575e+00,  9.12635811e+00,
         3.47889525e+00,  3.94770392e+00,  4.47140798e-03,
         5.24838051e+00,  7.03611228e+00, -2.23165528e+00,
         1.84003191e+00,  5.69639395e+00,  5.33233277e+00,
         9.98761823e+00,  6.99697482e+00,  3.35132225e+00,
         4.05865793e+00,  1.05857124e+01, -9.67889623e+00,
         1.30644225e+00,  1.97055711e+00, -4.34713310e+00,
         2.09180643e+00,  4.40295366e+00,  9.07594450e+00,
         6.25408810e-01,  2.73537789e+00,  7.98394952e+00,
         7.45065714e+00]])
Coordinates:
  * school   (school) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  * sample   (sample) object MultiIndex
  * chain    (sample) int64 0 3 0 1 2 2 0 0 1 2 2 3 ... 1 1 0 2 2 0 0 1 0 3 2 2
  * draw     (sample) int64 188 28 456 337 350 328 ... 323 277 196 275 474 434

xarray.DataArray

'theta'

• school: 8
• sample: 100

• 11.64 9.802 6.931 4.405 3.375 1.66 ... 9.076 0.6254 2.735 7.984 7.451

  array([[ 1.16356393e+01,  9.80188050e+00,  6.93102223e+00,
           4.40479805e+00,  3.37453765e+00,  1.66027661e+00,
           6.88803121e+00,  7.68194332e+00,  3.42527056e+00,
           3.24161727e+00,  9.76390798e+00,  2.66587615e+00,
           1.25933432e+01,  1.20017602e+01, -2.56500961e+00,
          -1.19450497e+00,  3.16324840e-01, -4.35165446e-01,
          -1.45575080e+00,  7.57620852e+00,  7.81921367e-01,
           9.80188050e+00,  2.61377333e-01, -5.10846749e+00,
           6.25495846e+00,  9.58014932e+00, -1.32636459e+00,
           8.27358881e-01,  1.92743101e+01,  9.85349670e-01,
           2.67717800e+00,  9.80188050e+00, -8.82173225e-01,
           4.27281294e+00,  9.71751869e+00,  5.21907152e+00,
           3.23315640e+00,  7.25006163e+00,  1.35503511e+01,
           1.12856232e+01,  3.88495610e+00, -6.98493081e-01,
           1.32968861e+01,  9.14079559e+00,  6.90074411e+00,
          -5.37572195e-01,  3.25808233e+00, -6.24775331e-01,
           9.76585447e+00,  9.88063755e+00,  1.04382282e+01,
           1.47359226e+01,  7.34034281e+00,  1.59326273e+00,
           2.37221256e+00,  5.29726546e-01,  1.31022009e+01,
          -1.33919123e-02,  5.36521047e+00, -3.75264208e+00,
  ...
           2.76188750e+00,  1.13577324e+01, -7.89320888e-01,
           5.26780466e+00,  3.35132225e+00, -3.55332084e+00,
           5.44159038e+00,  5.04095986e+00,  8.86523760e+00,
           1.25813350e+01,  4.83350971e+00,  2.22113452e+00,
           4.58559411e+00,  1.94926476e+00,  1.58236330e+01,
           2.86838389e-01,  6.99697482e+00,  2.27260207e+00,
          -2.26780401e+00,  5.83475486e-01,  6.55331619e+00,
          -3.15352465e+00,  5.49126949e+00,  4.12192432e+00,
           3.35132225e+00,  3.90203385e+00,  4.11797530e+00,
           9.31886541e+00,  9.03110639e+00,  7.63726943e+00,
           1.59962412e+00,  3.58661575e+00,  9.12635811e+00,
           3.47889525e+00,  3.94770392e+00,  4.47140798e-03,
           5.24838051e+00,  7.03611228e+00, -2.23165528e+00,
           1.84003191e+00,  5.69639395e+00,  5.33233277e+00,
           9.98761823e+00,  6.99697482e+00,  3.35132225e+00,
           4.05865793e+00,  1.05857124e+01, -9.67889623e+00,
           1.30644225e+00,  1.97055711e+00, -4.34713310e+00,
           2.09180643e+00,  4.40295366e+00,  9.07594450e+00,
           6.25408810e-01,  2.73537789e+00,  7.98394952e+00,
           7.45065714e+00]])

• Coordinates: (4)
  • school
    (school)
    object
    'Choate' ... 'Mt. Hermon'

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

  • sample
    (sample)
    object
    MultiIndex

    array([(0, 188), (3, 28), (0, 456), (1, 337), (2, 350), (2, 328), (0, 485),
           (0, 417), (1, 163), (2, 45), (2, 281), (3, 471), (0, 355), (1, 8),
           (3, 263), (2, 65), (0, 60), (0, 477), (3, 257), (3, 308), (0, 78),
           (3, 31), (3, 16), (2, 112), (2, 192), (0, 176), (2, 64), (0, 71),
           (3, 426), (1, 55), (1, 206), (3, 29), (2, 100), (1, 264), (2, 421),
           (0, 169), (2, 37), (3, 201), (1, 322), (0, 1), (1, 420), (3, 250),
           (3, 7), (1, 29), (1, 189), (2, 276), (1, 445), (2, 334), (0, 183),
           (1, 105), (0, 299), (0, 5), (2, 7), (1, 436), (0, 18), (2, 336),
           (1, 62), (0, 446), (3, 124), (3, 233), (0, 48), (2, 38), (2, 435),
           (2, 450), (1, 493), (3, 380), (1, 467), (0, 381), (3, 357), (1, 368),
           (2, 204), (0, 316), (0, 81), (2, 72), (1, 280), (1, 159), (2, 208),
           (0, 394), (0, 398), (1, 85), (2, 280), (1, 134), (3, 82), (1, 174),
           (0, 499), (3, 111), (1, 456), (3, 220), (1, 400), (1, 84), (0, 282),
           (2, 225), (2, 105), (0, 69), (0, 323), (1, 277), (0, 196), (3, 275),
           (2, 474), (2, 434)], dtype=object)

  • chain
    (sample)
    int64
    0 3 0 1 2 2 0 0 ... 2 0 0 1 0 3 2 2

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

  • draw
    (sample)
    int64
    188 28 456 337 ... 196 275 474 434

    array([188,  28, 456, 337, 350, 328, 485, 417, 163,  45, 281, 471, 355,   8,
           263,  65,  60, 477, 257, 308,  78,  31,  16, 112, 192, 176,  64,  71,
           426,  55, 206,  29, 100, 264, 421, 169,  37, 201, 322,   1, 420, 250,
             7,  29, 189, 276, 445, 334, 183, 105, 299,   5,   7, 436,  18, 336,
            62, 446, 124, 233,  48,  38, 435, 450, 493, 380, 467, 381, 357, 368,
           204, 316,  81,  72, 280, 159, 208, 394, 398,  85, 280, 134,  82, 174,
           499, 111, 456, 220, 400,  84, 282, 225, 105,  69, 323, 277, 196, 275,
           474, 434])

• Indexes: (2)
  • school
    PandasIndex

    PandasIndex(Index(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
           'Hotchkiss', 'Lawrenceville', 'St. Paul's', 'Mt. Hermon'],
          dtype='object', name='school'))

  • sample
    chain
    draw
    PandasMultiIndex

    PandasIndex(MultiIndex([(0, 188),
                (3,  28),
                (0, 456),
                (1, 337),
                (2, 350),
                (2, 328),
                (0, 485),
                (0, 417),
                (1, 163),
                (2,  45),
                (2, 281),
                (3, 471),
                (0, 355),
                (1,   8),
                (3, 263),
                (2,  65),
                (0,  60),
                (0, 477),
                (3, 257),
                (3, 308),
                (0,  78),
                (3,  31),
                (3,  16),
                (2, 112),
                (2, 192),
                (0, 176),
                (2,  64),
                (0,  71),
                (3, 426),
                (1,  55),
                (1, 206),
                (3,  29),
                (2, 100),
                (1, 264),
                (2, 421),
                (0, 169),
                (2,  37),
                (3, 201),
                (1, 322),
                (0,   1),
                (1, 420),
                (3, 250),
                (3,   7),
                (1,  29),
                (1, 189),
                (2, 276),
                (1, 445),
                (2, 334),
                (0, 183),
                (1, 105),
                (0, 299),
                (0,   5),
                (2,   7),
                (1, 436),
                (0,  18),
                (2, 336),
                (1,  62),
                (0, 446),
                (3, 124),
                (3, 233),
                (0,  48),
                (2,  38),
                (2, 435),
                (2, 450),
                (1, 493),
                (3, 380),
                (1, 467),
                (0, 381),
                (3, 357),
                (1, 368),
                (2, 204),
                (0, 316),
                (0,  81),
                (2,  72),
                (1, 280),
                (1, 159),
                (2, 208),
                (0, 394),
                (0, 398),
                (1,  85),
                (2, 280),
                (1, 134),
                (3,  82),
                (1, 174),
                (0, 499),
                (3, 111),
                (1, 456),
                (3, 220),
                (1, 400),
                (1,  84),
                (0, 282),
                (2, 225),
                (2, 105),
                (0,  69),
                (0, 323),
                (1, 277),
                (0, 196),
                (3, 275),
                (2, 474),
                (2, 434)],
               name='sample'))

• Attributes: (0)

To keep the chain and draw dimensions, use combined=False.

az.extract(idata, group="prior", combined=False)

<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) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
Data variables:
    tau      (chain, draw) float64 ...
    theta    (chain, draw, school) float64 ...
    mu       (chain, draw) float64 ...
Attributes:
    arviz_version:              0.13.0.dev0
    created_at:                 2022-10-13T14:37:26.602116
    inference_library:          pymc
    inference_library_version:  4.2.2

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)
    object
    'Choate' ... 'Mt. Hermon'

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

• Data variables: (3)
  • tau
    (chain, draw)
    float64
    ...

    [500 values with dtype=float64]

  • theta
    (chain, draw, school)
    float64
    ...

    [4000 values with dtype=float64]

  • mu
    (chain, draw)
    float64
    ...

    [500 values with dtype=float64]

• Indexes: (3)
  • chain
    PandasIndex

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

  • draw
    PandasIndex

    PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
           ...
           490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
          dtype='int64', name='draw', length=500))

  • school
    PandasIndex

    PandasIndex(Index(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
           'Hotchkiss', 'Lawrenceville', 'St. Paul's', 'Mt. Hermon'],
          dtype='object', name='school'))

• Attributes: (4)

  arviz_version :

  0.13.0.dev0

  created_at :

  2022-10-13T14:37:26.602116

  inference_library :

  pymc

  inference_library_version :

  4.2.2