arviz.extract#
- arviz.extract(data, group='posterior', combined=True, var_names=None, filter_vars=None, num_samples=None, keep_dataset=False, rng=None)[source]#
Extract an InferenceData group or subset of it.
- Parameters
- idata
InferenceData
orInferenceData_like
InferenceData from which to extract the data.
- group
str
, optional Which InferenceData data group to extract data from.
- combinedbool, optional
Combine
chain
anddraw
dimensions intosample
. Won’t work if a dimension namedsample
already exists.- var_names
str
orlist
ofstr
, 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 lapandas.filter
. Like with plotting, sometimes it’s easier to subset saying what to exclude instead of what to include- num_samples
int
, 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 ifnum_samples
is notNone
, 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.
- idata
- Returns
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
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([[ 5.00053080e-01, 8.85972309e+00, 1.49014212e+01, 4.62526861e+00, 3.44326320e+01, 4.37393819e+00, 1.11110725e+01, 7.19731695e+00, 1.92743101e+01, 1.99628075e+01, 1.01448164e+01, 5.36521047e+00, 3.56959557e+00, 7.96180932e+00, 3.24428943e+00, 9.80188050e+00, 6.32589200e+00, 7.15289257e-01, 1.39026471e+00, 1.09299330e+01, -1.74303095e+00, 2.81729489e+01, 5.36521047e+00, 1.10205196e+01, 6.44774439e+00, 6.13458430e+00, 1.54518943e+01, 1.04202606e+01, -1.94662756e+00, -1.45632453e+00, 6.19166018e+00, 9.97226578e+00, 5.14813384e+00, 1.64329887e+01, 4.06217766e+00, 6.98422466e+00, 1.97876889e+00, 1.39511471e+00, 1.12413177e+01, -3.71502476e+00, 1.68769505e+01, 2.53235961e+01, 1.99067780e+00, 3.96109671e+00, 8.06311068e+00, 9.75288135e+00, 3.34359524e+00, 8.65007191e+00, 7.43254100e-01, -3.17824163e+00, -8.10820282e-02, -1.45575080e+00, 5.06773825e+00, 1.24254132e+01, 2.01472848e+01, 4.64574154e+00, 7.23326740e+00, 9.47489537e+00, 4.83051016e+00, -3.44660755e+00, ... 3.65105149e+00, -2.38039401e+00, 2.84813834e-01, 6.05209366e+00, 3.45131075e+00, 5.23739698e+00, 4.28997121e+00, -8.11232046e+00, 3.12690796e-01, 6.61190575e-02, 7.08384221e+00, 5.23985192e+00, -3.19486996e-01, 2.44438035e+00, 1.18865961e+00, 5.83475486e-01, 4.43677457e+00, -1.43622823e+00, 1.75683067e+01, 3.35911194e+00, 1.13628835e+01, 7.90169447e+00, 9.24872843e+00, 5.26503444e+00, 9.65910877e+00, 8.75646362e+00, 1.67007553e+00, 5.31689027e+00, 5.29491415e+00, -2.47571242e+00, 3.65565286e+00, 1.37839425e+01, 3.35132225e+00, 4.38037389e+00, 2.86110208e+00, 8.86523760e+00, -5.93250093e-01, 5.55416881e+00, 6.76390558e+00, -2.15337193e+00, 4.11577504e-02, 1.49927153e+00, 7.33997694e+00, 3.35132225e+00, 2.39378808e+00, 8.86580029e+00, 9.39143626e+00, 3.31390557e+00, 1.80393980e+00, 5.75992457e+00, 3.35132225e+00, 3.58898903e+00, 9.12635811e+00, 3.35132225e+00, 6.23728051e+00, -1.78243701e+00, 7.25149063e+00, -7.34760704e+00]]) Coordinates: * school (school) object 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon' * sample (sample) object MultiIndex * chain (sample) int64 0 2 3 3 2 1 2 2 3 3 1 3 ... 3 2 1 1 1 3 1 1 1 0 2 0 * draw (sample) int64 449 418 67 148 251 257 ... 278 460 114 376 206 135
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