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)[source]#
Extract a group or group subset from a DataTree.
- Parameters:
- idataDataTree-like
DataTree from which to extract the data.
- group
str
, 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 torcParams["data.sample_dims"]
.- combinedbool, optional
Combine sample_dims dimensions into
sample
. Won’t work if a dimension namedsample
already exists. It is irrelevant and ignored when sample_dims is a single dimension.- 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 la pandas.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
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_method
str
, 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_seed
int
,numpy.Generator
, optional Random number generator or seed. Only used if
weights
is notNone
or ifnum_samples
is notNone
.
- Returns:
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: * school (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon' * 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 Data variables: mu (sample) float64 16kB 7.872 3.385 9.1 7.304 ... 1.767 3.486 3.404 theta (school, sample) float64 128kB 12.32 11.29 5.709 ... 8.452 1.295 tau (sample) float64 16kB 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)> Size: 6kB array([[ 3.09387429e+01, 5.71620186e+00, 8.14716772e+00, 1.22880450e+01, -8.12128199e+00, 6.32454995e+00, 1.25833151e+01, 7.35127382e+00, 6.61626771e+00, 6.97027172e+00, 3.34359524e+00, 3.21324974e-01, 2.86064613e+00, 8.15327277e+00, 2.28772366e+00, 2.10585946e+00, 7.80435210e+00, 9.54687800e+00, 2.07538479e+00, 7.68522224e+00, 3.11147235e+00, 7.03102581e+00, 7.88942950e+00, 8.64601026e+00, 1.59722640e+01, 4.07866337e+00, 3.98511536e-01, 8.06311068e+00, 1.26162675e+01, 9.44361810e+00, 7.57793225e-01, 1.64329887e+01, 1.20049894e+01, 2.97785044e+00, -1.13886614e+00, 1.04025013e+01, 4.53814410e+00, 8.03517762e+00, 9.15852024e+00, 5.83123935e+00, 1.00726585e+01, 8.65007191e+00, -3.17824163e+00, 3.54038165e+00, 6.10959984e+00, 7.21919702e+00, -6.50164631e-01, 7.05688130e+00, -2.02773042e+00, 1.83440742e+01, 3.06489534e+00, 3.33922737e+00, 8.06378285e-01, 1.69945398e+01, 1.68496500e+01, 6.03475550e+00, 1.13370369e+01, 8.93305013e+00, 7.03588032e+00, 1.22060355e+01, ... -8.11232046e+00, 5.79259069e+00, 7.63541882e+00, 2.56752133e+00, -1.02702293e+00, 5.88696762e+00, 3.78957862e-01, -2.76035595e+00, 4.53714710e+00, 7.18026928e+00, -3.58286664e+00, 6.65588464e+00, 7.86258400e+00, 5.30239912e+00, 7.84375303e+00, 8.79079955e+00, 7.42415492e+00, 1.82571418e+00, 3.77671414e+00, 8.77336892e+00, 5.59107632e+00, 1.50613658e+01, 7.47430176e+00, 1.78499767e+00, 3.41427310e+00, 2.94760112e+00, -1.16496750e-01, 1.75683067e+01, 2.85143502e+00, 5.16522809e+00, 1.15479225e+01, 6.57878387e+00, 3.14127906e+00, 2.34452622e+00, 2.69293481e+00, -2.07833386e+00, 3.93829921e+00, 7.20084454e+00, -6.25581924e+00, 8.52437956e+00, 3.14824203e+00, 9.55063157e-01, -4.86167897e+00, 5.08930303e+00, 7.92068098e+00, 4.18385037e+00, 3.35132225e+00, 6.80602415e+00, 9.24872843e+00, 1.40735354e+01, -2.58684632e+00, 6.99697482e+00, 4.95648908e+00, 5.24910668e+00, 6.99697482e+00, 1.06309246e+01, 6.30089763e+00, 3.59552609e+00]]) Coordinates: * school (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon' * sample (sample) object 800B MultiIndex * chain (sample) int64 800B 3 0 2 3 1 0 0 3 1 1 2 ... 2 3 3 2 3 1 1 3 3 2 1 * draw (sample) int64 800B 157 462 170 3 19 151 28 ... 487 72 123 40 5 260
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: tau (chain, draw) float64 4kB 1.941 3.388 4.208 ... 0.06893 2.145 theta (chain, draw, school) float64 32kB 4.866 4.59 ... -2.031 6.045 mu (chain, draw) float64 4kB 3.903 3.915 -1.751 ... 0.7908 2.869 Attributes: arviz_version: 0.13.0.dev0 created_at: 2022-10-13T14:37:26.602116 inference_library: pymc inference_library_version: 4.2.2