"""Stats, linear algebra and einops for xarray."""
from __future__ import annotations
from collections.abc import Callable, Hashable, Iterable, Sequence
from typing import Any
import numpy as np
import xarray as xr
from numpy.typing import DTypeLike
from .linalg import einsum, einsum_path, matmul
from .accessors import LinAlgAccessor, EinopsAccessor
from ._types import Dims, DimOrDims
__all__ = [
"einsum",
"einsum_path",
"matmul",
"zeros_ref",
"ones_ref",
"empty_ref",
"LinAlgAccessor",
"EinopsAccessor",
"Dims",
"DimOrDims",
]
__version__ = "0.11.0"
[docs]
def sort(
da: xr.DataArray,
dim: Hashable,
kind: str | None = None,
stable: bool | None = None,
**kwargs: Any,
) -> xr.DataArray:
"""Sort along dimension using DataArray values.
Wrapper around :func:`numpy.sort`
The returned DataArray has the same shape and dimensions as the original one
but the coordinate values along `dim` no longer make sense given each subset
along the other dimensions can have a different order along `dim` so they are removed.
Parameters
----------
da : DataArray
Input data
dim : hashable
Dimension along which to sort using dataarray values, not coordinates.
kind : str, optional
stable : bool, optional
**kwargs
Passed to :func:`xarray.apply_ufunc`
Returns
-------
DataArray
"""
sort_kwargs = {"axis": -1, "kind": kind, "stable": stable}
return xr.apply_ufunc(
np.sort,
da,
input_core_dims=[[dim]],
output_core_dims=[[dim]],
kwargs=sort_kwargs,
**kwargs,
).drop_vars(dim, errors="ignore")
def _remove_indexes_to_reduce(da: xr.DataArray, dims: Iterable[Hashable]) -> xr.DataArray:
"""Remove indexes related to provided dims.
Removes indexes related to dims on which we need to operate.
As many functions only support integer `axis` or None,
in order to have our functions operate on multiple dimensions
we need to stack/flatten them. If some of those dimensions
are already indexed by a multiindex this doesn't work, so we
remove the indexes. As they are reduced, that information
will end up being lost eventually either way.
"""
index_keys = list(da.indexes)
remove_indicator = [
(any(da.indexes[k] is index for k in index_keys if k in dims))
for name, index in da.indexes.items()
]
indexes_to_remove = [k for k, remove in zip(index_keys, remove_indicator) if remove]
da = da.drop_indexes(indexes_to_remove)
coords_to_remove = [coord for coord in da.coords if coord in indexes_to_remove or coord in dims]
return da.reset_coords(coords_to_remove, drop=True)
def _find_index(
elem: Hashable, to_search_in: Sequence[xr.DataArray | xr.Dataset]
) -> tuple[int, Hashable]:
for i, da in enumerate(to_search_in):
if elem in da.dims:
return (i, elem)
raise ValueError(f"{elem} not found in any reference object")
def _create_ref(
*args: xr.DataArray | xr.Dataset,
dims: Sequence[Hashable],
np_creator: Callable[..., Any],
dtype: DTypeLike | None = None,
) -> xr.DataArray:
if dtype is None and all(isinstance(arg, xr.DataArray) for arg in args):
dtype = np.result_type(*[arg.dtype for arg in args])
ref_idxs = [_find_index(dim, args) for dim in dims]
shape = [len(args[idx][dim]) for idx, dim in ref_idxs]
# TODO: keep non indexing coords?
coords = {dim: args[idx][dim] for idx, dim in ref_idxs if dim in args[idx].coords}
return xr.DataArray(
np_creator(shape, dtype=dtype),
dims=dims,
coords=coords,
)
[docs]
def zeros_ref(
*args: xr.DataArray | xr.Dataset, dims: Sequence[Hashable], dtype: DTypeLike | None = None
) -> xr.DataArray:
"""Create a zeros DataArray from reference object(s).
Creates a DataArray filled with zeros from reference
DataArrays or Datasets and a list with the desired dimensions.
Parameters
----------
*args : DataArray or Dataset
Reference objects from which the lengths and coordinate values (if any)
of the given `dims` will be taken.
dims : sequence of hashable
List of dimensions of the output DataArray. Passed as is to the
:class:`~xarray.DataArray` constructor.
dtype : dtype, optional
The dtype of the output array.
If it is not provided it will be inferred from the reference
DataArrays in `args` with :func:`numpy.result_type`.
Returns
-------
DataArray
See Also
--------
ones_ref, empty_ref
"""
return _create_ref(*args, dims=dims, np_creator=np.zeros, dtype=dtype)
[docs]
def empty_ref(
*args: xr.DataArray | xr.Dataset, dims: Sequence[Hashable], dtype: DTypeLike | None = None
) -> xr.DataArray:
"""Create an empty DataArray from reference object(s).
Creates an empty DataArray from reference
DataArrays or Datasets and a list with the desired dimensions.
Parameters
----------
*args : DataArray or Dataset
Reference objects from which the lengths and coordinate values (if any)
of the given `dims` will be taken.
dims : sequence of hashable
List of dimensions of the output DataArray. Passed as is to the
:class:`~xarray.DataArray` constructor.
dtype : dtype, optional
The dtype of the output array.
If it is not provided it will be inferred from the reference
DataArrays in `args` with :func:`numpy.result_type`.
Returns
-------
DataArray
See Also
--------
ones_ref, zeros_ref
"""
return _create_ref(*args, dims=dims, np_creator=np.empty, dtype=dtype)
[docs]
def ones_ref(
*args: xr.DataArray | xr.Dataset, dims: Sequence[Hashable], dtype: DTypeLike | None = None
) -> xr.DataArray:
"""Create a ones DataArray from reference object(s).
Creates a DataArray filled with ones from reference
DataArrays or Datasets and a list with the desired dimensions.
Parameters
----------
*args : DataArray or Dataset
Reference objects from which the lengths and coordinate values (if any)
of the given `dims` will be taken.
dims : sequence of hashable
List of dimensions of the output DataArray. Passed as is to the
:class:`~xarray.DataArray` constructor.
dtype : dtype, optional
The dtype of the output array.
If it is not provided it will be inferred from the reference
DataArrays in `args` with :func:`numpy.result_type`.
Returns
-------
DataArray
See Also
--------
empty_ref, zeros_ref
"""
return _create_ref(*args, dims=dims, np_creator=np.ones, dtype=dtype)