Getting started

Welcome to xarray-einstats!

xarray-einstats is an open source Python library part of the ArviZ project. It acts as a bridge between the xarray library for labelled arrays and libraries for raw arrays such as NumPy or SciPy.

Xarray has as “Compatibility with the broader ecosystem” as one of its main goals. Which is what allows xarray-einstats to perform this bridge role with minimal code and duplication.

Overview

xarray-einstats provides wrappers for:

These wrappers have the same names and functionality as the original functions. The difference in behaviour is that the wrappers will not make assumptions about the meaning of a dimension based on its position nor they have arguments like axis or axes. They will have dims argument that take dimension names instead of integers indicating the positions of the dimensions on which to act.

It also provides a handful of re-implemented functions:

These are partially reimplemented because the original function doesn’t yet support multidimensional and/or batched computations. They also share the name with a function in NumPy or SciPy, but they only implement a subset of the features. Moreover, the goal is for those to eventually be wrappers too.

Using xarray-einstats

DataArray inputs

Functions in xarray-einstats are designed to work on DataArray objects.

Let’s load some example data:

from xarray_einstats import linalg, stats, tutorial

da = tutorial.generate_matrices_dataarray(4)
da

<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB
3.799 0.4308 3.24 0.1412 0.9402 0.7951 ... 0.6156 1.124 0.8559 2.108 0.7637
Dimensions without coordinates: batch, experiment, dim, dim2

and show an example:

stats.skew(da, dims=["batch", "dim2"])

<xarray.DataArray (experiment: 3, dim: 4)> Size: 96B
1.256 1.432 0.9728 1.762 1.612 1.188 1.033 2.388 2.196 1.455 1.631 1.373
Dimensions without coordinates: experiment, dim

xarray-einstats uses dims as argument throughout the codebase as an alternative to both axis or axes indistinctively, also as alternative to the (..., M, M) convention used by NumPy.

The use of dims follows dot, instead of the singular dim argument used for example in mean. Both a single dimension or multiple are valid inputs, and using dims emphasizes the fact that operations and reductions can be performed over multiple dimensions at the same time. Moreover, in linear algebra functions, dims is often restricted to a 2 element list as it indicates which dimensions define the matrices, interpreting all the others as batch dimensions.

That means that the two calls below are equivalent, even if the dimension names of the inputs are not, because their dimension names are the same. Thus,

linalg.det(da, dims=["dim", "dim2"])

<xarray.DataArray (batch: 10, experiment: 3)> Size: 240B
23.55 2.033 0.3923 -7.374 0.06645 ... 1.804 -0.1599 8.875 -0.04935 -8.428
Dimensions without coordinates: batch, experiment

returns the same as:

linalg.det(da.transpose("dim2", "experiment", "dim", "batch"), dims=["dim", "dim2"])

<xarray.DataArray (experiment: 3, batch: 10)> Size: 240B
23.55 -7.374 -5.617 -12.29 1.77 -0.6289 ... -11.07 -0.5096 -28.77 -0.1599 -8.428
Dimensions without coordinates: experiment, batch

Important

In xarray_einstats only the dimension names matter, not their order.

Dataset and GroupBy inputs

While the DataArray is the base xarray object, there are also other xarray objects that are key while using the library. These other objects such as Dataset are implemented as a collection of DataArray objects, and all include a .map method in order to apply the same function to all its child DataArrays.

ds = tutorial.generate_mcmc_like_dataset(9438)
ds

<xarray.Dataset> Size: 6kB
Dimensions:  (plot_dim: 20, chain: 4, draw: 10, team: 6, match: 12)
Coordinates:
  * team     (team) <U1 24B 'a' 'b' 'c' 'd' 'e' 'f'
  * chain    (chain) int64 32B 0 1 2 3
  * draw     (draw) int64 80B 0 1 2 3 4 5 6 7 8 9
Dimensions without coordinates: plot_dim, match
Data variables:
    x_plot   (plot_dim) float64 160B 0.0 0.5263 1.053 1.579 ... 8.947 9.474 10.0
    mu       (chain, draw, team) float64 2kB 0.2691 0.1617 ... 0.4673 1.844
    sigma    (chain, draw) float64 320B 1.939 1.435 0.5109 ... 0.594 1.54 1.257
    score    (chain, draw, match) int64 4kB 0 2 3 0 0 0 0 0 ... 0 1 1 1 2 4 0 2

We can use map to apply the same function to all the 4 child DataArrays in ds, but this will not always be possible. When using .map, the function provided is applied to all child DataArrays with the same **kwargs.

If we try doing:

ds.map(stats.circmean, dims=("chain", "draw"))

Hide code cell output

 
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[6], line 1
----> 1 ds.map(stats.circmean, dims=("chain", "draw"))

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray/core/dataset.py:6936, in Dataset.map(self, func, keep_attrs, args, **kwargs)
   6933 if keep_attrs is None:
   6934     keep_attrs = _get_keep_attrs(default=False)
   6935 variables = {
-> 6936     k: maybe_wrap_array(v, func(v, *args, **kwargs))
   6937     for k, v in self.data_vars.items()
   6938 }
   6939 if keep_attrs:
   6940     for k, v in variables.items():

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray_einstats/stats.py:566, in circmean(da, dims, high, low, nan_policy, **kwargs)
    564 if nan_policy is not None:
    565     circmean_kwargs["nan_policy"] = nan_policy
--> 566 return _apply_reduce_func(stats.circmean, da, dims, kwargs, circmean_kwargs)

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray_einstats/stats.py:502, in _apply_reduce_func(func, da, dims, kwargs, func_kwargs)
    500 else:
    501     core_dims = [dims]
--> 502 out_da = xr.apply_ufunc(
    503     func, da, input_core_dims=[core_dims], output_core_dims=[[]], kwargs=func_kwargs, **kwargs
    504 )
    505 return out_da

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray/computation/apply_ufunc.py:1267, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, on_missing_core_dim, *args)
   1265 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
   1266 elif any(isinstance(a, DataArray) for a in args):
-> 1267     return apply_dataarray_vfunc(
   1268         variables_vfunc,
   1269         *args,
   1270         signature=signature,
   1271         join=join,
   1272         exclude_dims=exclude_dims,
   1273         keep_attrs=keep_attrs,
   1274     )
   1275 # feed Variables directly through apply_variable_ufunc
   1276 elif any(isinstance(a, Variable) for a in args):

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray/computation/apply_ufunc.py:310, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    305 result_coords, result_indexes = build_output_coords_and_indexes(
    306     args, signature, exclude_dims, combine_attrs=keep_attrs
    307 )
    309 data_vars = [getattr(a, "variable", a) for a in args]
--> 310 result_var = func(*data_vars)
    312 out: tuple[DataArray, ...] | DataArray
    313 if signature.num_outputs > 1:

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray/computation/apply_ufunc.py:730, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    723 broadcast_dims = tuple(
    724     dim for dim in dim_sizes if dim not in signature.all_core_dims
    725 )
    726 output_dims = [broadcast_dims + out for out in signature.output_core_dims]
    728 input_data = [
    729     (
--> 730         broadcast_compat_data(arg, broadcast_dims, core_dims)
    731         if isinstance(arg, Variable)
    732         else arg
    733     )
    734     for arg, core_dims in zip(args, signature.input_core_dims, strict=True)
    735 ]
    737 if any(is_chunked_array(array) for array in input_data):
    738     if dask == "forbidden":

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray/computation/apply_ufunc.py:675, in broadcast_compat_data(variable, broadcast_dims, core_dims)
    673 reordered_dims = old_broadcast_dims + core_dims
    674 if reordered_dims != old_dims:
--> 675     order = tuple(old_dims.index(d) for d in reordered_dims)
    676     data = duck_array_ops.transpose(data, order)
    678 if new_dims != reordered_dims:

File ~/checkouts/readthedocs.org/user_builds/xarray-einstats/envs/78/lib/python3.12/site-packages/xarray/computation/apply_ufunc.py:675, in <genexpr>(.0)
    673 reordered_dims = old_broadcast_dims + core_dims
    674 if reordered_dims != old_dims:
--> 675     order = tuple(old_dims.index(d) for d in reordered_dims)
    676     data = duck_array_ops.transpose(data, order)
    678 if new_dims != reordered_dims:

ValueError: tuple.index(x): x not in tuple

we get an exception. The chain and draw dimensions are not present in all child DataArrays. Instead, we could apply it only to the variables that have both chain and dim dimensions.

ds_samples = ds[["mu", "sigma", "score"]]
ds_samples.map(stats.circmean, dims=("chain", "draw"))

<xarray.Dataset> Size: 176B
Dimensions:  (team: 6, match: 12)
Coordinates:
  * team     (team) <U1 24B 'a' 'b' 'c' 'd' 'e' 'f'
Dimensions without coordinates: match
Data variables:
    mu       (team) float64 48B 0.8221 0.7376 0.6195 0.7485 0.7439 0.7818
    sigma    float64 8B 0.8134
    score    (match) float64 96B 0.7441 0.3923 0.9316 ... 0.5814 0.9538 0.94

Attention

In general, you should prefer using .map attribute over using non-DataArray objects as input to the xarray_einstats directly. .map will ensure no unexpected broadcasting between the multiple child DataArrays takes place. See the examples below for some examples.

However, if you are using functions that reduce dimensions on non-DataArray inputs whose child DataArrays all have all the dimensions to reduce you will not trigger any such broadcasting, and we have included that behaviour on our test suite to ensure it stays this way.

It is also possible to do

stats.circmean(ds_samples, dims=("chain", "draw"))

<xarray.Dataset> Size: 176B
Dimensions:  (team: 6, match: 12)
Coordinates:
  * team     (team) <U1 24B 'a' 'b' 'c' 'd' 'e' 'f'
Dimensions without coordinates: match
Data variables:
    mu       (team) float64 48B 0.8221 0.7376 0.6195 0.7485 0.7439 0.7818
    sigma    float64 8B 0.8134
    score    (match) float64 96B 0.7441 0.3923 0.9316 ... 0.5814 0.9538 0.94

Here, all child DataArrays have both chain and draw dimension, so as expected, the result is the same. There are some cases however, in which not using .map triggers some broadcasting operations which will generally not be the desired output.

If we use the .map attribute, the function is applied to each child DataArray independently from the others:

ds.map(stats.rankdata)

<xarray.Dataset> Size: 6kB
Dimensions:  (plot_dim: 20, chain: 4, draw: 10, team: 6, match: 12)
Dimensions without coordinates: plot_dim, chain, draw, team, match
Data variables:
    x_plot   (plot_dim) float64 160B 1.0 2.0 3.0 4.0 5.0 ... 17.0 18.0 19.0 20.0
    mu       (chain, draw, team) float64 2kB 65.0 41.0 89.0 ... 55.0 97.0 205.0
    sigma    (chain, draw) float64 320B 33.0 30.0 15.0 5.0 ... 18.0 31.0 29.0
    score    (chain, draw, match) float64 4kB 105.0 401.0 457.0 ... 105.0 401.0

whereas without using the .map attribute, extra broadcasting can happen:

stats.rankdata(ds)

<xarray.Dataset> Size: 2MB
Dimensions:  (plot_dim: 20, chain: 4, draw: 10, team: 6, match: 12)
Dimensions without coordinates: plot_dim, chain, draw, team, match
Data variables:
    x_plot   (plot_dim, chain, draw, team, match) float64 461kB 1.44e+03 ... ...
    mu       (plot_dim, chain, draw, team, match) float64 461kB 1.548e+04 ......
    sigma    (plot_dim, chain, draw, team, match) float64 461kB 4.68e+04 ... ...
    score    (plot_dim, chain, draw, team, match) float64 461kB 1.254e+04 ......

The behaviour on DataArrayGroupBy for example is very similar to the examples we have shown for Datasets:

da = ds["mu"].assign_coords(team=["a", "b", "b", "a", "c", "b"])
da

<xarray.DataArray 'mu' (chain: 4, draw: 10, team: 6)> Size: 2kB
0.2691 0.1617 0.4371 0.4885 0.1836 2.149 ... 2.037 0.09032 0.2221 0.4673 1.844
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3
  * draw     (draw) int64 80B 0 1 2 3 4 5 6 7 8 9
  * team     (team) <U1 24B 'a' 'b' 'b' 'a' 'c' 'b'

when we apply a “group by” operation over the team dimension, we generate a DataArrayGroupBy with 3 groups.

gb = da.groupby("team")
gb

<DataArrayGroupBy, grouped over 1 grouper(s), 3 groups in total:
    'team': UniqueGrouper('team'), 3/3 groups with labels 'a', 'b', 'c'>

on which we can use .map to apply a function from xarray-einstats over all groups independently:

gb.map(stats.median_abs_deviation, dims=["draw", "team"])

<xarray.DataArray 'mu' (chain: 4, team: 3)> Size: 96B
0.3436 0.3758 0.2351 0.5221 0.5937 0.4158 ... 0.4314 0.212 0.3479 0.5708 0.2288
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3
  * team     (team) object 24B 'a' 'b' 'c'

which as expected has performed the operation group-wise, yielding a different result than either

stats.median_abs_deviation(da, dims=["draw", "team"])

<xarray.DataArray 'mu' (chain: 4)> Size: 32B
0.3444 0.5968 0.4553 0.4069
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3

or

stats.median_abs_deviation(da, dims="draw")

<xarray.DataArray 'mu' (chain: 4, team: 6)> Size: 192B
0.3452 0.3788 0.09536 0.3892 0.2351 ... 0.6554 0.2451 0.3832 0.2288 0.5281
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3
  * team     (team) <U1 24B 'a' 'b' 'b' 'a' 'c' 'b'

See also

Check out the GroupBy: Group and Bin Data page on xarray’s documentation.