arviz_base.extract

arviz_base.extract(data, group='posterior', sample_dims=None, *, combined=True, var_names=None, filter_vars=None, num_samples=None, weights=None, resampling_method=None, keep_dataset=False, random_seed=None)

Extract a group or group subset from a DataTree.

Parameters:

idataDataTree-like

DataTree from which to extract the data.

groupstr, optional

Which group to extract data from.

sample_dimssequence of hashable, optional

List of dimensions that should be considered sampling dimensions. Random subsets and potential stacking if combine=True happen over these dimensions only. Defaults to rcParams["data.sample_dims"].

combinedbool, optional

Combine sample_dims dimensions into sample. Won’t work if a dimension named sample already exists. It is irrelevant and ignored when sample_dims is a single dimension.

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 or sample_dims represents a single dimension.

weightsarray_like, optional

Extract a weighted subset of the samples. Only valid if num_samples is not None.

resampling_methodstr, optional

Method to use for resampling. Default is “multinomial”. Options are “multinomial” and “stratified”. For stratified resampling, weights must be provided. Default is “stratified” if weights are provided, “multinomial” otherwise.

keep_datasetbool, optional

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

random_seedint, numpy.Generator, optional

Random number generator or seed. Only used if weights is not None or if num_samples is not None.

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_base as az
idata = az.load_arviz_data("centered_eight")
az.extract(idata)

<xarray.Dataset> Size: 209kB
Dimensions:  (sample: 2000, school: 8)
Coordinates:
  * sample   (sample) object 16kB MultiIndex
  * chain    (sample) int64 16kB 0 0 0 0 0 0 0 0 0 0 0 ... 3 3 3 3 3 3 3 3 3 3 3
  * draw     (sample) int64 16kB 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
  * school   (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'
Data variables:
    mu       (sample) float64 16kB 1.716 1.903 1.903 1.903 ... 5.409 7.721 10.24
    theta    (school, sample) float64 128kB 2.317 0.8892 0.8892 ... 9.754 14.02
    tau      (sample) float64 16kB 0.8775 0.8027 0.8027 ... 2.236 2.99 3.052
Attributes:
    created_at:                 2025-01-19T14:32:33.071271+00:00
    arviz_version:              0.20.0
    inference_library:          pymc
    inference_library_version:  5.20.0
    sampling_time:              3.159093141555786
    tuning_steps:               1000

xarray.Dataset

• Dimensions:
  • sample: 2000
  • school: 8
• Coordinates: (4)
  • sample
    (sample)
    object
    MultiIndex

    [2000 values with dtype=object]

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

    [2000 values with dtype=int64]

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

    [2000 values with dtype=int64]

  • 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
    (sample)
    float64
    1.716 1.903 1.903 ... 7.721 10.24

    array([ 1.71572331,  1.90348113,  1.90348113, ...,  5.40883573,
            7.72143998, 10.23715678], shape=(2000,))

  • theta
    (school, sample)
    float64
    2.317 0.8892 0.8892 ... 9.754 14.02

    array([[ 2.31739097,  0.8891702 ,  0.8891702 , ..., 11.40639049,
             7.08613879, 10.46439019],
           [ 1.45017375,  0.74294874,  0.74294874, ...,  4.44693711,
            12.31188923, 13.71430563],
           [ 2.08554978,  3.12586895,  3.12586895, ...,  9.21077455,
             6.5843008 , 10.26166615],
           ...,
           [ 2.712972  ,  1.55893929,  1.55893929, ...,  4.81230178,
            11.85993774, 15.07090015],
           [ 3.08376362,  2.48750253,  2.48750253, ...,  9.69325672,
             7.95226848, 14.92321025],
           [ 1.46044828,  1.98437916,  1.98437916, ...,  4.91465554,
             9.75446752, 14.02312893]], shape=(8, 2000))

  • tau
    (sample)
    float64
    0.8775 0.8027 0.8027 ... 2.99 3.052

    array([0.8774941 , 0.80271439, 0.80271439, ..., 2.23648602, 2.98965551,
           3.05155936], shape=(2000,))

• Attributes: (6)

  created_at :

  2025-01-19T14:32:33.071271+00:00

  arviz_version :

  0.20.0

  inference_library :

  pymc

  inference_library_version :

  5.20.0

  sampling_time :

  3.159093141555786

  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)> Size: 6kB
array([[ 6.68471590e+00,  2.57249107e+00,  2.99172113e+00,
         5.81653080e+00,  1.55119685e+00, -2.13347506e+00,
         8.12191423e+00,  3.78557008e+00,  1.90457029e+01,
         1.17409423e+01,  4.43616091e+00,  2.92038998e+00,
         4.71273898e+00, -7.02366780e-01,  4.27337284e+00,
         1.32758740e+01,  9.37066645e+00,  1.10972038e+01,
         1.91515897e+00,  5.37903339e+00,  1.09675811e+01,
         6.24439453e+00,  5.07493796e+00,  7.87700852e+00,
         6.70987225e+00, -7.95927409e+00, -2.28779404e+00,
         7.16565453e+00,  6.69068199e+00,  3.33481048e+00,
         4.87010661e+00,  2.35939213e+01,  5.87166557e+00,
         1.54030624e+01,  5.88693731e+00,  1.29617378e+01,
         4.29774953e+00,  4.34980010e+00,  1.13150317e+01,
         7.47838234e+00,  1.78722902e+01,  7.20466974e-01,
         2.29337646e+01,  2.05025159e+00,  1.52621772e+00,
         6.28214233e+00,  1.30671870e+00,  1.05316393e+01,
         4.14922272e+00,  1.47220436e+01,  6.69690701e+00,
        -6.20691096e+00,  6.06725954e+00,  2.71076287e+00,
         3.35049267e+00,  1.27021144e+01,  1.19449175e+01,
         1.00552225e+01,  6.28250875e+00,  9.55189677e+00,
...
         1.59686541e+01, -1.06621507e-01,  3.49501735e+00,
        -1.05716230e+01,  4.45165075e+00,  1.11020636e+00,
         9.69629831e+00,  6.21592516e+00,  3.67578666e+00,
         4.18452816e+00,  3.03240922e+00,  5.27501577e+00,
         2.74181051e+00,  7.05670536e+00,  5.81672064e+00,
         1.18808099e+00,  2.60479843e+00,  4.00740765e+00,
        -1.53972551e-01,  9.28109060e+00,  8.99880439e+00,
         3.07043795e+00,  1.46427630e+01,  1.83572572e+01,
        -6.72358848e+00,  6.88936888e+00,  3.82680610e+00,
         9.42741342e+00,  1.10157749e+00,  4.24093460e+00,
        -5.16899739e+00,  7.48235996e+00,  4.33901016e+00,
         1.02765035e+01,  2.64341240e+00,  6.81192833e+00,
         6.95804318e-01,  1.97179248e+00, -8.68273653e-01,
        -2.17777716e+00,  3.78542983e+00,  5.86804626e+00,
         9.40307510e+00,  2.56739760e+00,  1.62283160e+00,
         1.04302129e+01,  3.62829097e+00,  3.00465019e+00,
         6.68384338e+00,  7.86257298e+00, -1.21359955e+00,
         9.87984716e+00,  1.72647854e+00,  7.70648424e+00,
         1.52228819e+01,  7.42756817e+00,  2.88561246e+00,
         6.11788501e+00]])
Coordinates:
  * school   (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'
  * sample   (sample) object 800B MultiIndex
  * chain    (sample) int64 800B 0 0 1 1 1 0 0 3 3 3 0 ... 1 0 3 3 1 0 2 1 1 1 3
  * draw     (sample) int64 800B 28 438 482 406 240 43 ... 390 412 28 364 263

xarray.DataArray

'theta'

• school: 8
• sample: 100

• 6.685 2.572 2.992 5.817 1.551 -2.133 ... 7.706 15.22 7.428 2.886 6.118

  array([[ 6.68471590e+00,  2.57249107e+00,  2.99172113e+00,
           5.81653080e+00,  1.55119685e+00, -2.13347506e+00,
           8.12191423e+00,  3.78557008e+00,  1.90457029e+01,
           1.17409423e+01,  4.43616091e+00,  2.92038998e+00,
           4.71273898e+00, -7.02366780e-01,  4.27337284e+00,
           1.32758740e+01,  9.37066645e+00,  1.10972038e+01,
           1.91515897e+00,  5.37903339e+00,  1.09675811e+01,
           6.24439453e+00,  5.07493796e+00,  7.87700852e+00,
           6.70987225e+00, -7.95927409e+00, -2.28779404e+00,
           7.16565453e+00,  6.69068199e+00,  3.33481048e+00,
           4.87010661e+00,  2.35939213e+01,  5.87166557e+00,
           1.54030624e+01,  5.88693731e+00,  1.29617378e+01,
           4.29774953e+00,  4.34980010e+00,  1.13150317e+01,
           7.47838234e+00,  1.78722902e+01,  7.20466974e-01,
           2.29337646e+01,  2.05025159e+00,  1.52621772e+00,
           6.28214233e+00,  1.30671870e+00,  1.05316393e+01,
           4.14922272e+00,  1.47220436e+01,  6.69690701e+00,
          -6.20691096e+00,  6.06725954e+00,  2.71076287e+00,
           3.35049267e+00,  1.27021144e+01,  1.19449175e+01,
           1.00552225e+01,  6.28250875e+00,  9.55189677e+00,
  ...
           1.59686541e+01, -1.06621507e-01,  3.49501735e+00,
          -1.05716230e+01,  4.45165075e+00,  1.11020636e+00,
           9.69629831e+00,  6.21592516e+00,  3.67578666e+00,
           4.18452816e+00,  3.03240922e+00,  5.27501577e+00,
           2.74181051e+00,  7.05670536e+00,  5.81672064e+00,
           1.18808099e+00,  2.60479843e+00,  4.00740765e+00,
          -1.53972551e-01,  9.28109060e+00,  8.99880439e+00,
           3.07043795e+00,  1.46427630e+01,  1.83572572e+01,
          -6.72358848e+00,  6.88936888e+00,  3.82680610e+00,
           9.42741342e+00,  1.10157749e+00,  4.24093460e+00,
          -5.16899739e+00,  7.48235996e+00,  4.33901016e+00,
           1.02765035e+01,  2.64341240e+00,  6.81192833e+00,
           6.95804318e-01,  1.97179248e+00, -8.68273653e-01,
          -2.17777716e+00,  3.78542983e+00,  5.86804626e+00,
           9.40307510e+00,  2.56739760e+00,  1.62283160e+00,
           1.04302129e+01,  3.62829097e+00,  3.00465019e+00,
           6.68384338e+00,  7.86257298e+00, -1.21359955e+00,
           9.87984716e+00,  1.72647854e+00,  7.70648424e+00,
           1.52228819e+01,  7.42756817e+00,  2.88561246e+00,
           6.11788501e+00]])

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

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

  • sample
    (sample)
    object
    MultiIndex

    [100 values with dtype=object]

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

    [100 values with dtype=int64]

  • draw
    (sample)
    int64
    28 438 482 406 ... 412 28 364 263

    [100 values with dtype=int64]

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

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

<xarray.Dataset> Size: 45kB
Dimensions:  (chain: 1, draw: 500, school: 8)
Coordinates:
  * chain    (chain) int64 8B 0
  * draw     (draw) int64 4kB 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
  * school   (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'
Data variables:
    theta    (chain, draw, school) float64 32kB -8.435 24.12 ... 54.57 52.29
    tau      (chain, draw) float64 4kB 11.93 17.76 4.732 ... 2.231 3.319 93.69
    mu       (chain, draw) float64 4kB 4.714 3.853 1.709 ... -2.245 -2.435
Attributes:
    created_at:                 2025-01-19T14:32:29.212688+00:00
    arviz_version:              0.20.0
    inference_library:          pymc
    inference_library_version:  5.20.0

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], shape=(500,))

  • 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)
  • theta
    (chain, draw, school)
    float64
    -8.435 24.12 -2.102 ... 54.57 52.29

    array([[[ -8.43479341,  24.11879974,  -2.10195465, ...,  -1.70694855,
               1.25196206,  -3.40573897],
            [  7.42603349,  23.52089082,  11.05302151, ..., -15.93784847,
               2.66888813,  16.38098276],
            [  1.04915511,  -0.39067339,   7.76801253, ...,  -1.16853509,
               4.56453375,   6.83548422],
            ...,
            [ -5.55254317,  -0.97569016,  -4.23918563, ...,  -2.79594888,
              -5.25232216,  -3.59953474],
            [ -4.48610579,   3.77498029,  -7.49497175, ...,  -5.74934752,
              -9.98939404,  -3.81731412],
            [-76.14397452, -17.75933052,  12.60574946, ...,  -3.30466192,
              54.56836811,  52.29205731]]], shape=(1, 500, 8))

  • tau
    (chain, draw)
    float64
    11.93 17.76 4.732 ... 3.319 93.69

    array([[1.19322460e+01, 1.77566998e+01, 4.73212149e+00, 2.16132188e+00,
            5.57015912e+00, 1.04851175e+01, 2.19971852e+02, 1.64780625e+02,
            2.10307370e+00, 8.22994287e+00, 1.71883239e-01, 1.43799798e+01,
            3.71794967e+00, 2.19182845e+01, 3.22629960e+02, 3.29922564e+00,
            6.69658595e+00, 1.23762267e+00, 2.61458413e+00, 5.15602170e+00,
            9.30222177e+00, 7.67626550e+01, 2.24244028e+01, 2.03757828e+01,
            1.47167831e-01, 5.04831050e+01, 3.96713841e+00, 3.33037092e+00,
            2.10355017e+01, 5.64329632e-01, 2.92851927e+01, 4.86251535e+00,
            1.17373746e+00, 4.76050639e+00, 8.53913195e+00, 2.47705238e+00,
            1.71943641e+01, 1.17424729e+00, 2.16860748e-01, 1.24681716e+01,
            1.84994793e+01, 3.00572482e-01, 4.93071048e+02, 2.03439781e+01,
            4.56339479e+00, 4.74418776e+01, 2.16490353e+00, 4.28303757e-01,
            2.99704232e-01, 1.69271358e+00, 3.24505034e+00, 1.39050629e+01,
            1.12994181e+02, 2.80980946e+00, 1.00634123e+01, 5.65106207e+00,
            9.26796894e+01, 1.93683523e-01, 6.29547490e+00, 4.82283938e+00,
            9.46578025e+01, 9.10546723e-01, 2.32891011e+01, 4.11572708e+00,
            8.32026071e+00, 1.37395405e+01, 1.38413755e+00, 1.55702522e+01,
            4.54943794e-01, 1.50397012e+00, 4.20255065e-01, 5.18204782e-01,
            3.31773610e+00, 1.39206811e+00, 4.90588326e+00, 5.89085606e+00,
            2.63364930e+00, 9.14185200e-01, 1.35276087e+00, 2.95616485e-01,
    ...
            1.52922694e+01, 6.16450636e+01, 1.33352851e+00, 2.49794513e+02,
            1.94919214e+01, 5.24205734e+03, 1.54788325e+01, 7.13263945e+00,
            4.76592936e+00, 1.46147744e+00, 3.90488879e+01, 3.70721793e-01,
            3.38793333e+01, 4.17562743e+00, 6.60334061e+00, 7.43319504e+00,
            3.11836307e+00, 8.58036190e+00, 2.00649133e+00, 8.58370645e+00,
            1.13492459e+01, 5.22163907e+00, 1.38400163e+01, 1.98471540e+00,
            1.77330602e+00, 5.87338404e+01, 5.90409977e+00, 4.56269749e+00,
            3.11700746e+00, 8.85169497e+01, 4.86723371e+00, 2.78797683e+00,
            2.79966137e+00, 1.80177121e+01, 6.99530549e+00, 5.02858615e-01,
            1.10270287e+01, 2.24329036e+00, 3.29851707e+00, 3.03650097e+00,
            8.73411682e+00, 2.41052280e+01, 4.02762874e+00, 5.11117057e-01,
            2.92595859e+00, 4.48875365e+00, 2.24202338e+00, 1.97857537e+00,
            1.29892809e+01, 1.50965825e+01, 3.62772871e+00, 4.74486601e+00,
            7.06815559e+00, 4.22735035e+00, 3.45605723e+00, 2.69918973e-01,
            3.14868764e+00, 1.49207938e+00, 2.87394279e+01, 8.97828161e+00,
            5.36977910e+00, 9.39655114e-01, 3.98205727e+00, 1.55118468e+00,
            2.66358248e+00, 2.33890554e+00, 1.08287688e+01, 2.64621753e+00,
            2.33804822e+01, 7.40617929e+00, 4.12230945e-01, 2.27288462e+00,
            1.12266584e+00, 1.81114811e+01, 1.35868276e+00, 9.71133532e+00,
            8.62411767e+00, 2.23056304e+00, 3.31874334e+00, 9.36923606e+01]])

  • mu
    (chain, draw)
    float64
    4.714 3.853 1.709 ... -2.245 -2.435

    array([[ 4.71433677e+00,  3.85251933e+00,  1.70896382e+00,
             1.59049845e+00, -2.97655593e+00, -3.28385699e+00,
            -8.89306943e-01,  4.68764431e+00, -1.19103766e+00,
            -7.60510660e+00,  4.15990175e+00, -4.45677605e+00,
             9.23138508e+00, -9.81360739e+00,  6.25962477e-01,
            -5.33796788e+00, -8.08727065e-01, -6.46692117e+00,
             8.14474118e+00, -2.46931047e+00, -3.66738899e+00,
             1.48283847e+00,  4.92928574e+00,  2.49849984e+00,
             2.91064110e+00,  4.69223639e-01,  3.84094459e-01,
            -1.26251993e+00,  1.77169222e-02, -4.35815216e+00,
             6.01055775e-01,  3.18989407e+00,  2.60966829e+00,
            -6.99259146e+00, -7.33583378e+00,  2.32429068e+00,
            -5.61139693e+00,  4.97772984e+00, -4.78539983e+00,
            -1.45305354e+00, -3.12854339e+00, -6.54954091e-01,
             2.09993123e+00, -4.71636592e+00,  1.19239096e+00,
             1.55885986e+00,  6.67866070e-01, -2.91741806e+00,
            -4.50219646e+00, -4.04282504e+00,  3.09119867e+00,
            -3.27877893e+00, -9.29314785e+00, -1.09178431e+00,
            -5.80679729e+00,  2.27732919e-01, -3.70559974e-01,
            -2.04604720e+00, -3.05746397e+00, -5.04011789e-01,
    ...
             6.16423864e+00, -2.61412750e-01,  1.65874754e+00,
            -7.12988744e+00,  4.95542022e+00, -4.89118576e+00,
             8.58312388e+00,  2.07240051e+00,  2.93747083e+00,
             7.78938524e+00,  1.70402618e+00,  6.18287018e-02,
             9.73669066e+00,  8.57462290e+00, -8.33799334e-01,
             5.66134633e+00, -6.31137040e+00,  1.62239492e-01,
             6.97770653e+00,  3.01296591e+00, -2.69087685e+00,
             7.20246084e+00, -2.14756395e+00,  3.58812866e+00,
             8.38896597e-01,  7.27977518e+00, -9.67591617e+00,
             1.91152462e+00,  3.39336925e+00, -3.62713034e+00,
             4.55111147e+00, -4.80339449e+00, -1.52242861e+00,
             6.60796118e+00,  8.75744733e+00, -6.92713657e+00,
            -2.89050931e+00, -1.10220253e+00,  1.85109463e+00,
            -7.65998246e+00, -5.52490783e+00,  5.57889433e+00,
            -5.62910237e-01, -2.57122648e-01, -1.24526050e+00,
            -9.95815030e-01,  9.99065065e+00,  1.31241234e+00,
             6.28519614e+00,  4.23234582e+00,  1.02096797e+01,
            -1.48178688e+00,  2.25202083e-03, -6.82222001e+00,
             2.62980361e+00, -4.19319392e+00, -2.32917260e+00,
            -2.24545308e+00, -2.43458445e+00]])

• Attributes: (4)

  created_at :

  2025-01-19T14:32:29.212688+00:00

  arviz_version :

  0.20.0

  inference_library :

  pymc

  inference_library_version :

  5.20.0