arviz.extract_dataset

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

Extract an InferenceData group or subset of it as a xarray.Dataset.

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 plotted, two variables are required. Prefix the variables by ~ when you want to exclude them from the plot.

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

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.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_dataset(idata)

<xarray.Dataset>
Dimensions:  (school: 8, sample: 2000)
Coordinates:
  * school   (school) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  * sample   (sample) MultiIndex
  - chain    (sample) int64 0 0 0 0 0 0 0 0 0 0 0 0 0 3 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 -3.477 -2.456 -2.826 -1.996 ... 4.597 5.899 0.1614
    theta    (school, sample) float64 1.669 -6.239 2.195 ... -1.095 4.013 4.523
    tau      (sample) float64 3.73 2.075 3.703 4.146 ... 8.589 8.346 7.711 5.407
Attributes:
    created_at:                 2019-06-21T17:36:34.398087
    inference_library:          pymc3
    inference_library_version:  3.7

xarray.Dataset

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

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

  • sample
    (sample)
    MultiIndex
    (chain, draw)

    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
    -3.477 -2.456 ... 5.899 0.1614

    array([-3.47698606, -2.45587061, -2.82625433, ...,  4.59705819,
            5.89850592,  0.16138927])

  • theta
    (school, sample)
    float64
    1.669 -6.239 2.195 ... 4.013 4.523

    array([[ 1.66865435, -6.23935854,  2.19509753, ..., 14.18629932,
            -1.42094639, -0.05015941],
           [-8.53740136,  1.0714115 , -2.9208433 , ..., 11.76271563,
            -4.0344055 ,  0.06353824],
           [-2.622619  , -4.51792731, -3.49420117, ...,  4.68403648,
            16.62006419,  7.60845423],
           ...,
           [-9.17401935,  0.39954568, -2.72030236, ...,  3.35749116,
            -4.90192473, -1.10614164],
           [ 0.15523406, -4.46252755, -3.4870049 , ..., -0.58651408,
            15.85064806, 10.59293283],
           [-6.81825139, -1.11076126, -1.02643024, ..., -1.09507072,
             4.01339719,  4.52338885]])

  • tau
    (sample)
    float64
    3.73 2.075 3.703 ... 7.711 5.407

    array([3.7301011 , 2.07538342, 3.70299344, ..., 8.34563098, 7.71079012,
           5.40679834])

• Attributes: (3)

  created_at :

  2019-06-21T17:36:34.398087

  inference_library :

  pymc3

  inference_library_version :

  3.7

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

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

<xarray.Dataset>
Dimensions:  (school: 8, sample: 100)
Coordinates:
  * school   (school) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
  * sample   (sample) MultiIndex
  - chain    (sample) int64 2 1 1 1 3 0 0 0 3 1 0 1 ... 1 1 1 3 1 2 1 0 1 2 0 2
  - draw     (sample) int64 21 292 439 86 441 11 209 ... 156 92 168 2 180 278
Data variables:
    theta    (school, sample) float64 1.609 3.862 1.92 ... 6.226 4.652 8.521
Attributes:
    created_at:                 2019-06-21T17:36:34.398087
    inference_library:          pymc3
    inference_library_version:  3.7

xarray.Dataset

• Dimensions:
  • school: 8
  • sample: 100
• Coordinates: (2)
  • school
    (school)
    object
    'Choate' ... 'Mt. Hermon'

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

  • sample
    (sample)
    MultiIndex
    (chain, draw)

    array([(2, 21), (1, 292), (1, 439), (1, 86), (3, 441), (0, 11), (0, 209),
           (0, 230), (3, 298), (1, 260), (0, 302), (1, 55), (1, 65), (0, 16),
           (3, 414), (0, 487), (1, 136), (3, 29), (3, 48), (2, 462), (2, 375),
           (1, 362), (2, 449), (1, 402), (3, 382), (1, 46), (2, 323), (1, 463),
           (3, 489), (1, 461), (3, 33), (3, 172), (2, 458), (1, 274), (2, 233),
           (0, 72), (1, 338), (1, 239), (2, 20), (1, 498), (2, 224), (0, 470),
           (2, 154), (1, 296), (3, 167), (2, 216), (1, 375), (0, 214), (2, 75),
           (2, 68), (3, 81), (0, 472), (0, 390), (1, 34), (1, 334), (2, 96),
           (0, 177), (3, 120), (2, 380), (2, 163), (1, 479), (2, 443), (1, 71),
           (1, 146), (1, 78), (0, 344), (3, 341), (2, 241), (1, 79), (0, 221),
           (1, 88), (2, 322), (0, 299), (0, 351), (2, 41), (3, 350), (1, 20),
           (1, 356), (2, 170), (2, 12), (2, 63), (1, 470), (0, 170), (0, 313),
           (3, 321), (2, 189), (3, 126), (1, 51), (1, 91), (1, 174), (1, 126),
           (3, 72), (1, 160), (2, 276), (1, 156), (0, 92), (1, 168), (2, 2),
           (0, 180), (2, 278)], dtype=object)

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

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

  • draw
    (sample)
    int64
    21 292 439 86 441 ... 168 2 180 278

    array([ 21, 292, 439,  86, 441,  11, 209, 230, 298, 260, 302,  55,  65,  16,
           414, 487, 136,  29,  48, 462, 375, 362, 449, 402, 382,  46, 323, 463,
           489, 461,  33, 172, 458, 274, 233,  72, 338, 239,  20, 498, 224, 470,
           154, 296, 167, 216, 375, 214,  75,  68,  81, 472, 390,  34, 334,  96,
           177, 120, 380, 163, 479, 443,  71, 146,  78, 344, 341, 241,  79, 221,
            88, 322, 299, 351,  41, 350,  20, 356, 170,  12,  63, 470, 170, 313,
           321, 189, 126,  51,  91, 174, 126,  72, 160, 276, 156,  92, 168,   2,
           180, 278])

• Data variables: (1)
  • theta
    (school, sample)
    float64
    1.609 3.862 1.92 ... 4.652 8.521

    array([[ 1.60947010e+00,  3.86220366e+00,  1.91987282e+00,
             8.82747613e+00, -4.35899109e+00,  6.46381152e+00,
             7.01145928e+00,  5.57430393e+00,  8.95113097e+00,
            -4.42034200e+00, -6.92380205e-01,  9.20767531e+00,
             7.37422510e+00,  5.41452506e-01,  2.66031654e+00,
             2.45837957e-01,  1.23545262e+01,  2.01866402e+01,
             5.81663070e+00,  8.50640830e+00,  2.75467802e+00,
             1.22599853e+01, -1.22202607e+00,  1.05999917e+01,
             7.34972366e+00,  1.24706824e+01,  5.74349870e-01,
             2.69174280e+01,  9.41245303e+00,  1.32075065e+01,
             3.81408610e+00,  2.16185611e+01,  8.99625990e+00,
             1.08054481e+01,  7.61543508e+00,  3.62243774e-02,
             1.13179551e+01,  5.81318657e+00,  1.25674436e+01,
             1.47355006e+01,  3.16876956e+00, -1.92393695e+00,
             2.30718848e+00,  6.20856882e+00,  2.12712371e+01,
             8.41368745e+00,  3.72479161e+00,  9.97375373e+00,
             1.00856814e+01,  1.61280425e+01,  7.20697719e+00,
             7.67092618e+00,  5.82779748e+00,  8.07394570e+00,
             9.89112465e+00,  1.73021909e+01,  8.86478392e+00,
             5.71545586e+00,  1.65462811e+00,  9.18939731e+00,
    ...
             3.95812734e+00,  8.96002329e+00,  1.20800920e+01,
             3.80102237e+00, -2.67017929e+00,  1.18843286e+01,
             4.78436309e+00, -1.43584099e+00, -2.40389662e+00,
             3.70209177e+00,  6.92568215e+00,  5.33327355e+00,
             7.35998825e+00,  1.98826185e+01,  7.39880466e+00,
             8.39945054e+00,  1.74458783e+00,  9.63995576e+00,
             3.00499476e-01,  5.22128129e+00,  4.28686247e+00,
             7.06564133e+00,  1.46244316e+00,  1.44790884e+00,
             3.61811574e+00, -2.44105834e+00,  1.84658703e+00,
             4.02676123e+00,  5.47905110e+00, -3.47680716e+00,
             1.90317384e+00,  1.44790884e+00, -2.32220043e+00,
             2.08761951e+00,  8.90871247e+00,  5.36209092e+00,
            -3.92643267e+00,  1.58236733e+00,  4.19363612e+00,
             6.43494803e+00,  6.57849583e+00,  5.94367026e+00,
             5.67085820e-01,  1.01591078e+01,  1.10237378e+01,
             1.39292642e+00,  1.35422312e+00,  1.13649937e+01,
             5.30899651e+00,  6.80481785e+00, -8.28402578e+00,
             8.52082240e+00,  3.86996099e+00,  4.58517229e-01,
             1.34176528e+01,  6.22563869e+00,  4.65187937e+00,
             8.52082240e+00]])

• Attributes: (3)

  created_at :

  2019-06-21T17:36:34.398087

  inference_library :

  pymc3

  inference_library_version :

  3.7

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

az.extract_dataset(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 ... 492 493 494 495 496 497 498 499
  * school     (school) object 'Choate' 'Deerfield' ... 'Mt. Hermon'
Data variables:
    tau        (chain, draw) float64 6.561 1.016 68.91 ... 1.56 5.949 0.7631
    tau_log__  (chain, draw) float64 1.881 0.01593 4.233 ... 1.783 -0.2704
    mu         (chain, draw) float64 5.293 0.8137 0.7122 ... -1.658 -3.273
    theta      (chain, draw, school) float64 2.357 7.371 7.251 ... -3.775 -3.555
    obs        (chain, draw, school) float64 -3.54 6.769 19.68 ... -21.16 -6.071
Attributes:
    created_at:                 2019-06-21T17:36:34.490387
    inference_library:          pymc3
    inference_library_version:  3.7

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: (5)
  • tau
    (chain, draw)
    float64
    ...

    array([[ 6.560633,  1.016055, 68.91391 , ...,  1.560098,  5.948734,  0.763063]])

  • tau_log__
    (chain, draw)
    float64
    ...

    array([[ 1.881087,  0.015927,  4.232858, ...,  0.444748,  1.783178, -0.270415]])

  • mu
    (chain, draw)
    float64
    ...

    array([[ 5.29345 ,  0.813724,  0.712223, ..., -0.979857, -1.657547, -3.272668]])

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

    array([[[ 2.357357,  7.371371, ...,  6.135082,  3.984435],
            [ 0.258399, -0.752515, ...,  1.73084 , -0.034163],
            ...,
            [-4.353289,  2.194643, ..., -7.819076, -6.21613 ],
            [-4.131344, -4.093318, ..., -3.775218, -3.555126]]])

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

    array([[[ -3.539971,   6.769448, ...,   8.26964 ,  -8.569042],
            [-21.166369,   1.14605 , ..., -13.157913,  -8.5424  ],
            ...,
            [ 29.354582,  -5.511382, ..., -17.892521,  46.28878 ],
            [ -6.379747,   6.538907, ..., -21.155214,  -6.070767]]])

• Attributes: (3)

  created_at :

  2019-06-21T17:36:34.490387

  inference_library :

  pymc3

  inference_library_version :

  3.7

previous

arviz.concat

next

arviz.list_datasets