xarray_einstats_einops_rearrange_reduce_and_stats_quickstart.py

1import xarray as xr
2import numpy as np
3from xarray_einstats import einops, stats
4
5# Create a sample xarray DataArray
6data = xr.DataArray(
7    np.random.normal(size=(4, 10, 3)),
8    coords={"chain": range(4), "draw": range(10), "team": ["A", "B", "C"]},
9    dims=("chain", "draw", "team"),
10)
11
12# Use einops.rearrange to reshape and transpose using dimension names
13# Here we combine chain and draw into a single 'sample' dimension
14rearranged = einops.rearrange(data, "(chain draw) team -> team sample")
15
16# Use einops.reduce to perform reductions (like mean) over specific dimensions
17reduced = einops.reduce(data, "chain draw team -> team", reduction="mean")
18
19# Use stats module to compute wrapped scipy-like statistics
20# For example, computing the rank along the 'draw' dimension
21ranks = stats.rankdata(data, dim="draw")
22
23print("Rearranged shape:", rearranged.dims)
24print("Reduced values:", reduced.values)
25print("Ranks head:", ranks.sel(chain=0, team="A").values[:5])