diff --git a/.github/workflows/array-api-tests.yml b/.github/workflows/array-api-tests.yml new file mode 100644 index 00000000000..efeeaecbf66 --- /dev/null +++ b/.github/workflows/array-api-tests.yml @@ -0,0 +1,46 @@ +name: NamedArray._array_api + +on: [push, pull_request] + +env: + PYTEST_ARGS: "-v -rxXfE --ci --hypothesis-disable-deadline --max-examples 5" + API_VERSIONS: "2023.12" + +jobs: + array-api-tests: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.12'] + exclude: + - python-version: '3.8' + + steps: + - name: Checkout xarray + uses: actions/checkout@v4 + with: + path: xarray + - name: Checkout array-api-tests + uses: actions/checkout@v4 + with: + repository: data-apis/array-api-tests + submodules: 'true' + path: array-api-tests + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install Dependencies + run: | + python -m pip install --upgrade pip + python -m pip install ${GITHUB_WORKSPACE}/xarray + python -m pip install -r ${GITHUB_WORKSPACE}/array-api-tests/requirements.txt + python -m pip install hypothesis + python -m pip install array-api-strict + python -m pip install array-api-compat + - name: Run the array API testsuite + env: + ARRAY_API_TESTS_MODULE: xarray.namedarray._array_api + run: | + cd ${GITHUB_WORKSPACE}/array-api-tests + pytest array_api_tests/ --skips-file ${GITHUB_WORKSPACE}/xarray/xarray/tests/namedarray_array_api_skips.txt ${PYTEST_ARGS} diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 07113d66b5b..ee866fdd36c 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -370,7 +370,7 @@ def _as_array_or_item(data): return data -class Variable(NamedArray, AbstractArray, VariableArithmetic): +class Variable(AbstractArray, VariableArithmetic, NamedArray): """A netcdf-like variable consisting of dimensions, data and attributes which describe a single Array. A single Variable object is not fully described outside the context of its parent Dataset (if you want such a diff --git a/xarray/namedarray/_array_api.py b/xarray/namedarray/_array_api.py deleted file mode 100644 index 9cd064b110e..00000000000 --- a/xarray/namedarray/_array_api.py +++ /dev/null @@ -1,221 +0,0 @@ -from __future__ import annotations - -from types import ModuleType -from typing import Any - -import numpy as np - -from xarray.namedarray._typing import ( - Default, - _arrayapi, - _Axes, - _Axis, - _default, - _Dim, - _DType, - _ScalarType, - _ShapeType, - _SupportsImag, - _SupportsReal, -) -from xarray.namedarray.core import NamedArray - - -def _get_data_namespace(x: NamedArray[Any, Any]) -> ModuleType: - if isinstance(x._data, _arrayapi): - return x._data.__array_namespace__() - - return np - - -# %% Creation Functions - - -def astype( - x: NamedArray[_ShapeType, Any], dtype: _DType, /, *, copy: bool = True -) -> NamedArray[_ShapeType, _DType]: - """ - Copies an array to a specified data type irrespective of Type Promotion Rules rules. - - Parameters - ---------- - x : NamedArray - Array to cast. - dtype : _DType - Desired data type. - copy : bool, optional - Specifies whether to copy an array when the specified dtype matches the data - type of the input array x. - If True, a newly allocated array must always be returned. - If False and the specified dtype matches the data type of the input array, - the input array must be returned; otherwise, a newly allocated array must be - returned. Default: True. - - Returns - ------- - out : NamedArray - An array having the specified data type. The returned array must have the - same shape as x. - - Examples - -------- - >>> narr = NamedArray(("x",), np.asarray([1.5, 2.5])) - >>> narr - Size: 16B - array([1.5, 2.5]) - >>> astype(narr, np.dtype(np.int32)) - Size: 8B - array([1, 2], dtype=int32) - """ - if isinstance(x._data, _arrayapi): - xp = x._data.__array_namespace__() - return x._new(data=xp.astype(x._data, dtype, copy=copy)) - - # np.astype doesn't exist yet: - return x._new(data=x._data.astype(dtype, copy=copy)) # type: ignore[attr-defined] - - -# %% Elementwise Functions - - -def imag( - x: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var] - /, -) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: - """ - Returns the imaginary component of a complex number for each element x_i of the - input array x. - - Parameters - ---------- - x : NamedArray - Input array. Should have a complex floating-point data type. - - Returns - ------- - out : NamedArray - An array containing the element-wise results. The returned array must have a - floating-point data type with the same floating-point precision as x - (e.g., if x is complex64, the returned array must have the floating-point - data type float32). - - Examples - -------- - >>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) - >>> imag(narr) - Size: 16B - array([2., 4.]) - """ - xp = _get_data_namespace(x) - out = x._new(data=xp.imag(x._data)) - return out - - -def real( - x: NamedArray[_ShapeType, np.dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var] - /, -) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: - """ - Returns the real component of a complex number for each element x_i of the - input array x. - - Parameters - ---------- - x : NamedArray - Input array. Should have a complex floating-point data type. - - Returns - ------- - out : NamedArray - An array containing the element-wise results. The returned array must have a - floating-point data type with the same floating-point precision as x - (e.g., if x is complex64, the returned array must have the floating-point - data type float32). - - Examples - -------- - >>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) - >>> real(narr) - Size: 16B - array([1., 2.]) - """ - xp = _get_data_namespace(x) - out = x._new(data=xp.real(x._data)) - return out - - -# %% Manipulation functions -def expand_dims( - x: NamedArray[Any, _DType], - /, - *, - dim: _Dim | Default = _default, - axis: _Axis = 0, -) -> NamedArray[Any, _DType]: - """ - Expands the shape of an array by inserting a new dimension of size one at the - position specified by dims. - - Parameters - ---------- - x : - Array to expand. - dim : - Dimension name. New dimension will be stored in the axis position. - axis : - (Not recommended) Axis position (zero-based). Default is 0. - - Returns - ------- - out : - An expanded output array having the same data type as x. - - Examples - -------- - >>> x = NamedArray(("x", "y"), np.asarray([[1.0, 2.0], [3.0, 4.0]])) - >>> expand_dims(x) - Size: 32B - array([[[1., 2.], - [3., 4.]]]) - >>> expand_dims(x, dim="z") - Size: 32B - array([[[1., 2.], - [3., 4.]]]) - """ - xp = _get_data_namespace(x) - dims = x.dims - if dim is _default: - dim = f"dim_{len(dims)}" - d = list(dims) - d.insert(axis, dim) - out = x._new(dims=tuple(d), data=xp.expand_dims(x._data, axis=axis)) - return out - - -def permute_dims(x: NamedArray[Any, _DType], axes: _Axes) -> NamedArray[Any, _DType]: - """ - Permutes the dimensions of an array. - - Parameters - ---------- - x : - Array to permute. - axes : - Permutation of the dimensions of x. - - Returns - ------- - out : - An array with permuted dimensions. The returned array must have the same - data type as x. - - """ - - dims = x.dims - new_dims = tuple(dims[i] for i in axes) - if isinstance(x._data, _arrayapi): - xp = _get_data_namespace(x) - out = x._new(dims=new_dims, data=xp.permute_dims(x._data, axes)) - else: - out = x._new(dims=new_dims, data=x._data.transpose(axes)) # type: ignore[attr-defined] - return out diff --git a/xarray/namedarray/_array_api/__init__.py b/xarray/namedarray/_array_api/__init__.py new file mode 100644 index 00000000000..957032b09b4 --- /dev/null +++ b/xarray/namedarray/_array_api/__init__.py @@ -0,0 +1,371 @@ +__all__ = [] + +__array_api_version__ = "2023.12" + +__all__ += ["__array_api_version__"] + +# from xarray.namedarray.core import NamedArray as Array + +# __all__ += ["Array"] + +from xarray.namedarray._array_api._constants import e, inf, nan, newaxis, pi + +__all__ += ["e", "inf", "nan", "newaxis", "pi"] + +from xarray.namedarray._array_api._creation_functions import ( + arange, + asarray, + empty, + empty_like, + eye, + from_dlpack, + full, + full_like, + linspace, + meshgrid, + ones, + ones_like, + tril, + triu, + zeros, + zeros_like, +) + +__all__ += [ + "arange", + "asarray", + "empty", + "empty_like", + "eye", + "from_dlpack", + "full", + "full_like", + "linspace", + "meshgrid", + "ones", + "ones_like", + "tril", + "triu", + "zeros", + "zeros_like", +] + +from xarray.namedarray._array_api._data_type_functions import ( + astype, + can_cast, + finfo, + iinfo, + isdtype, + result_type, +) + +__all__ += [ + "astype", + "can_cast", + "finfo", + "iinfo", + "isdtype", + "result_type", +] + +from xarray.namedarray._array_api._dtypes import ( + bool, + complex64, + complex128, + float32, + float64, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, +) + +__all__ += [ + "bool", + "complex64", + "complex128", + "float32", + "float64", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", +] + +from xarray.namedarray._array_api._elementwise_functions import ( + abs, + acos, + acosh, + add, + asin, + asinh, + atan, + atan2, + atanh, + bitwise_and, + bitwise_invert, + bitwise_left_shift, + bitwise_or, + bitwise_right_shift, + bitwise_xor, + ceil, + clip, + conj, + copysign, + cos, + cosh, + divide, + equal, + exp, + expm1, + floor, + floor_divide, + greater, + greater_equal, + hypot, + imag, + isfinite, + isinf, + isnan, + less, + less_equal, + log, + log1p, + log2, + log10, + logaddexp, + logical_and, + logical_not, + logical_or, + logical_xor, + maximum, + minimum, + multiply, + negative, + not_equal, + positive, + pow, + real, + remainder, + round, + sign, + signbit, + sin, + sinh, + sqrt, + square, + subtract, + tan, + tanh, + trunc, +) + +__all__ += [ + "abs", + "acos", + "acosh", + "add", + "asin", + "asinh", + "atan", + "atan2", + "atanh", + "bitwise_and", + "bitwise_invert", + "bitwise_left_shift", + "bitwise_or", + "bitwise_right_shift", + "bitwise_xor", + "ceil", + "clip", + "conj", + "copysign", + "cos", + "cosh", + "divide", + "equal", + "exp", + "expm1", + "floor", + "floor_divide", + "greater", + "greater_equal", + "hypot", + "imag", + "isfinite", + "isinf", + "isnan", + "less", + "less_equal", + "log", + "log1p", + "log2", + "log10", + "logaddexp", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "maximum", + "minimum", + "multiply", + "negative", + "not_equal", + "positive", + "pow", + "real", + "remainder", + "round", + "sign", + "signbit", + "sin", + "sinh", + "sqrt", + "square", + "subtract", + "tan", + "tanh", + "trunc", +] + +import xarray.namedarray._array_api._fft as fft + +__all__ = ["fft"] + +from xarray.namedarray._array_api._indexing_functions import take + +__all__ += ["take"] + +from xarray.namedarray._array_api._info import __array_namespace_info__ + +__all__ += [ + "__array_namespace_info__", +] + +import xarray.namedarray._array_api._linalg as linalg + +__all__ = ["linalg"] + + +from xarray.namedarray._array_api._linear_algebra_functions import ( + matmul, + matrix_transpose, + tensordot, + vecdot, +) + +__all__ += [ + "matmul", + "matrix_transpose", + "tensordot", + "vecdot", +] + +from xarray.namedarray._array_api._manipulation_functions import ( + broadcast_arrays, + broadcast_to, + concat, + expand_dims, + flip, + moveaxis, + permute_dims, + repeat, + reshape, + roll, + squeeze, + stack, + tile, + unstack, +) + +__all__ += [ + "broadcast_arrays", + "broadcast_to", + "concat", + "expand_dims", + "flip", + "moveaxis", + "permute_dims", + "repeat", + "reshape", + "roll", + "squeeze", + "stack", + "tile", + "unstack", +] + +from xarray.namedarray._array_api._searching_functions import ( + argmax, + argmin, + nonzero, + searchsorted, + where, +) + +__all__ += [ + "argmax", + "argmin", + "nonzero", + "searchsorted", + "where", +] + +from xarray.namedarray._array_api._set_functions import ( + unique_all, + unique_counts, + unique_inverse, + unique_values, +) + +__all__ += [ + "unique_all", + "unique_counts", + "unique_inverse", + "unique_values", +] + + +from xarray.namedarray._array_api._sorting_functions import argsort, sort + +__all__ += ["argsort", "sort"] + +from xarray.namedarray._array_api._statistical_functions import ( + cumulative_sum, + max, + mean, + min, + prod, + std, + sum, + var, +) + +__all__ += [ + "cumulative_sum", + "max", + "mean", + "min", + "prod", + "std", + "sum", + "var", +] + +from xarray.namedarray._array_api._utility_functions import ( + all, + any, +) + +__all__ += [ + "all", + "any", +] diff --git a/xarray/namedarray/_array_api/_constants.py b/xarray/namedarray/_array_api/_constants.py new file mode 100644 index 00000000000..8b8da5ffbe8 --- /dev/null +++ b/xarray/namedarray/_array_api/_constants.py @@ -0,0 +1,9 @@ +from xarray.namedarray._array_api._utils import _maybe_default_namespace + +_xp = _maybe_default_namespace() + +e = _xp.e +inf = _xp.inf +nan = _xp.nan +newaxis = _xp.newaxis +pi = _xp.pi diff --git a/xarray/namedarray/_array_api/_creation_functions.py b/xarray/namedarray/_array_api/_creation_functions.py new file mode 100644 index 00000000000..dbacbc407a3 --- /dev/null +++ b/xarray/namedarray/_array_api/_creation_functions.py @@ -0,0 +1,656 @@ +from __future__ import annotations + +from typing import Any, overload + +from xarray.namedarray._array_api._utils import ( + _broadcast_dims, + _get_data_namespace, + _get_namespace, + _get_namespace_dtype, + _infer_dims, + _move_dims, + _normalize_dimensions, +) +from xarray.namedarray._typing import ( + Default, + _ArrayLike, + _default, + _Device, + _DimsLike2, + _DType, + _Shape, + _Shape1D, + _ShapeType, + duckarray, +) +from xarray.namedarray.core import NamedArray + + +def arange( + start: int | float, + /, + stop: int | float | None = None, + step: int | float = 1, + *, + dtype: _DType | None = None, + device: _Device | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_Shape1D, _DType]: + """ + Returns evenly spaced values within the half-open interval [start, stop) as a one-dimensional array. + + Examples + -------- + >>> arange(3) + Size: 24B + array([0, 1, 2]) + >>> arange(3.0) + Size: 24B + array([0., 1., 2.]) + >>> arange(3, 7) + Size: 32B + array([3, 4, 5, 6]) + >>> arange(3, 7, 2) + Size: 16B + array([3, 5]) + + If dims is set: + + >>> arange(3, dims="x") + Size: 24B + array([0, 1, 2]) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.arange(start, stop=stop, step=step, dtype=dtype, device=device) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +@overload +def asarray( + obj: duckarray[_ShapeType, Any], + /, + *, + dtype: _DType, + device: _Device | None = ..., + copy: bool | None = ..., + dims: _DimsLike2 | Default = ..., +) -> NamedArray[_ShapeType, _DType]: ... +@overload +def asarray( + obj: _ArrayLike, + /, + *, + dtype: _DType, + device: _Device | None = ..., + copy: bool | None = ..., + dims: _DimsLike2 | Default = ..., +) -> NamedArray[Any, _DType]: ... +@overload +def asarray( + obj: duckarray[_ShapeType, _DType], + /, + *, + dtype: None, + device: _Device | None = None, + copy: bool | None = None, + dims: _DimsLike2 | Default = ..., +) -> NamedArray[_ShapeType, _DType]: ... +@overload +def asarray( + obj: _ArrayLike, + /, + *, + dtype: None, + device: _Device | None = ..., + copy: bool | None = ..., + dims: _DimsLike2 | Default = ..., +) -> NamedArray[Any, _DType]: ... +def asarray( + obj: duckarray[_ShapeType, _DType] | _ArrayLike, + /, + *, + dtype: _DType | None = None, + device: _Device | None = None, + copy: bool | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType] | NamedArray[Any, Any]: + """ + Create a NamedArray from an array-like object. + + Examples + -------- + Convert a list into an array: + + >>> x = [1, 2] + >>> asarray(x) + Size: 16B + array([1, 2]) + + Existing arrays are not copied: + + >>> import numpy as np + >>> x = NamedArray(("x",), np.array([1, 2])) + >>> asarray(x) is x + True + + If dtype is set, array is copied only if dtype does not match: + + >>> x = NamedArray(("x",), np.array([1, 2], dtype=np.float32)) + >>> asarray(x, dtype=np.float32) is x + True + >>> asarray(x, dtype=np.float64) is x + False + + If dims is set: + + >>> x = [1, 2] + >>> asarray(x, dims="x") + Size: 16B + array([1, 2]) + + If dims is set, array is copied only if dims does not match: + + >>> x = NamedArray(("x",), np.array([1, 2], dtype=np.float32)) + >>> asarray(x, dims="x") is x + True + >>> asarray(x, dims=("x",)) is x + True + >>> asarray(x, dims=("y",)) is x + False + """ + data = obj + if isinstance(data, NamedArray): + xp = _get_data_namespace(data) + _dtype = data.dtype if dtype is None else dtype + new_data = xp.asarray(data._data, dtype=_dtype, device=device, copy=copy) + if new_data is data._data and ( + isinstance(dims, Default) or _normalize_dimensions(dims) == data.dims + ): + return data + else: + return NamedArray( + data.dims if isinstance(dims, Default) else dims, new_data, data.attrs + ) + + xp = _get_namespace(data) + _data = xp.asarray(data, dtype=dtype, device=device, copy=copy) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def empty( + shape: _ShapeType, + *, + dtype: _DType | None = None, + device: _Device | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns an uninitialized array having a specified shape. + + Examples + -------- + >>> import numpy as np + >>> x = empty((2, 2)) + >>> x.dims, x.shape, x.dtype + (('dim_1', 'dim_0'), (2, 2), dtype('float64')) + >>> x = empty([2, 2], dtype=np.int64) + >>> x.dims, x.shape, x.dtype + (('dim_1', 'dim_0'), (2, 2), dtype('int64')) + + If dims is set: + + >>> x = empty((2, 2), dims=("x", "y")) + >>> x.dims, x.shape, x.dtype + (('x', 'y'), (2, 2), dtype('float64')) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.empty(shape, dtype=dtype, device=device) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def empty_like( + x: NamedArray[_ShapeType, _DType], + /, + *, + dtype: _DType | None = None, + device: _Device | None = None, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns an uninitialized array with the same shape as an input array x. + + Examples + -------- + >>> import numpy as np + >>> x = empty_like(NamedArray(("x", "y"), np.array([[1, 2, 3], [4, 5, 6]]))) + >>> x.dims, x.shape, x.dtype + (('x', 'y'), (2, 3), dtype('int64')) + """ + xp = _get_data_namespace(x) + _data = xp.empty_like(x._data, dtype=dtype, device=device) + return x._new(data=_data) + + +def eye( + n_rows: int, + n_cols: int | None = None, + /, + *, + k: int = 0, + dtype: _DType | None = None, + device: _Device | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a two-dimensional array with ones on the kth diagonal and zeros elsewhere. + + Examples + -------- + >>> import numpy as np + >>> eye(2, dtype=np.int64) + Size: 32B + array([[1, 0], + [0, 1]]) + >>> eye(3, k=1) + Size: 72B + array([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) + + If dims is set: + + >>> eye(2, dims=("x", "y")) + Size: 32B + array([[1., 0.], + [0., 1.]]) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.eye(n_rows, n_cols, k=k, dtype=dtype, device=device) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def from_dlpack( + x: object, + /, + *, + device: _Device | None = None, + copy: bool | None = None, +) -> NamedArray[Any, Any]: + """ + Returns a new array containing the data from another (array) object with a __dlpack__ method. + """ + if isinstance(x, NamedArray): + xp = _get_data_namespace(x) + _device = x.device if device is None else device + _data = xp.from_dlpack(x, device=_device, copy=copy) + _dims = x.dims + else: + xp = _get_namespace(x) + _device = device + _data = xp.from_dlpack(x, device=_device, copy=copy) + _dims = _infer_dims(_data.shape) + return NamedArray(_dims, _data) + + +def full( + shape: _Shape, + fill_value: bool | int | float | complex, + *, + dtype: _DType | None = None, + device: _Device | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a new array having a specified shape and filled with fill_value. + + Examples + -------- + >>> import numpy as np + >>> full((2, 2), np.inf) + Size: 32B + array([[inf, inf], + [inf, inf]]) + >>> full((2, 2), 10) + Size: 32B + array([[10, 10], + [10, 10]]) + + If dims is set: + >>> full((2, 2), np.inf, dims=("x", "y")) + Size: 32B + array([[inf, inf], + [inf, inf]]) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.full(shape, fill_value, dtype=dtype, device=device) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def full_like( + x: NamedArray[_ShapeType, _DType], + /, + fill_value: bool | int | float | complex, + *, + dtype: _DType | None = None, + device: _Device | None = None, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a new array filled with fill_value and having the same shape as an input array x. + + Examples + -------- + >>> import numpy as np + >>> x = arange(6, dtype=np.int64, dims=("x",)) + >>> full_like(x, 1) + Size: 48B + array([1, 1, 1, 1, 1, 1]) + >>> full_like(x, 0.1) + Size: 48B + array([0, 0, 0, 0, 0, 0]) + >>> full_like(x, 0.1, dtype=np.float32) + Size: 24B + array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1], dtype=float32) + >>> full_like(x, np.nan, dtype=np.float32) + Size: 24B + array([nan, nan, nan, nan, nan, nan], dtype=float32) + """ + xp = _get_data_namespace(x) + _data = xp.full_like(x._data, fill_value, dtype=dtype, device=device) + return x._new(data=_data) + + +def linspace( + start: int | float | complex, + stop: int | float | complex, + /, + num: int, + *, + dtype: _DType | None = None, + device: _Device | None = None, + endpoint: bool = True, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_Shape1D, _DType]: + """ + Returns evenly spaced numbers over a specified interval. + + Examples + -------- + >>> import numpy as np + >>> linspace(2.0, 3.0, num=5) + Size: 40B + array([2. , 2.25, 2.5 , 2.75, 3. ]) + >>> linspace(2.0, 3.0, num=5, endpoint=False) + Size: 40B + array([2. , 2.2, 2.4, 2.6, 2.8]) + + If dims is set: + + >>> linspace(2.0, 3.0, num=5, dims=("x",)) + Size: 40B + array([2. , 2.25, 2.5 , 2.75, 3. ]) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.linspace( + start, + stop, + num=num, + dtype=dtype, + device=device, + endpoint=endpoint, + ) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def meshgrid( + *arrays: NamedArray[Any, Any], indexing: str = "xy" +) -> list[NamedArray[Any, Any]]: + """ + Returns coordinate matrices from coordinate vectors. + + Note: broadcast_arrays might be a better option for namedarrays. + + Parameters + ---------- + *arrays : NamedArray[Any, Any] + An arbitrary number of one-dimensional arrays representing grid coordinates. + Each array should have the same numeric data type. + indexing : str, optional + Cartesian 'xy' or matrix 'ij' indexing of output. If provided zero or one + one-dimensional vector(s) (i.e., the zero- and one-dimensional cases, respectively), + the indexing keyword has no effect and should be ignored. Default: 'xy'. + + + Examples + -------- + >>> nx, ny = (3, 2) + >>> x = linspace(0, 1, nx, dims="x") + >>> meshgrid(x, indexing="xy")[0] + Size: 24B + array([0. , 0.5, 1. ]) + >>> y = linspace(0, 1, ny, dims="y") + >>> xv, yv = meshgrid(x, y, indexing="xy") + >>> xv + Size: 48B + array([[0. , 0.5, 1. ], + [0. , 0.5, 1. ]]) + >>> yv + Size: 48B + array([[0., 0., 0.], + [1., 1., 1.]]) + + With matrix indexing + + >>> nx, ny = (3, 2) + >>> x = linspace(0, 1, nx, dims="x") + >>> meshgrid(x, indexing="ij")[0] + Size: 24B + array([0. , 0.5, 1. ]) + >>> y = linspace(0, 1, ny, dims="y") + >>> xv, yv = meshgrid(x, y, indexing="ij") + >>> xv + Size: 48B + array([[0. , 0. ], + [0.5, 0.5], + [1. , 1. ]]) + >>> yv + Size: 48B + array([[0., 1.], + [0., 1.], + [0., 1.]]) + + Empty arrays: + + >>> x0, x1, x2 = ones(()), ones((2,), dims=("x",)), ones((3,), dims=("y",)) + >>> meshgrid() + [] + >>> meshgrid(x0) + [ Size: 8B + array([1.])] + >>> meshgrid(x0, x1) + [ Size: 16B + array([[1.], + [1.]]), Size: 16B + array([[1.], + [1.]])] + >>> meshgrid(x0, x1, x2, indexing="ij") + [ Size: 48B + array([[[1., 1., 1.], + [1., 1., 1.]]]), Size: 48B + array([[[1., 1., 1.], + [1., 1., 1.]]]), Size: 48B + array([[[1., 1., 1.], + [1., 1., 1.]]])] + + TODO: Problems with meshgrid: + + >>> # Adds another dimension even though only 2 distinct dims are used: + >>> x212 = meshgrid(x2, x1, x2, indexing="ij") + >>> x212[0].shape, x212[0].dims + ((3, 2, 3), ('dim_0', 'x', 'y')) + """ + if len(arrays) == 0: + return [] + + arr = arrays[0] + xp = _get_data_namespace(arr) + _datas = xp.meshgrid(*[a._data for a in arrays], indexing=indexing) + _dims, _ = _broadcast_dims(*arrays) + _dims = _infer_dims(_datas[0].shape, _dims) + if indexing == "xy" and len(_dims) > 1: + _dims = _move_dims(_dims, 0, 1) + return [arr._new(_dims, _data) for _data in _datas] + + +def ones( + shape: _Shape, + *, + dtype: _DType | None = None, + device: _Device | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a new array having a specified shape and filled with ones. + + Examples + >>> import numpy as np + >>> ones((2, 2)) + Size: 32B + array([[1., 1.], + [1., 1.]]) + + If dims is set: + >>> ones((2, 2), dims=("x", "y")) + Size: 32B + array([[1., 1.], + [1., 1.]]) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.ones(shape, dtype=dtype, device=device) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def ones_like( + x: NamedArray[_ShapeType, _DType], + /, + *, + dtype: _DType | None = None, + device: _Device | None = None, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a new array filled with ones and having the same shape as an input array x. + + Examples + -------- + >>> import numpy as np + >>> ones_like(NamedArray(("x", "y"), np.array([[1, 2, 3], [4, 5, 6]]))) + Size: 48B + array([[1, 1, 1], + [1, 1, 1]]) + """ + xp = _get_data_namespace(x) + _data = xp.ones_like(x._data, dtype=dtype, device=device) + return x._new(data=_data) + + +def tril( + x: NamedArray[_ShapeType, _DType], /, *, k: int = 0 +) -> NamedArray[_ShapeType, _DType]: + """ + Returns the lower triangular part of a matrix (or a stack of matrices) x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y"), np.array([[1, 2, 3], [4, 5, 6]])) + >>> tril(x) + Size: 48B + array([[1, 0, 0], + [4, 5, 0]]) + """ + xp = _get_data_namespace(x) + _data = xp.tril(x._data, k=k) + return NamedArray(x.dims, _data) + + +def triu( + x: NamedArray[_ShapeType, _DType], /, *, k: int = 0 +) -> NamedArray[_ShapeType, _DType]: + """ + Returns the upper triangular part of a matrix (or a stack of matrices) x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y"), np.array([[1, 2, 3], [4, 5, 6]])) + >>> triu(x) + Size: 48B + array([[1, 2, 3], + [0, 5, 6]]) + """ + xp = _get_data_namespace(x) + _data = xp.triu(x._data, k=k) + return NamedArray(x.dims, _data) + + +def zeros( + shape: _Shape, + *, + dtype: _DType | None = None, + device: _Device | None = None, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a new array having a specified shape and filled with ones. + + Examples + >>> import numpy as np + >>> zeros((2, 2)) + Size: 32B + array([[0., 0.], + [0., 0.]]) + + If dims is set: + >>> zeros((2, 2), dims=("x", "y")) + Size: 32B + array([[0., 0.], + [0., 0.]]) + """ + xp = _get_namespace_dtype(dtype) + _data = xp.zeros(shape, dtype=dtype, device=device) + _dims = _infer_dims(_data.shape, dims) + return NamedArray(_dims, _data) + + +def zeros_like( + x: NamedArray[_ShapeType, _DType], + /, + *, + dtype: _DType | None = None, + device: _Device | None = None, +) -> NamedArray[_ShapeType, _DType]: + """ + Returns a new array filled with zeros and having the same shape as an input array x. + + Examples + -------- + >>> import numpy as np + >>> zeros_like(NamedArray(("x", "y"), np.array([[1, 2, 3], [4, 5, 6]]))) + Size: 48B + array([[0, 0, 0], + [0, 0, 0]]) + """ + xp = _get_data_namespace(x) + _data = xp.zeros_like(x._data, dtype=dtype, device=device) + return x._new(data=_data) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_array_api/_data_type_functions.py b/xarray/namedarray/_array_api/_data_type_functions.py new file mode 100644 index 00000000000..fc728be4a2c --- /dev/null +++ b/xarray/namedarray/_array_api/_data_type_functions.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from typing import Any, cast + +from xarray.namedarray._array_api._utils import ( + _get_data_namespace, + _get_namespace_dtype, +) +from xarray.namedarray._typing import ( + _arrayapi, + _Device, + _DType, + _dtype, + _FInfo, + _IInfo, + _ShapeType, +) +from xarray.namedarray.core import NamedArray + + +def astype( + x: NamedArray[_ShapeType, Any], + dtype: _DType, + /, + *, + copy: bool = True, + device: _Device | None = None, +) -> NamedArray[_ShapeType, _DType]: + """ + Copies an array to a specified data type irrespective of Type Promotion Rules rules. + + Parameters + ---------- + x : NamedArray + Array to cast. + dtype : _DType + Desired data type. + copy : bool, optional + Specifies whether to copy an array when the specified dtype matches the data + type of the input array x. + If True, a newly allocated array must always be returned. + If False and the specified dtype matches the data type of the input array, + the input array must be returned; otherwise, a newly allocated array must be + returned. Default: True. + + Returns + ------- + out : NamedArray + An array having the specified data type. The returned array must have the + same shape as x. + + Examples + -------- + >>> narr = NamedArray(("x",), np.asarray([1.5, 2.5])) + >>> narr + Size: 16B + array([1.5, 2.5]) + >>> astype(narr, np.dtype(np.int32)) + Size: 8B + array([1, 2], dtype=int32) + """ + if isinstance(x._data, _arrayapi): + xp = x._data.__array_namespace__() + return x._new(data=xp.astype(x._data, dtype, copy=copy)) + + # np.astype doesn't exist yet: + return x._new(data=x._data.astype(dtype, copy=copy)) # type: ignore[attr-defined] + + +def can_cast(from_: _dtype[Any] | NamedArray[Any, Any], to: _dtype[Any], /) -> bool: + if isinstance(from_, NamedArray): + xp = _get_data_namespace(from_) + from_ = from_.dtype + return cast(bool, xp.can_cast(from_, to)) # TODO: Why is cast necessary? + else: + xp = _get_namespace_dtype(from_) + return cast(bool, xp.can_cast(from_, to)) # TODO: Why is cast necessary? + + +def finfo(type: _dtype[Any] | NamedArray[Any, Any], /) -> _FInfo: + if isinstance(type, NamedArray): + xp = _get_data_namespace(type) + return cast(_FInfo, xp.finfo(type._data)) # TODO: Why is cast necessary? + else: + xp = _get_namespace_dtype(type) + return cast(_FInfo, xp.finfo(type)) # TODO: Why is cast necessary? + + +def iinfo(type: _dtype[Any] | NamedArray[Any, Any], /) -> _IInfo: + if isinstance(type, NamedArray): + xp = _get_data_namespace(type) + return cast(_IInfo, xp.iinfo(type._data)) # TODO: Why is cast necessary? + else: + xp = _get_namespace_dtype(type) + return cast(_IInfo, xp.iinfo(type)) # TODO: Why is cast necessary? + + +def isdtype( + dtype: _dtype[Any], kind: _dtype[Any] | str | tuple[_dtype[Any] | str, ...] +) -> bool: + xp = _get_namespace_dtype(dtype) + return cast(bool, xp.isdtype(dtype, kind)) # TODO: Why is cast necessary? + + +def result_type(*arrays_and_dtypes: NamedArray[Any, Any] | _dtype[Any]) -> _dtype[Any]: + # TODO: Empty arg? + arr_or_dtype = arrays_and_dtypes[0] + if isinstance(arr_or_dtype, NamedArray): + xp = _get_data_namespace(arr_or_dtype) + else: + xp = _get_namespace_dtype(arr_or_dtype) + + return cast( + _dtype[Any], + xp.result_type( + *(a.dtype if isinstance(a, NamedArray) else a for a in arrays_and_dtypes) + ), + ) # TODO: Why is cast necessary? diff --git a/xarray/namedarray/_array_api/_dtypes.py b/xarray/namedarray/_array_api/_dtypes.py new file mode 100644 index 00000000000..c796400329a --- /dev/null +++ b/xarray/namedarray/_array_api/_dtypes.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from xarray.namedarray._array_api._utils import _maybe_default_namespace + +_xp = _maybe_default_namespace() +int8 = _xp.int8 +int16 = _xp.int16 +int32 = _xp.int32 +int64 = _xp.int64 +uint8 = _xp.uint8 +uint16 = _xp.uint16 +uint32 = _xp.uint32 +uint64 = _xp.uint64 +float32 = _xp.float32 +float64 = _xp.float64 +complex64 = _xp.complex64 +complex128 = _xp.complex128 +bool = _xp.bool + +_all_dtypes = ( + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + complex64, + complex128, + bool, +) +_boolean_dtypes = (bool,) +_real_floating_dtypes = (float32, float64) +_floating_dtypes = (float32, float64, complex64, complex128) +_complex_floating_dtypes = (complex64, complex128) +_integer_dtypes = (int8, int16, int32, int64, uint8, uint16, uint32, uint64) +_signed_integer_dtypes = (int8, int16, int32, int64) +_unsigned_integer_dtypes = (uint8, uint16, uint32, uint64) +_integer_or_boolean_dtypes = ( + bool, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, +) +_real_numeric_dtypes = ( + float32, + float64, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, +) +_numeric_dtypes = ( + float32, + float64, + complex64, + complex128, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, +) + +_dtype_categories = { + "all": _all_dtypes, + "real numeric": _real_numeric_dtypes, + "numeric": _numeric_dtypes, + "integer": _integer_dtypes, + "integer or boolean": _integer_or_boolean_dtypes, + "boolean": _boolean_dtypes, + "real floating-point": _floating_dtypes, + "complex floating-point": _complex_floating_dtypes, + "floating-point": _floating_dtypes, +} diff --git a/xarray/namedarray/_array_api/_elementwise_functions.py b/xarray/namedarray/_array_api/_elementwise_functions.py new file mode 100644 index 00000000000..8517f3f91e6 --- /dev/null +++ b/xarray/namedarray/_array_api/_elementwise_functions.py @@ -0,0 +1,607 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np + +from xarray.namedarray._array_api._manipulation_functions import _arithmetic_broadcast +from xarray.namedarray._array_api._utils import ( + _get_data_namespace, +) +from xarray.namedarray._typing import ( + _DType, + _ScalarType, + _ShapeType, + _SupportsImag, + _SupportsReal, +) +from xarray.namedarray.core import NamedArray + + +def abs(x: NamedArray[_ShapeType, _DType], /) -> NamedArray[_ShapeType, _DType]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.abs(x._data) + return x._new(_dims, _data) + + +def acos(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.acos(x._data) + return x._new(_dims, _data) + + +def acosh(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.acosh(x._data) + return x._new(_dims, _data) + + +def add(x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.add(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def asin(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.asin(x._data) + return x._new(_dims, _data) + + +def asinh(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.asinh(x._data) + return x._new(_dims, _data) + + +def atan(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.atan(x._data) + return x._new(_dims, _data) + + +def atan2( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.atan2(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def atanh(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.atanh(x._data) + return x._new(_dims, _data) + + +def bitwise_and( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.bitwise_and(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def bitwise_invert(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.bitwise_invert(x._data) + return x._new(_dims, _data) + + +def bitwise_left_shift( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.bitwise_left_shift(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def bitwise_or( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.bitwise_or(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def bitwise_right_shift( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.bitwise_right_shift(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def bitwise_xor( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.bitwise_xor(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def ceil(x: NamedArray[_ShapeType, _DType], /) -> NamedArray[_ShapeType, _DType]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.ceil(x._data) + return x._new(_dims, _data) + + +def clip( + x: NamedArray[Any, Any], + /, + min: int | float | NamedArray[Any, Any] | None = None, + max: int | float | NamedArray[Any, Any] | None = None, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.clip(x._data, min=min, max=max) + return x._new(_dims, _data) + + +def conj(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.conj(x._data) + return x._new(_dims, _data) + + +def copysign( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.copysign(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def cos(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.cos(x._data) + return x._new(_dims, _data) + + +def cosh(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.cosh(x._data) + return x._new(_dims, _data) + + +def divide( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.divide(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def exp(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.exp(x._data) + return x._new(_dims, _data) + + +def expm1(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.expm1(x._data) + return x._new(_dims, _data) + + +def equal( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.equal(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def floor(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.floor(x._data) + return x._new(_dims, _data) + + +def floor_divide( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.floor_divide(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def greater( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.greater(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def greater_equal( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.greater_equal(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def hypot( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.hypot(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def imag( + x: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], + /, # type: ignore[type-var] +) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: + """ + Returns the imaginary component of a complex number for each element x_i of the + input array x. + + Parameters + ---------- + x : NamedArray + Input array. Should have a complex floating-point data type. + + Returns + ------- + out : NamedArray + An array containing the element-wise results. The returned array must have a + floating-point data type with the same floating-point precision as x + (e.g., if x is complex64, the returned array must have the floating-point + data type float32). + + Examples + -------- + >>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) + >>> imag(narr) + Size: 16B + array([2., 4.]) + """ + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.imag(x._data) + return x._new(_dims, _data) + + +def isfinite(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.isfinite(x._data) + return x._new(_dims, _data) + + +def isinf(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.isinf(x._data) + return x._new(_dims, _data) + + +def isnan(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.isnan(x._data) + return x._new(_dims, _data) + + +def less(x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.less(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def less_equal( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.less_equal(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def log(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.log(x._data) + return x._new(_dims, _data) + + +def log1p(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.log1p(x._data) + return x._new(_dims, _data) + + +def log2(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.log2(x._data) + return x._new(_dims, _data) + + +def log10(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.log10(x._data) + return x._new(_dims, _data) + + +def logaddexp( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.logaddexp(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def logical_and( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.logical_and(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def logical_not(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.logical_not(x._data) + return x._new(_dims, _data) + + +def logical_or( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.logical_or(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def logical_xor( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.logical_xor(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def maximum( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.maximum(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def minimum( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.minimum(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def multiply( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.multiply(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def negative(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.negative(x._data) + return x._new(_dims, _data) + + +def not_equal( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.not_equal(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def positive(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.positive(x._data) + return x._new(_dims, _data) + + +def pow(x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.pow(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def real( + x: NamedArray[_ShapeType, np.dtype[_SupportsReal[_ScalarType]]], + /, # type: ignore[type-var] +) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]: + """ + Returns the real component of a complex number for each element x_i of the + input array x. + + Parameters + ---------- + x : NamedArray + Input array. Should have a complex floating-point data type. + + Returns + ------- + out : NamedArray + An array containing the element-wise results. The returned array must have a + floating-point data type with the same floating-point precision as x + (e.g., if x is complex64, the returned array must have the floating-point + data type float32). + + Examples + -------- + >>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) + >>> real(narr) + Size: 16B + array([1., 2.]) + """ + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.real(x._data) + return x._new(_dims, _data) + + +def remainder( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.remainder(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def round(x: NamedArray[_ShapeType, _DType], /) -> NamedArray[_ShapeType, _DType]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.round(x._data) + return x._new(_dims, _data) + + +def sign(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.sign(x._data) + return x._new(_dims, _data) + + +def signbit(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.signbit(x._data) + return x._new(_dims, _data) + + +def sin(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.sin(x._data) + return x._new(_dims, _data) + + +def sinh(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.sinh(x._data) + return x._new(_dims, _data) + + +def sqrt(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.sqrt(x._data) + return x._new(_dims, _data) + + +def square(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.square(x._data) + return x._new(_dims, _data) + + +def subtract( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + x1_new, x2_new = _arithmetic_broadcast(x1, x2) + _dims = x1_new.dims + _data = xp.subtract(x1_new._data, x2_new._data) + return x1._new(_dims, _data) + + +def tan(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.tan(x._data) + return x._new(_dims, _data) + + +def tanh(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.tanh(x._data) + return x._new(_dims, _data) + + +def trunc(x: NamedArray[_ShapeType, _DType], /) -> NamedArray[_ShapeType, _DType]: + xp = _get_data_namespace(x) + _dims = x.dims + _data = xp.trunc(x._data) + return x._new(_dims, _data) diff --git a/xarray/namedarray/_array_api/_fft/__init__.py b/xarray/namedarray/_array_api/_fft/__init__.py new file mode 100644 index 00000000000..ac2fa15deff --- /dev/null +++ b/xarray/namedarray/_array_api/_fft/__init__.py @@ -0,0 +1,35 @@ +__all__ = [] + +from xarray.namedarray._array_api._fft._fft import ( + fft, + fftfreq, + fftn, + fftshift, + hfft, + ifft, + ifftn, + ifftshift, + ihfft, + irfft, + irfftn, + rfft, + rfftfreq, + rfftn, +) + +__all__ = [ + "fft", + "ifft", + "fftn", + "ifftn", + "rfft", + "irfft", + "rfftn", + "irfftn", + "hfft", + "ihfft", + "fftfreq", + "rfftfreq", + "fftshift", + "ifftshift", +] diff --git a/xarray/namedarray/_array_api/_fft/_fft.py b/xarray/namedarray/_array_api/_fft/_fft.py new file mode 100644 index 00000000000..9830581080f --- /dev/null +++ b/xarray/namedarray/_array_api/_fft/_fft.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal + +from xarray.namedarray._array_api._utils import ( + _get_data_namespace, + _infer_dims, + _maybe_default_namespace, +) +from xarray.namedarray.core import NamedArray + +if TYPE_CHECKING: + from xarray.namedarray._typing import _Axes, _Axis, _Device, _DType + + _Norm = Literal["backward", "ortho", "forward"] + + +def fft( + x: NamedArray[Any, _DType], + /, + *, + n: int | None = None, + axis: _Axis = -1, + norm: _Norm = "backward", +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.fft.fft(x._data, n=n, axis=axis, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def ifft( + x: NamedArray[Any, _DType], + /, + *, + n: int | None = None, + axis: _Axis = -1, + norm: _Norm = "backward", +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.fft.ifft(x._data, n=n, axis=axis, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def fftn( + x: NamedArray[Any, _DType], + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.fft.fftn(x._data, s=s, axes=axes, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def ifftn( + x: NamedArray[Any, _DType], + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.fft.ifftn(x._data, s=s, axes=axes, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def rfft( + x: NamedArray[Any, Any], + /, + *, + n: int | None = None, + axis: _Axis = -1, + norm: _Norm = "backward", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.fft.rfft(x._data, n=n, axis=axis, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def irfft( + x: NamedArray[Any, Any], + /, + *, + n: int | None = None, + axis: _Axis = -1, + norm: _Norm = "backward", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.fft.irfft(x._data, n=n, axis=axis, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def rfftn( + x: NamedArray[Any, Any], + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.fft.rfftn(x._data, s=s, axes=axes, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def irfftn( + x: NamedArray[Any, Any], + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.fft.irfftn(x._data, s=s, axes=axes, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def hfft( + x: NamedArray[Any, Any], + /, + *, + n: int | None = None, + axis: _Axis = -1, + norm: _Norm = "backward", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.fft.hfft(x._data, n=n, axis=axis, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def ihfft( + x: NamedArray[Any, Any], + /, + *, + n: int | None = None, + axis: _Axis = -1, + norm: _Norm = "backward", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.fft.ihfft(x._data, n=n, axis=axis, norm=norm) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def fftfreq( + n: int, /, *, d: float = 1.0, device: _Device | None = None +) -> NamedArray[Any, Any]: + xp = _maybe_default_namespace() # TODO: Can use device? + _data = xp.fft.fftfreq(n, d=d, device=device) + _dims = _infer_dims(_data.shape) + return NamedArray(_dims, _data) + + +def rfftfreq( + n: int, /, *, d: float = 1.0, device: _Device | None = None +) -> NamedArray[Any, Any]: + xp = _maybe_default_namespace() # TODO: Can use device? + _data = xp.fft.rfftfreq(n, d=d, device=device) + _dims = _infer_dims(_data.shape) + return NamedArray(_dims, _data) + + +def fftshift( + x: NamedArray[Any, _DType], /, *, axes: _Axes | None = None +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.fft.fftshift(x._data, axes=axes) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def ifftshift( + x: NamedArray[Any, _DType], /, *, axes: _Axes | None = None +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.fft.ifftshift(x._data, axes=axes) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) diff --git a/xarray/namedarray/_array_api/_indexing_functions.py b/xarray/namedarray/_array_api/_indexing_functions.py new file mode 100644 index 00000000000..1a4dac96ce8 --- /dev/null +++ b/xarray/namedarray/_array_api/_indexing_functions.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from typing import Any + +from xarray.namedarray._array_api._utils import _dims_to_axis, _get_data_namespace +from xarray.namedarray._typing import ( + Default, + _default, + _Dim, + _DType, +) +from xarray.namedarray.core import NamedArray + + +def take( + x: NamedArray[Any, _DType], + indices: NamedArray[Any, Any], + /, + *, + dim: _Dim | Default = _default, + axis: int | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dim, axis)[0] + # TODO: Handle attrs? will get x1 now + out = x._new(data=xp.take(x._data, indices._data, axis=_axis)) + return out diff --git a/xarray/namedarray/_array_api/_info.py b/xarray/namedarray/_array_api/_info.py new file mode 100644 index 00000000000..9876a2fea01 --- /dev/null +++ b/xarray/namedarray/_array_api/_info.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from xarray.namedarray._array_api._dtypes import ( + bool, + complex64, + complex128, + float32, + float64, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, +) + +if TYPE_CHECKING: + from types import ModuleType + + from xarray.namedarray._typing import ( + _Capabilities, + _DataTypes, + _DefaultDataTypes, + _Device, + ) + + +def __array_namespace_info__() -> ModuleType: + import xarray.namedarray._array_api._info + + return xarray.namedarray._array_api._info + + +def capabilities() -> _Capabilities: + return { + "boolean indexing": False, + "data-dependent shapes": False, + } + + +def default_device() -> _Device: + from xarray.namedarray._array_api._utils import _maybe_default_namespace + + xp = _maybe_default_namespace() + info = xp.__array_namespace_info__() + return info.default_device + + +def default_dtypes( + *, + device: _Device | None = None, +) -> _DefaultDataTypes: + return { + "real floating": float64, + "complex floating": complex128, + "integral": int64, + "indexing": int64, + } + + +def dtypes( + *, + device: _Device | None = None, + kind: str | tuple[str, ...] | None = None, +) -> _DataTypes: + if kind is None: + return { + "bool": bool, + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + "uint16": uint16, + "uint32": uint32, + "uint64": uint64, + "float32": float32, + "float64": float64, + "complex64": complex64, + "complex128": complex128, + } + if kind == "bool": + return {"bool": bool} + if kind == "signed integer": + return { + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + } + if kind == "unsigned integer": + return { + "uint8": uint8, + "uint16": uint16, + "uint32": uint32, + "uint64": uint64, + } + if kind == "integral": + return { + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + "uint16": uint16, + "uint32": uint32, + "uint64": uint64, + } + if kind == "real floating": + return { + "float32": float32, + "float64": float64, + } + if kind == "complex floating": + return { + "complex64": complex64, + "complex128": complex128, + } + if kind == "numeric": + return { + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + "uint16": uint16, + "uint32": uint32, + "uint64": uint64, + "float32": float32, + "float64": float64, + "complex64": complex64, + "complex128": complex128, + } + if isinstance(kind, tuple): + res: _DataTypes = {} + for k in kind: + res.update(dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + +def devices() -> list[_Device]: + return [default_device()] + + +__all__ = [ + "capabilities", + "default_device", + "default_dtypes", + "devices", + "dtypes", +] diff --git a/xarray/namedarray/_array_api/_linalg/__init__.py b/xarray/namedarray/_array_api/_linalg/__init__.py new file mode 100644 index 00000000000..d86858242e4 --- /dev/null +++ b/xarray/namedarray/_array_api/_linalg/__init__.py @@ -0,0 +1,53 @@ +__all__ = [] + +from xarray.namedarray._array_api._linalg._linalg import ( + cholesky, + cross, + det, + diagonal, + eigh, + eigvalsh, + inv, + matmul, + matrix_norm, + matrix_power, + matrix_rank, + matrix_transpose, + outer, + pinv, + qr, + slogdet, + solve, + svd, + svdvals, + tensordot, + trace, + vecdot, + vector_norm, +) + +__all__ = [ + "cholesky", + "cross", + "det", + "diagonal", + "eigh", + "eigvalsh", + "inv", + "matmul", + "matrix_norm", + "matrix_power", + "matrix_rank", + "matrix_transpose", + "outer", + "pinv", + "qr", + "slogdet", + "solve", + "svd", + "svdvals", + "tensordot", + "trace", + "vecdot", + "vector_norm", +] diff --git a/xarray/namedarray/_array_api/_linalg/_linalg.py b/xarray/namedarray/_array_api/_linalg/_linalg.py new file mode 100644 index 00000000000..0c643e870f3 --- /dev/null +++ b/xarray/namedarray/_array_api/_linalg/_linalg.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, overload + +from xarray.namedarray._array_api._utils import _get_data_namespace, _infer_dims +from xarray.namedarray.core import NamedArray + +if TYPE_CHECKING: + from xarray.namedarray._typing import _Axes, _Axis, _DType, _ShapeType + + +class EighResult(NamedTuple): + eigenvalues: NamedArray[Any, Any] + eigenvectors: NamedArray[Any, Any] + + +class QRResult(NamedTuple): + Q: NamedArray[Any, Any] + R: NamedArray[Any, Any] + + +class SlogdetResult(NamedTuple): + sign: NamedArray[Any, Any] + logabsdet: NamedArray[Any, Any] + + +class SVDResult(NamedTuple): + U: NamedArray[Any, Any] + S: NamedArray[Any, Any] + Vh: NamedArray[Any, Any] + + +def cholesky( + x: NamedArray[_ShapeType, Any], /, *, upper: bool = False +) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.cholesky(x._data, upper=upper) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +# Note: cross is the numpy top-level namespace, not np.linalg +def cross( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], /, *, axis: _Axis = -1 +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + _data = xp.linalg.cross(x1._data, x2._data, axis=axis) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x1._new(_dims, _data) + + +def det(x: NamedArray[Any, _DType], /) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.linalg.det(x._data) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def diagonal( + x: NamedArray[Any, _DType], /, *, offset: int = 0 +) -> NamedArray[Any, _DType]: + # Note: diagonal is the numpy top-level namespace, not np.linalg + xp = _get_data_namespace(x) + _data = xp.linalg.diagonal(x._data, offset=offset) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def eigh(x: NamedArray[Any, Any], /) -> EighResult: + xp = _get_data_namespace(x) + eigvals, eigvecs = xp.linalg.eigh(x._data) + _dims_vals = _infer_dims(eigvals.shape) # TODO: Fix dims + _dims_vecs = _infer_dims(eigvecs.shape) # TODO: Fix dims + return EighResult( + x._new(_dims_vals, eigvals), + x._new(_dims_vecs, eigvecs), + ) + + +def eigvalsh(x: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.eigvalsh(x._data) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def inv(x: NamedArray[_ShapeType, Any], /) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.inv(x._data) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def matrix_norm( + x: NamedArray[Any, Any], + /, + *, + keepdims: bool = False, + ord: int | float | Literal["fro", "nuc"] | None = "fro", +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.matrix_norm(x._data, keepdims=keepdims, ord=ord) # ckeck xp.mean + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def matrix_power( + x: NamedArray[_ShapeType, Any], n: int, / +) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.matrix_power(x._data, n=n) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def matrix_rank( + x: NamedArray[Any, Any], /, *, rtol: float | NamedArray[Any, Any] | None = None +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.matrix_rank(x._data, rtol=rtol) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def matrix_transpose(x: NamedArray[Any, _DType], /) -> NamedArray[Any, _DType]: + from xarray.namedarray._array_api._linear_algebra_functions import matrix_transpose + + return matrix_transpose(x) + + +def outer( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + _data = xp.linalg.outer(x1._data, x2._data) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x1._new(_dims, _data) + + +def pinv( + x: NamedArray[Any, Any], /, *, rtol: float | NamedArray[Any, Any] | None = None +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.pinv(x._data, rtol=rtol) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def qr( + x: NamedArray[Any, Any], /, *, mode: Literal["reduced", "complete"] = "reduced" +) -> QRResult: + xp = _get_data_namespace(x) + q, r = xp.linalg.qr(x._data) + _dims_q = _infer_dims(q.shape) # TODO: Fix dims + _dims_r = _infer_dims(r.shape) # TODO: Fix dims + return QRResult( + x._new(_dims_q, q), + x._new(_dims_r, r), + ) + + +def slogdet(x: NamedArray[Any, Any], /) -> SlogdetResult: + xp = _get_data_namespace(x) + sign, logabsdet = xp.linalg.slogdet(x._data) + _dims_sign = _infer_dims(sign.shape) # TODO: Fix dims + _dims_logabsdet = _infer_dims(logabsdet.shape) # TODO: Fix dims + return SlogdetResult( + x._new(_dims_sign, sign), + x._new(_dims_logabsdet, logabsdet), + ) + + +def solve( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + _data = xp.linalg.solve(x1._data, x2._data) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x1._new(_dims, _data) + + +def svd(x: NamedArray[Any, Any], /, *, full_matrices: bool = True) -> SVDResult: + xp = _get_data_namespace(x) + u, s, vh = xp.linalg.svd(x._data, full_matrices=full_matrices) + _dims_u = _infer_dims(u.shape) # TODO: Fix dims + _dims_s = _infer_dims(s.shape) # TODO: Fix dims + _dims_vh = _infer_dims(vh.shape) # TODO: Fix dims + return SVDResult( + x._new(_dims_u, u), + x._new(_dims_s, s), + x._new(_dims_vh, vh), + ) + + +def svdvals(x: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.svdvals(x._data) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +@overload +def trace( + x: NamedArray[Any, Any], /, *, offset: int = 0, dtype: _DType +) -> NamedArray[Any, _DType]: ... +@overload +def trace( + x: NamedArray[Any, _DType], /, *, offset: int = 0, dtype: None +) -> NamedArray[Any, _DType]: ... +def trace( + x: NamedArray[Any, _DType | Any], /, *, offset: int = 0, dtype: _DType | None = None +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.linalg.trace(x._data, offset=offset) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def matmul( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + from xarray.namedarray._array_api._linear_algebra_functions import matmul + + return matmul(x1, x2) + + +def tensordot( + x1: NamedArray[Any, Any], + x2: NamedArray[Any, Any], + /, + *, + axes: int | tuple[Sequence[int], Sequence[int]] = 2, +) -> NamedArray[Any, Any]: + from xarray.namedarray._array_api._linear_algebra_functions import tensordot + + return tensordot(x1, x2, axes=axes) + + +def vecdot( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], /, *, axis: _Axis = -1 +) -> NamedArray[Any, Any]: + from xarray.namedarray._array_api._linear_algebra_functions import vecdot + + return vecdot(x1, x2, axis=axis) + + +def vector_norm( + x: NamedArray[Any, Any], + /, + *, + axis: _Axes | None = None, + keepdims: bool = False, + ord: int | float | None = 2, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _data = xp.linalg.vector_norm(x._data, axis=axis, keepdims=keepdims, ord=ord) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) diff --git a/xarray/namedarray/_array_api/_linear_algebra_functions.py b/xarray/namedarray/_array_api/_linear_algebra_functions.py new file mode 100644 index 00000000000..5fb435e5387 --- /dev/null +++ b/xarray/namedarray/_array_api/_linear_algebra_functions.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from xarray.namedarray._array_api._utils import ( + _broadcast_dims, + _get_data_namespace, + _infer_dims, + _reduce_dims, +) +from xarray.namedarray.core import NamedArray + + +def matmul( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / +) -> NamedArray[Any, Any]: + """ + Matrix product of two arrays. + + Examples + -------- + For 2-D arrays it is the matrix product: + + >>> import numpy as np + >>> a = NamedArray(("y", "x"), np.array([[1, 0], [0, 1]])) + >>> b = NamedArray(("y", "x"), np.array([[4, 1], [2, 2]])) + >>> matmul(a, b) + Size: 32B + array([[4, 1], + [2, 2]]) + + For 2-D mixed with 1-D, the result is the usual. + + >>> a = NamedArray(("y", "x"), np.array([[1, 0], [0, 1]])) + >>> b = NamedArray(("x",), np.array([1, 2])) + >>> matmul(a, b) + + Broadcasting is conventional for stacks of arrays + + >>> a = NamedArray(("z", "y", "x"), np.arange(2 * 2 * 4).reshape((2, 2, 4))) + >>> b = NamedArray(("z", "y", "x"), np.arange(2 * 2 * 4).reshape((2, 4, 2))) + >>> matmul(a, b) + Size: 64B + array([[[ 28, 34], + [ 76, 98]], + + [[428, 466], + [604, 658]]]) + """ + xp = _get_data_namespace(x1) + _data = xp.matmul(x1._data, x2._data) + # TODO: Figure out a better way: + _dims = x1.dims # _infer_dims(_data.shape) + return NamedArray(_dims, _data) + + +def tensordot( + x1: NamedArray[Any, Any], + x2: NamedArray[Any, Any], + /, + *, + axes: int | tuple[Sequence[int], Sequence[int]] = 2, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + _data = xp.tensordot(x1._data, x2._data, axes=axes) + # TODO: Figure out a better way: + _dims = _infer_dims(_data.shape) + return NamedArray(_dims, _data) + + +def matrix_transpose(x: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + """ + Transposes a matrix (or a stack of matrices) x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y", "z"), np.zeros((1, 2, 3))) + >>> matrix_transpose(x) + Size: 48B + array([[[0., 0.], + [0., 0.], + [0., 0.]]]) + + >>> x = NamedArray(("x", "y"), np.zeros((2, 3))) + >>> matrix_transpose(x) + Size: 48B + array([[0., 0.], + [0., 0.], + [0., 0.]]) + """ + xp = _get_data_namespace(x) + _data = xp.matrix_transpose(x._data) + d = x.dims + _dims = d[:-2] + d[-2:][::-1] # (..., M, N) -> (..., N, M) + return NamedArray(_dims, _data) + + +def vecdot( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], /, *, axis: int = -1 +) -> NamedArray[Any, Any]: + """ + Computes the (vector) dot product of two arrays. + + Examples + -------- + >>> import numpy as np + >>> v = NamedArray( + ... ("y", "x"), + ... np.array( + ... [[0.0, 5.0, 0.0], [0.0, 0.0, 10.0], [0.0, 6.0, 8.0], [0.0, 6.0, 8.0]] + ... ), + ... ) + >>> n = NamedArray(("x",), np.array([0.0, 0.6, 0.8])) + >>> vecdot(v, n) + Size: 32B + array([ 3., 8., 10., 10.]) + """ + xp = _get_data_namespace(x1) + _data = xp.vecdot(x1._data, x2._data, axis=axis) + d, _ = _broadcast_dims(x1, x2) + _dims = _reduce_dims(d, axis=axis, keepdims=False) + return NamedArray(_dims, _data) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_array_api/_manipulation_functions.py b/xarray/namedarray/_array_api/_manipulation_functions.py new file mode 100644 index 00000000000..17276fc2787 --- /dev/null +++ b/xarray/namedarray/_array_api/_manipulation_functions.py @@ -0,0 +1,671 @@ +from __future__ import annotations + +import math +from typing import Any + +from xarray.namedarray._array_api._data_type_functions import result_type +from xarray.namedarray._array_api._utils import ( + _atleast1d_dims, + _broadcast_dims, + _dim_to_optional_axis, + _dims_from_tuple_indexing, + _dims_to_axis, + _flatten_dims, + _get_data_namespace, + _infer_dims, + _insert_dim, + _move_dims, + _new_unique_dim_name, + _squeeze_dims, +) +from xarray.namedarray._typing import ( + Default, + _Axes, + _Axis, + _AxisLike, + _default, + _Dim, + _Dims, + _DimsLike2, + _DType, + _Shape, + _ShapeType, +) +from xarray.namedarray.core import NamedArray + + +def broadcast_arrays(*arrays: NamedArray[Any, Any]) -> list[NamedArray[Any, Any]]: + """ + Broadcasts one or more arrays against one another. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.zeros((3,))) + >>> y = NamedArray(("y", "x"), np.zeros((2, 1))) + >>> x_new, y_new = broadcast_arrays(x, y) + >>> x_new.dims, x_new.shape, y_new.dims, y_new.shape + (('y', 'x'), (2, 3), ('y', 'x'), (2, 3)) + + Errors + + >>> x = NamedArray(("x",), np.zeros((3,))) + >>> y = NamedArray(("x",), np.zeros((2))) + >>> x_new, y_new = broadcast_arrays(x, y) + Traceback (most recent call last): + ... + ValueError: operands could not be broadcast together with dims = (('x',), ('x',)) and shapes = ((3,), (2,)) + """ + x = arrays[0] + xp = _get_data_namespace(x) + _dims, _ = _broadcast_dims(*arrays) + _arrays = tuple(a._data for a in arrays) + _datas = xp.broadcast_arrays(*_arrays) + return [arr._new(_dims, _data) for arr, _data in zip(arrays, _datas, strict=False)] + + +def broadcast_to( + x: NamedArray[Any, _DType], + /, + shape: _ShapeType, + *, + dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.arange(0, 3)) + >>> broadcast_to(x, (1, 1, 3)) + Size: 24B + array([[[0, 1, 2]]]) + + >>> broadcast_to(x, shape=(1, 1, 3), dims=("y", "x")) + Size: 24B + array([[[0, 1, 2]]]) + """ + xp = _get_data_namespace(x) + _data = xp.broadcast_to(x._data, shape=shape) + _dims = _infer_dims(_data.shape, x.dims if isinstance(dims, Default) else dims) + return x._new(_dims, _data) + + +def concat( + arrays: tuple[NamedArray[Any, Any], ...] | list[NamedArray[Any, Any]], + /, + *, + axis: _Axis | None = 0, +) -> NamedArray[Any, Any]: + """ + Joins a sequence of arrays along an existing axis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.zeros((3,))) + >>> concat((x, 1 + x)) + Size: 48B + array([0., 0., 0., 1., 1., 1.]) + + >>> x = NamedArray(("x", "y"), np.zeros((1, 3))) + >>> concat((x, 1 + x)) + Size: 48B + array([[0., 0., 0.], + [1., 1., 1.]]) + + 0D + + >>> x1 = NamedArray((), np.array(0)) + >>> x2 = NamedArray((), np.array(1)) + >>> concat((x1, x2), axis=None) + Size: 16B + array([0, 1]) + """ + x = arrays[0] + xp = _get_data_namespace(x) + _axis = axis # TODO: add support for dim? + dtype = result_type(*arrays) + _arrays = tuple(a._data for a in arrays) + _data = xp.concat(_arrays, axis=_axis, dtype=dtype) + _dims = _atleast1d_dims(x.dims) if axis is None else x.dims + return NamedArray(_dims, _data) + + +def expand_dims( + x: NamedArray[Any, _DType], + /, + *, + axis: _Axis = 0, + dim: _Dim | Default = _default, +) -> NamedArray[Any, _DType]: + """ + Expands the shape of an array by inserting a new dimension of size one at the + position specified by dims. + + Parameters + ---------- + x : + Array to expand. + dim : + Dimension name. New dimension will be stored in the axis position. + axis : + Axis position (zero-based). Default is 0. + + Returns + ------- + out : + An expanded output array having the same data type as x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y"), np.asarray([[1.0, 2.0], [3.0, 4.0]])) + >>> expand_dims(x) + Size: 32B + array([[[1., 2.], + [3., 4.]]]) + + Specify dimension name + + >>> expand_dims(x, dim="z") + Size: 32B + array([[[1., 2.], + [3., 4.]]]) + """ + # Array Api does not support multiple axes, but maybe in the future: + # https://github.com/data-apis/array-api/issues/760 + # xref: https://github.com/numpy/numpy/blob/3b246c6488cf246d488bbe5726ca58dc26b6ea74/numpy/lib/_shape_base_impl.py#L509C17-L509C24 + xp = _get_data_namespace(x) + _data = xp.expand_dims(x._data, axis=axis) + _dims = _insert_dim(x.dims, dim, axis) + return x._new(_dims, _data) + + +def flip( + x: NamedArray[_ShapeType, _DType], /, *, axis: _Axes | None = None +) -> NamedArray[_ShapeType, _DType]: + """ + Reverse the order of elements in an array along the given axis. + + Examples + -------- + >>> import numpy as np + >>> A = NamedArray(("x",), np.arange(8)) + >>> flip(A, axis=0) + Size: 64B + array([7, 6, 5, 4, 3, 2, 1, 0]) + >>> A = NamedArray(("z", "y", "x"), np.reshape(np.arange(8), (2, 2, 2))) + >>> flip(A, axis=0) + Size: 64B + array([[[4, 5], + [6, 7]], + + [[0, 1], + [2, 3]]]) + """ + xp = _get_data_namespace(x) + _data = xp.flip(x._data, axis=axis) + return x._new(x.dims, _data) + + +def moveaxis( + x: NamedArray[Any, _DType], source: _Axes, destination: _Axes, / +) -> NamedArray[Any, _DType]: + """ + Moves array axes (dimensions) to new positions, while leaving other axes + in their original positions. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("z", "y", "x"), np.zeros((3, 2, 1))) + >>> moveaxis(x, 0, -1) + Size: 48B + array([[[0., 0., 0.]], + + [[0., 0., 0.]]]) + >>> moveaxis(x, (0, 1), (-1, -2)) + Size: 48B + array([[[0., 0., 0.], + [0., 0., 0.]]]) + """ + # TODO: Add a dims argument? + xp = _get_data_namespace(x) + _data = xp.moveaxis(x._data, source=source, destination=destination) + _dims = _move_dims(x.dims, source, destination) + return x._new(_dims, _data) + + +def permute_dims( + x: NamedArray[Any, _DType], + /, + axes: _Axes | None = None, + *, + dims: _Dims | Default = _default, +) -> NamedArray[Any, _DType]: + """ + Permutes the dimensions of an array. + + Parameters + ---------- + x : + Array to permute. + axes : + Permutation of the dimensions of x. + + Returns + ------- + out : + An array with permuted dimensions. The returned array must have the same + data type as x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y", "z"), np.zeros((1, 2, 3))) + >>> permute_dims(x, (2, 1, 0)) + Size: 48B + array([[[0.], + [0.]], + + [[0.], + [0.]], + + [[0.], + [0.]]]) + + >>> permute_dims(x, dims=("y", "x", "z")) + Size: 48B + array([[[0., 0., 0.]], + + [[0., 0., 0.]]]) + """ + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axes) + if _axis is None: + raise TypeError("permute_dims missing argument axes or dims") + old_dims = x.dims + _data = xp.permute_dims(x._data, _axis) + _dims = tuple(old_dims[i] for i in _axis) + return x._new(_dims, _data) + + +def repeat( + x: NamedArray[Any, _DType], + repeats: int | NamedArray[Any, Any], + /, + *, + axis: _Axis | None = None, + dim: _Dim | Default = _default, +) -> NamedArray[Any, _DType]: + """ + Repeats each element of an array a specified number of times on a per-element basis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray((), np.asarray(3)) + >>> repeat(x, 4) + Size: 32B + array([3, 3, 3, 3]) + + >>> x = NamedArray(("y", "x"), np.reshape(np.arange(1, 5), (2, 2))) + >>> repeat(x, 2) + Size: 64B + array([1, 1, 2, 2, 3, 3, 4, 4]) + >>> repeat(x, 3, axis=1) + Size: 96B + array([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 4, 4, 4]]) + >>> repeat(x, NamedArray(("x",), np.asarray([1, 2])), axis=0) # TODD: Correct? + Size: 48B + array([[1, 2], + [3, 4], + [3, 4]]) + """ + xp = _get_data_namespace(x) + _axis = _dim_to_optional_axis(x, dim, axis=axis) + _repeats = repeats if isinstance(repeats, int) else repeats._data + _data = xp.repeat(x._data, _repeats, axis=_axis) + if _axis is None: + _dims = _flatten_dims(_atleast1d_dims(x.dims)) + else: + _dims = _atleast1d_dims(x.dims) + return x._new(_dims, _data) + + +def reshape( + x: NamedArray[Any, _DType], /, shape: _ShapeType, *, copy: bool | None = None +) -> NamedArray[_ShapeType, _DType]: + """ + Reshapes an array without changing its data. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.zeros((3,))) + >>> reshape(x, (-1,)) + Size: 24B + array([0., 0., 0.]) + + To N-dimensions + + >>> reshape(x, (1, -1, 1)) + Size: 24B + array([[[0.], + [0.], + [0.]]]) + + >>> x = NamedArray(("x", "y"), np.zeros((3, 4))) + >>> reshape(x, (-1,)) + Size: 96B + array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]) + """ + xp = _get_data_namespace(x) + _data = xp.reshape(x._data, shape, copy=copy) + + if math.prod(shape) == -1: + # Flattening operations merges all dimensions to 1: + dims_flattened = _flatten_dims(x.dims) + dim = dims_flattened[0] + d = [] + for v in shape: + d.append(dim if v == -1 else _new_unique_dim_name(tuple(d))) + _dims = tuple(d) + else: + # _dims = _infer_dims(_data.shape, x.dims) + _dims = _infer_dims(_data.shape) + + return x._new(_dims, _data) + + +def roll( + x: NamedArray[_ShapeType, _DType], + /, + shift: int | tuple[int, ...], + *, + axis: _AxisLike | None = None, + # dims: _DimsLike2 | Default = _default, +) -> NamedArray[_ShapeType, _DType]: + """ + Rolls array elements along a specified axis. Array elements that roll + beyond the last position are re-introduced at the first position. + Array elements that roll beyond the first position are re-introduced + at the last position. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.arange(10)) + >>> roll(x, 2) + Size: 80B + array([8, 9, 0, 1, 2, 3, 4, 5, 6, 7]) + >>> roll(x, -2) + Size: 80B + array([2, 3, 4, 5, 6, 7, 8, 9, 0, 1]) + + >>> x2 = NamedArray(("y", "x"), np.reshape(np.arange(10), (2, 5))) + >>> roll(x2, 1) + Size: 80B + array([[9, 0, 1, 2, 3], + [4, 5, 6, 7, 8]]) + >>> roll(x2, -1) + Size: 80B + array([[1, 2, 3, 4, 5], + [6, 7, 8, 9, 0]]) + >>> roll(x2, 1, axis=0) + Size: 80B + array([[5, 6, 7, 8, 9], + [0, 1, 2, 3, 4]]) + >>> roll(x2, -1, axis=0) + Size: 80B + array([[5, 6, 7, 8, 9], + [0, 1, 2, 3, 4]]) + >>> roll(x2, 1, axis=1) + Size: 80B + array([[4, 0, 1, 2, 3], + [9, 5, 6, 7, 8]]) + >>> roll(x2, -1, axis=1) + Size: 80B + array([[1, 2, 3, 4, 0], + [6, 7, 8, 9, 5]]) + >>> roll(x2, (1, 1), axis=(1, 0)) + Size: 80B + array([[9, 5, 6, 7, 8], + [4, 0, 1, 2, 3]]) + >>> roll(x2, (2, 1), axis=(1, 0)) + Size: 80B + array([[8, 9, 5, 6, 7], + [3, 4, 0, 1, 2]]) + """ + xp = _get_data_namespace(x) + # _axis = _dims_to_axis(x, dims, axis) + _data = xp.roll(x._data, shift=shift, axis=axis) + return x._new(data=_data) + + +def squeeze(x: NamedArray[Any, _DType], /, axis: _AxisLike) -> NamedArray[Any, _DType]: + """ + Removes singleton dimensions (axes) from x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y", "z"), np.reshape(np.arange(1 * 2 * 3), (1, 2, 3))) + >>> squeeze(x, axis=0) + Size: 48B + array([[0, 1, 2], + [3, 4, 5]]) + """ + xp = _get_data_namespace(x) + _data = xp.squeeze(x._data, axis=axis) + _dims = _squeeze_dims(x.dims, x.shape, axis) + return x._new(_dims, _data) + + +def stack( + arrays: tuple[NamedArray[Any, Any], ...] | list[NamedArray[Any, Any]], + /, + *, + axis: _Axis = 0, +) -> NamedArray[Any, Any]: + x = arrays[0] + xp = _get_data_namespace(x) + _arrays = tuple(a._data for a in arrays) + _data = xp.stack(_arrays, axis=axis) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def tile( + x: NamedArray[Any, _DType], repetitions: tuple[int, ...], / +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _data = xp.tile(x._data, repetitions) + _dims = _infer_dims(_data.shape) # TODO: Fix dims + return x._new(_dims, _data) + + +def unstack( + x: NamedArray[Any, Any], /, *, axis: _Axis = 0 +) -> tuple[NamedArray[Any, Any], ...]: + """ + Splits an array into a sequence of arrays along the given axis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y", "z"), np.reshape(np.arange(1 * 2 * 3), (1, 2, 3))) + >>> x_y0, x_y1 = unstack(x, axis=1) + >>> x_y0 + Size: 24B + array([[0, 1, 2]]) + >>> x_y1 + Size: 24B + array([[3, 4, 5]]) + """ + xp = _get_data_namespace(x) + _datas = xp.unstack(x._data, axis=axis) + key = [slice(None)] * x.ndim + key[axis] = 0 + _dims = _dims_from_tuple_indexing(x.dims, tuple(key)) + out: tuple[NamedArray[Any, Any], ...] = () + for _data in _datas: + out += (x._new(_dims, _data),) + return out + + +# %% Automatic broadcasting +_OPTIONS = {} +_OPTIONS["arithmetic_broadcast"] = True + + +def _set_dims( + x: NamedArray[Any, Any], dim: _Dims, shape: _Shape | None +) -> NamedArray[Any, Any]: + """ + Return a new array with given set of dimensions. + This method might be used to attach new dimension(s) to array. + + When possible, this operation does not copy this variable's data. + + Parameters + ---------- + dim : + Dimensions to include on the new variable. + shape : + Shape of the dimensions. If None, new dimensions are inserted with length 1. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.asarray([1, 2, 3])) + >>> _set_dims(x, ("y", "x"), None) + Size: 24B + array([[1, 2, 3]]) + >>> _set_dims(x, ("x", "y"), None) + Size: 24B + array([[1], + [2], + [3]]) + + With shape: + + >>> _set_dims(x, ("y", "x"), (2, 3)) + Size: 48B + array([[1, 2, 3], + [1, 2, 3]]) + + No operation + + >>> _set_dims(x, ("x",), None) + Size: 24B + array([1, 2, 3]) + + Unordered dims + + >>> x = NamedArray(("y", "x"), np.zeros((2, 3))) + >>> _set_dims(x, ("x", "y"), None) + Size: 48B + array([[0., 0.], + [0., 0.], + [0., 0.]]) + + Errors + + >>> x = NamedArray(("x",), np.asarray([1, 2, 3])) + >>> _set_dims(x, (), None) + Traceback (most recent call last): + ... + ValueError: new dimensions () must be a superset of existing dimensions ('x',) + """ + if x.dims == dim: + # No operation. Don't use broadcast_to unless necessary so the result + # remains writeable as long as possible: + return x + + missing_dims = set(x.dims) - set(dim) + if missing_dims: + raise ValueError( + f"new dimensions {dim!r} must be a superset of " + f"existing dimensions {x.dims!r}" + ) + + extra_dims = tuple(d for d in dim if d not in x.dims) + + if shape is not None: + # Add dimensions, with same size as shape: + dims_map = dict(zip(dim, shape, strict=False)) + expanded_dims = extra_dims + x.dims + tmp_shape = tuple(dims_map[d] for d in expanded_dims) + return permute_dims(broadcast_to(x, tmp_shape, dims=expanded_dims), dims=dim) + else: + # Add dimensions, with size 1 only: + out = x + for d in extra_dims: + out = expand_dims(out, dim=d) + return permute_dims(out, dims=dim) + + +def _broadcast_arrays(*arrays: NamedArray[Any, Any]) -> NamedArray[Any, Any]: + """ + TODO: Can this become xp.broadcast_arrays? + + Given any number of variables, return variables with matching dimensions + and broadcast data. + + The data on the returned variables may be a view of the data on the + corresponding original arrays but dimensions will be reordered and + inserted so that both broadcast arrays have the same dimensions. The new + dimensions are sorted in order of appearance in the first variable's + dimensions followed by the second variable's dimensions. + """ + dims, shape = _broadcast_dims(*arrays) + return tuple(_set_dims(var, dims, shape) for var in arrays) + + +def _broadcast_arrays_with_minimal_size( + *arrays: NamedArray[Any, Any], +) -> NamedArray[Any, Any]: + """ + Given any number of variables, return variables with matching dimensions. + + Unlike the result of broadcast_variables(), variables with missing dimensions + will have them added with size 1 instead of the size of the broadcast dimension. + """ + dims, _ = _broadcast_dims(*arrays) + return tuple(_set_dims(var, dims, None) for var in arrays) + + +def _arithmetic_broadcast( + x1: NamedArray[Any, Any], x2: NamedArray[Any, Any] +) -> NamedArray[Any, Any]: + """ + Fu + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.asarray([1, 2, 3])) + >>> y = NamedArray(("y",), np.asarray([4, 5])) + >>> x_new, y_new = _arithmetic_broadcast(x, y) + >>> x_new.dims, x_new.shape, y_new.dims, y_new.shape + (('x', 'y'), (3, 1), ('x', 'y'), (1, 2)) + """ + if not _OPTIONS["arithmetic_broadcast"]: + if x1.dims != x2.dims: + raise ValueError( + "Broadcasting is necessary but automatic broadcasting is disabled " + "via global option `'arithmetic_broadcast'`. " + "Use `xr.set_options(arithmetic_broadcast=True)` to enable " + "automatic broadcasting." + ) + + return _broadcast_arrays_with_minimal_size(x1, x2) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_array_api/_searching_functions.py b/xarray/namedarray/_array_api/_searching_functions.py new file mode 100644 index 00000000000..755ee16b588 --- /dev/null +++ b/xarray/namedarray/_array_api/_searching_functions.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from xarray.namedarray._array_api._utils import ( + _dim_to_optional_axis, + _get_data_namespace, + _infer_dims, + _reduce_dims, +) +from xarray.namedarray._typing import ( + Default, + _arrayapi, + _default, + _Dim, +) +from xarray.namedarray.core import ( + NamedArray, +) + +if TYPE_CHECKING: + from typing import Literal + + +def argmax( + x: NamedArray[Any, Any], + /, + *, + dim: _Dim | Default = _default, + keepdims: bool = False, + axis: int | None = None, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _axis = _dim_to_optional_axis(x, dim, axis) + _data = xp.argmax(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def argmin( + x: NamedArray[Any, Any], + /, + *, + dim: _Dim | Default = _default, + keepdims: bool = False, + axis: int | None = None, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x) + _axis = _dim_to_optional_axis(x, dim, axis) + _data = xp.argmin(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def nonzero(x: NamedArray[Any, Any], /) -> tuple[NamedArray[Any, Any], ...]: + xp = _get_data_namespace(x) + _datas: tuple[_arrayapi[Any, Any], ...] = xp.nonzero(x._data) + # TODO: Verify that dims and axis matches here: + return tuple( + x._new((dim,), data) for dim, data in zip(x.dims, _datas, strict=False) + ) + + +def searchsorted( + x1: NamedArray[Any, Any], + x2: NamedArray[Any, Any], + /, + *, + side: Literal["left", "right"] = "left", + sorter: NamedArray[Any, Any] | None = None, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + _data = xp.searchsorted(x1._data, x2._data, side=side, sorter=sorter) + # TODO: Check dims, probably can do it smarter: + _dims = _infer_dims(_data.shape) + return NamedArray(_dims, _data) + + +def where( + condition: NamedArray[Any, Any], + x1: NamedArray[Any, Any], + x2: NamedArray[Any, Any], + /, +) -> NamedArray[Any, Any]: + xp = _get_data_namespace(x1) + _data = xp.where(condition._data, x1._data, x2._data) + # TODO: Wrong, _dims should be either of the arguments. How to choose? + _dims = _infer_dims(_data.shape) + return NamedArray(_dims, _data) diff --git a/xarray/namedarray/_array_api/_set_functions.py b/xarray/namedarray/_array_api/_set_functions.py new file mode 100644 index 00000000000..37590fb409e --- /dev/null +++ b/xarray/namedarray/_array_api/_set_functions.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from typing import Any, NamedTuple + +from xarray.namedarray._array_api._utils import ( + _atleast1d_dims, + _flatten_dims, + _get_data_namespace, +) +from xarray.namedarray.core import NamedArray + + +class UniqueAllResult(NamedTuple): + values: NamedArray[Any, Any] + indices: NamedArray[Any, Any] + inverse_indices: NamedArray[Any, Any] + counts: NamedArray[Any, Any] + + +class UniqueCountsResult(NamedTuple): + values: NamedArray[Any, Any] + counts: NamedArray[Any, Any] + + +class UniqueInverseResult(NamedTuple): + values: NamedArray[Any, Any] + inverse_indices: NamedArray[Any, Any] + + +def unique_all(x: NamedArray[Any, Any], /) -> UniqueAllResult: + xp = _get_data_namespace(x) + values, indices, inverse_indices, counts = xp.unique_all(x._data) + _dims = _flatten_dims(_atleast1d_dims(x.dims)) + return UniqueAllResult( + NamedArray(_dims, values), + NamedArray(_dims, indices), + NamedArray(_flatten_dims(x.dims), inverse_indices), + NamedArray(_dims, counts), + ) + + +def unique_counts(x: NamedArray[Any, Any], /) -> UniqueCountsResult: + """ + Returns the unique elements of an input array x and the corresponding + counts for each unique element in x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.array([0, 1, 2, 2], dtype=int)) + >>> x_unique = unique_counts(x) + >>> x_unique.values + Size: 24B + array([0, 1, 2]) + >>> x_unique.counts + Size: 24B + array([1, 1, 2]) + + >>> x = NamedArray(("x", "y"), np.array([0, 1, 2, 2], dtype=int).reshape((2, 2))) + >>> x_unique = unique_counts(x) + >>> x_unique.values + Size: 24B + array([0, 1, 2]) + >>> x_unique.counts + Size: 24B + array([1, 1, 2]) + """ + xp = _get_data_namespace(x) + values, counts = xp.unique_counts(x._data) + _dims = _flatten_dims(_atleast1d_dims(x.dims)) + return UniqueCountsResult( + NamedArray(_dims, values), + NamedArray(_dims, counts), + ) + + +def unique_inverse(x: NamedArray[Any, Any], /) -> UniqueInverseResult: + """ + Returns the unique elements of an input array x and the indices + from the set of unique elements that reconstruct x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.array([0, 1, 2, 2], dtype=int)) + >>> x_unique = unique_inverse(x) + >>> x_unique.values + Size: 24B + array([0, 1, 2]) + >>> x_unique.inverse_indices + Size: 32B + array([0, 1, 2, 2]) + + >>> x = NamedArray(("x", "y"), np.array([0, 1, 2, 2], dtype=int).reshape((2, 2))) + >>> x_unique = unique_inverse(x) + >>> x_unique.values + Size: 24B + array([0, 1, 2]) + >>> x_unique.inverse_indices + Size: 32B + array([[0, 1], + [2, 2]]) + """ + xp = _get_data_namespace(x) + values, inverse_indices = xp.unique_inverse(x._data) + return UniqueInverseResult( + NamedArray(_flatten_dims(_atleast1d_dims(x.dims)), values), + NamedArray(x.dims, inverse_indices), + ) + + +def unique_values(x: NamedArray[Any, Any], /) -> NamedArray[Any, Any]: + """ + Returns the unique elements of an input array x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.array([0, 1, 2, 2], dtype=int)) + >>> unique_values(x) + Size: 24B + array([0, 1, 2]) + >>> x = NamedArray(("x", "y"), np.array([0, 1, 2, 2], dtype=int).reshape((2, 2))) + >>> unique_values(x) + Size: 24B + array([0, 1, 2]) + + # Scalars becomes 1-dimensional + + >>> x = NamedArray((), np.array(0, dtype=int)) + >>> unique_values(x) + Size: 8B + array([0]) + """ + xp = _get_data_namespace(x) + _data = xp.unique_values(x._data) + _dims = _flatten_dims(_atleast1d_dims(x.dims)) + return x._new(_dims, _data) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_array_api/_sorting_functions.py b/xarray/namedarray/_array_api/_sorting_functions.py new file mode 100644 index 00000000000..12554597e3c --- /dev/null +++ b/xarray/namedarray/_array_api/_sorting_functions.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from typing import Any + +from xarray.namedarray._array_api._utils import ( + _dim_to_axis, + _get_data_namespace, +) +from xarray.namedarray._typing import ( + Default, + _default, + _Dim, + _DType, + _ShapeType, +) +from xarray.namedarray.core import NamedArray + + +def argsort( + x: NamedArray[_ShapeType, Any], + /, + *, + dim: _Dim | Default = _default, + descending: bool = False, + stable: bool = True, + axis: int = -1, +) -> NamedArray[_ShapeType, Any]: + xp = _get_data_namespace(x) + _axis = _dim_to_axis(x, dim, axis) + + # TODO: As NumPy currently has no native descending sort, we imitate it here: + if not descending: + _data = xp.argsort(x._data, axis=_axis, stable=stable) + else: + _data = xp.flip( + xp.argsort(xp.flip(x._data, axis=_axis), stable=stable, axis=_axis), + axis=_axis, + ) + # Rely on flip()/argsort() to validate axis + normalised_axis = _axis if _axis >= 0 else x.ndim + _axis + max_i = x.shape[normalised_axis] - 1 + _data = max_i - _data + return x._new(data=_data) + + +def sort( + x: NamedArray[_ShapeType, _DType], + /, + *, + dim: _Dim | Default = _default, + descending: bool = False, + stable: bool = True, + axis: int = -1, +) -> NamedArray[_ShapeType, _DType]: + xp = _get_data_namespace(x) + _axis = _dim_to_axis(x, dim, axis) + + _data = xp.sort(x._data, axis=_axis, stable=stable) + # TODO: As NumPy currently has no native descending sort, we imitate it here: + if descending: + _data = xp.flip(_data, axis=_axis) + return x._new(data=_data) diff --git a/xarray/namedarray/_array_api/_statistical_functions.py b/xarray/namedarray/_array_api/_statistical_functions.py new file mode 100644 index 00000000000..294cfa6012c --- /dev/null +++ b/xarray/namedarray/_array_api/_statistical_functions.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +from typing import Any + +from xarray.namedarray._array_api._utils import ( + _dims_to_axis, + _get_data_namespace, + _reduce_dims, +) +from xarray.namedarray._typing import ( + Default, + _AxisLike, + _default, + _Dim, + _Dims, + _DType, + _ShapeType, +) +from xarray.namedarray.core import NamedArray + + +def cumulative_sum( + x: NamedArray[_ShapeType, _DType], + /, + *, + dim: _Dim | Default = _default, + dtype: _DType | None = None, + include_initial: bool = False, + axis: int | None = None, +) -> NamedArray[_ShapeType, _DType]: + xp = _get_data_namespace(x) + + a = _dims_to_axis(x, dim, axis) + _axis_none: int | None + if a is None: + _axis_none = a + else: + _axis_none = a[0] + + # TODO: The standard is not clear about what should happen when x.ndim == 0. + _axis: int + if _axis_none is None: + if x.ndim > 1: + raise ValueError( + "axis must be specified in cumulative_sum for more than one dimension" + ) + _axis = 0 + else: + _axis = _axis_none + + try: + _data = xp.cumulative_sum( + x._data, axis=_axis, dtype=dtype, include_initial=include_initial + ) + except AttributeError: + # Use np.cumsum until new name is introduced: + # np.cumsum does not support include_initial + if include_initial: + if _axis < 0: + _axis += x.ndim + d = xp.concat( + [ + xp.zeros( + x.shape[:_axis] + (1,) + x.shape[_axis + 1 :], dtype=x.dtype + ), + x._data, + ], + axis=_axis, + ) + else: + d = x._data + _data = xp.cumsum(d, axis=_axis, dtype=dtype) + return x._new(dims=x.dims, data=_data) + + +def max( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.max(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def mean( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + """ + Calculates the arithmetic mean of the input array x. + + Parameters + ---------- + x : + Should have a real-valued floating-point data type. + dims : + Dim or dims along which arithmetic means must be computed. By default, + the mean must be computed over the entire array. If a tuple of hashables, + arithmetic means must be computed over multiple axes. + Default: None. + keepdims : + if True, the reduced axes (dimensions) must be included in the result + as singleton dimensions, and, accordingly, the result must be compatible + with the input array (see Broadcasting). Otherwise, if False, the + reduced axes (dimensions) must not be included in the result. + Default: False. + axis : + Axis or axes along which arithmetic means must be computed. By default, + the mean must be computed over the entire array. If a tuple of integers, + arithmetic means must be computed over multiple axes. + Default: None. + + Returns + ------- + out : + If the arithmetic mean was computed over the entire array, + a zero-dimensional array containing the arithmetic mean; otherwise, + a non-zero-dimensional array containing the arithmetic means. + The returned array must have the same data type as x. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y"), np.asarray([[1.0, 2.0], [3.0, 4.0]])) + >>> mean(x) + Size: 8B + np.float64(2.5) + >>> mean(x, dims=("x",)) + Size: 16B + array([2., 3.]) + + Using keepdims: + + >>> mean(x, dims=("x",), keepdims=True) + Size: 16B + array([[2., 3.]]) + >>> mean(x, dims=("y",), keepdims=True) + Size: 16B + array([[1.5], + [3.5]]) + """ + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.mean(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def min( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.min(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def prod( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + dtype: _DType | None = None, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.prod(x._data, axis=_axis, dtype=dtype, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def std( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + correction: int | float = 0.0, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.std(x._data, axis=_axis, correction=correction, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def sum( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + dtype: _DType | None = None, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.sum(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +def var( + x: NamedArray[Any, _DType], + /, + *, + dims: _Dims | Default = _default, + correction: int | float = 0.0, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, _DType]: + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.var(x._data, axis=_axis, correction=correction, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return x._new(dims=_dims, data=_data) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_array_api/_utility_functions.py b/xarray/namedarray/_array_api/_utility_functions.py new file mode 100644 index 00000000000..07f1a4fdd60 --- /dev/null +++ b/xarray/namedarray/_array_api/_utility_functions.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from typing import Any + +from xarray.namedarray._array_api._utils import ( + _dim_to_axis, + _dims_to_axis, + _get_data_namespace, + _reduce_dims, +) +from xarray.namedarray._typing import ( + Default, + _Axis, + _AxisLike, + _default, + _Dim, + _DimsLike2, +) +from xarray.namedarray.core import NamedArray + + +def all( + x: NamedArray[Any, Any], + /, + *, + dims: _DimsLike2 | Default = _default, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, Any]: + """ + Tests whether all input array elements evaluate to True along a specified axis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y"), np.array([[True, False], [True, True]])) + >>> all(x) + Size: 1B + np.False_ + + >>> all(x, axis=0) + Size: 2B + array([ True, False]) + >>> all(x, dims="x") + Size: 2B + array([ True, False]) + + >>> x = NamedArray(("x",), np.array([-1, 4, 5])) + >>> all(x) + Size: 1B + np.True_ + """ + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.all(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return NamedArray(_dims, _data) + + +def any( + x: NamedArray[Any, Any], + /, + *, + dims: _DimsLike2 | Default = _default, + keepdims: bool = False, + axis: _AxisLike | None = None, +) -> NamedArray[Any, Any]: + """ + Tests whether any input array element evaluates to True along a specified axis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y"), np.array([[True, False], [True, True]])) + >>> any(x) + Size: 1B + np.True_ + + >>> x = NamedArray(("x", "y"), np.array([[True, False], [False, False]])) + >>> any(x, axis=0) + Size: 2B + array([ True, False]) + >>> any(x, dims="x") + Size: 2B + array([ True, False]) + + >>> x = NamedArray(("x",), np.array([-1, 4, 5])) + >>> any(x) + Size: 1B + np.True_ + """ + xp = _get_data_namespace(x) + _axis = _dims_to_axis(x, dims, axis) + _data = xp.any(x._data, axis=_axis, keepdims=keepdims) + _dims = _reduce_dims(x.dims, axis=_axis, keepdims=keepdims) + return NamedArray(_dims, _data) + + +def diff( + x: NamedArray[Any, Any], + /, + *, + axis: _Axis = -1, + n: int = 1, + prepend: NamedArray[Any, Any] | None = None, + append: NamedArray[Any, Any] | None = None, + dims: _Dim | Default = _default, +) -> NamedArray[Any, Any]: + """ + Calculates the n-th discrete forward difference along a specified axis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x",), np.array([1, 2, 4, 7, 0])) + >>> diff(x) + Size: 32B + array([ 1, 2, 3, -7]) + >>> diff(x, n=2) + Size: 24B + array([ 1, 1, -10]) + + >>> x = NamedArray(("x", "y"), np.array([[1, 3, 6, 10], [0, 5, 6, 8]])) + >>> diff(x) + Size: 48B + array([[2, 3, 4], + [5, 1, 2]]) + >>> diff(x, axis=0) + Size: 32B + array([[-1, 2, 0, -2]]) + >>> diff(x, dims="x") + Size: 32B + array([[-1, 2, 0, -2]]) + """ + xp = _get_data_namespace(x) + _axis = _dim_to_axis(x, dims, axis, axis_default=-1) + _data = xp.diff(x._data, axis=_axis, n=n) # TODO: , prepend=prepend, append=append + return NamedArray(x.dims, _data) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_array_api/_utils.py b/xarray/namedarray/_array_api/_utils.py new file mode 100644 index 00000000000..48521e756cf --- /dev/null +++ b/xarray/namedarray/_array_api/_utils.py @@ -0,0 +1,741 @@ +from __future__ import annotations + +import math +from collections.abc import Callable, Iterable, Iterator +from itertools import zip_longest +from types import ModuleType +from typing import Any, TypeGuard, cast + +from xarray.namedarray._typing import ( + _T, + Default, + _arrayapi, + _Axes, + _Axis, + _AxisLike, + _default, + _Dim, + _Dims, + _DimsLike2, + _dtype, + _IndexKeys, + _IndexKeysNoEllipsis, + _Shape, +) +from xarray.namedarray.core import NamedArray + + +def _maybe_default_namespace(xp: ModuleType | None = None) -> ModuleType: + if xp is None: + # import array_api_strict as xpd + + # import array_api_compat.numpy as xpd + + import numpy as xpd + + return xpd + else: + return xp + + +def _get_namespace(x: Any) -> ModuleType: + if isinstance(x, _arrayapi): + return x.__array_namespace__() + + return _maybe_default_namespace() + + +def _get_data_namespace(x: NamedArray[Any, Any]) -> ModuleType: + return _get_namespace(x._data) + + +def _get_namespace_dtype(dtype: _dtype[Any] | None = None) -> ModuleType: + """ + Attempt to get the namespace from dtype. + + Examples + -------- + >>> import numpy as np + >>> _get_namespace_dtype(None) is np + True + >>> _get_namespace_dtype(np.int64) is np + True + """ + if dtype is None: + return _maybe_default_namespace() + + try: + xp = __import__(dtype.__module__) + except AttributeError: + # TODO: Fix this. + # FAILED array_api_tests/test_searching_functions.py::test_searchsorted - AttributeError: 'numpy.dtypes.Float64DType' object has no attribute '__module__'. Did you mean: '__mul__'? + # Falsifying example: test_searchsorted( + # data=data(...), + # ) + return _maybe_default_namespace() + return xp + + +def _is_single_dim(dims: _DimsLike2) -> TypeGuard[_Dim]: + # TODO: https://peps.python.org/pep-0742/ + return isinstance(dims, str) or not isinstance(dims, Iterable) + + +def _normalize_dimensions(dims: _DimsLike2) -> _Dims: + """ + Normalize dimensions. + + Examples + -------- + >>> _normalize_dimensions(None) + (None,) + >>> _normalize_dimensions(1) + (1,) + >>> _normalize_dimensions("2") + ('2',) + >>> _normalize_dimensions(("time",)) + ('time',) + >>> _normalize_dimensions(["time"]) + ('time',) + >>> _normalize_dimensions([("time", "x", "y")]) + (('time', 'x', 'y'),) + """ + if _is_single_dim(dims): + return (dims,) + else: + return tuple(cast(_Dims, dims)) + + +def _infer_dims( + shape: _Shape, + dims: _DimsLike2 | Default = _default, +) -> _Dims: + """ + Create default dim names if no dims were supplied. + + Examples + -------- + >>> _infer_dims(()) + () + >>> _infer_dims((1,)) + ('dim_0',) + >>> _infer_dims((3, 1)) + ('dim_1', 'dim_0') + + >>> _infer_dims((), ()) + () + >>> _infer_dims((1,), "x") + ('x',) + >>> _infer_dims((1,), None) + (None,) + >>> _infer_dims((1,), ("x",)) + ('x',) + >>> _infer_dims((1, 3), ("x",)) + ('dim_0', 'x') + >>> _infer_dims((1, 1, 3), ("x",)) + ('dim_1', 'dim_0', 'x') + """ + if isinstance(dims, Default): + ndim = len(shape) + return tuple(f"dim_{ndim - 1 - n}" for n in range(ndim)) + + _dims = _normalize_dimensions(dims) + diff = len(shape) - len(_dims) + if diff > 0: + # TODO: Leads to ('dim_0', 'x'), should it be ('dim_1', 'x')? + return _infer_dims(shape[:diff], _default) + _dims + else: + return _dims + + +def _normalize_axis_index(axis: int, ndim: int) -> int: + """ + Normalize axis index to positive values. + + Parameters + ---------- + axis : int + The un-normalized index of the axis. Can be negative + ndim : int + The number of dimensions of the array that `axis` should be normalized + against + + Returns + ------- + normalized_axis : int + The normalized axis index, such that `0 <= normalized_axis < ndim` + + + Examples + -------- + >>> _normalize_axis_index(0, ndim=3) + 0 + >>> _normalize_axis_index(1, ndim=3) + 1 + >>> _normalize_axis_index(2, ndim=3) + 2 + >>> _normalize_axis_index(-1, ndim=3) + 2 + >>> _normalize_axis_index(-2, ndim=3) + 1 + >>> _normalize_axis_index(-3, ndim=3) + 0 + + Errors + + >>> _normalize_axis_index(3, ndim=3) + Traceback (most recent call last): + ... + ValueError: axis 3 is out of bounds for array of dimension 3 + >>> _normalize_axis_index(-4, ndim=3) + Traceback (most recent call last): + ... + ValueError: axis -4 is out of bounds for array of dimension 3 + """ + + if -ndim > axis or axis >= ndim: + raise ValueError(f"axis {axis} is out of bounds for array of dimension {ndim}") + + return axis % ndim + + +def _normalize_axis_tuple( + axis: _AxisLike, + ndim: int, + argname: str | None = None, + allow_duplicate: bool = False, +) -> _Axes: + """ + Normalizes an axis argument into a tuple of non-negative integer axes. + + This handles shorthands such as ``1`` and converts them to ``(1,)``, + as well as performing the handling of negative indices covered by + `normalize_axis_index`. + + By default, this forbids axes from being specified multiple times. + + + Parameters + ---------- + axis : int, iterable of int + The un-normalized index or indices of the axis. + ndim : int + The number of dimensions of the array that `axis` should be normalized + against. + argname : str, optional + A prefix to put before the error message, typically the name of the + argument. + allow_duplicate : bool, optional + If False, the default, disallow an axis from being specified twice. + + Returns + ------- + normalized_axes : tuple of int + The normalized axis index, such that `0 <= normalized_axis < ndim` + + Raises + ------ + AxisError + If any axis provided is out of range + ValueError + If an axis is repeated + """ + if isinstance(axis, tuple): + _axis = axis + else: + _axis = (axis,) + + # Going via an iterator directly is slower than via list comprehension. + _axis = tuple([_normalize_axis_index(ax, ndim) for ax in _axis]) + if not allow_duplicate and len(set(_axis)) != len(_axis): + if argname: + raise ValueError(f"repeated axis in `{argname}` argument") + else: + raise ValueError("repeated axis") + return _axis + + +def _assert_either_dim_or_axis( + dims: _DimsLike2 | Default, axis: _AxisLike | None +) -> None: + if dims is not _default and axis is not None: + raise ValueError("cannot supply both 'axis' and 'dim(s)' arguments") + + +# @overload +# def _dims_to_axis(x: NamedArray[Any, Any], dims: Default, axis: None) -> None: ... +# @overload +# def _dims_to_axis(x: NamedArray[Any, Any], dims: _DimsLike2, axis: None) -> _Axes: ... +# @overload +# def _dims_to_axis(x: NamedArray[Any, Any], dims: Default, axis: _AxisLike) -> _Axes: ... +def _dims_to_axis( + x: NamedArray[Any, Any], dims: _DimsLike2 | Default, axis: _AxisLike | None +) -> _Axes | None: + """ + Convert dims to axis indices. + + Examples + -------- + + Convert to dims to axis values + >>> import numpy as np + >>> x = NamedArray(("x", "y", "z"), np.zeros((1, 2, 3))) + >>> _dims_to_axis(x, ("y", "x"), None) + (1, 0) + >>> _dims_to_axis(x, ("y",), None) + (1,) + >>> _dims_to_axis(x, _default, 0) + (0,) + >>> type(_dims_to_axis(x, _default, None)) + + + Normalizes negative integers + + >>> _dims_to_axis(x, _default, -1) + (2,) + >>> _dims_to_axis(x, _default, (-2, -1)) + (1, 2) + + Using Hashable dims + + >>> x = NamedArray(("x", None), np.zeros((1, 2))) + >>> _dims_to_axis(x, None, None) + (1,) + + Defining both dims and axis raises an error + + >>> _dims_to_axis(x, "x", 1) + Traceback (most recent call last): + ... + ValueError: cannot supply both 'axis' and 'dim(s)' arguments + """ + _assert_either_dim_or_axis(dims, axis) + if not isinstance(dims, Default): + _dims = _normalize_dimensions(dims) + + axis = () + for dim in _dims: + try: + axis += (x.dims.index(dim),) + except ValueError as err: + raise ValueError( + f"{dim!r} not found in array dimensions {x.dims!r}" + ) from err + return axis + + if axis is None: + return axis + + return _normalize_axis_tuple(axis, x.ndim) + + +def _dim_to_optional_axis( + x: NamedArray[Any, Any], dim: _Dim | Default, axis: int | None +) -> int | None: + a = _dims_to_axis(x, dim, axis) + if a is None: + return a + + return a[0] + + +def _dim_to_axis( + x: NamedArray[Any, Any], + dim: _Dim | Default, + axis: int, + *, + axis_default: int | None = None, +) -> int: + """ + Convert dim to axis. + + Examples + -------- + >>> import numpy as np + >>> x = NamedArray(("x", "y", "z"), np.zeros((1, 2, 3))) + >>> _dim_to_axis(x, _default, 0) + 0 + >>> _dim_to_axis(x, _default, -1) + 2 + >>> _dim_to_axis(x, "x", -1) + 0 + >>> _dim_to_axis(x, "z", 0) + 2 + + Add a default value + + >>> _dim_to_axis(x, _default, 0, axis_default=-1) + 0 + >>> _dim_to_axis(x, "z", -1, axis_default=-1) + 2 + + Defining both dims and axis value different to the default raises an error + + >>> _dim_to_axis(x, "x", 0, axis_default=-1) + Traceback (most recent call last): + ... + ValueError: cannot supply both 'axis' and 'dim(s)' arguments + """ + if axis_default is not None: + a = None if axis is axis_default else axis + _assert_either_dim_or_axis(dim, a) + + _dim: _Dim = x.dims[axis] if isinstance(dim, Default) else dim + _axis = _dim_to_optional_axis(x, _dim, None) + assert _axis is not None # Not supposed to happen. + return _axis + + +def _new_unique_dim_name(dims: _Dims, i: int | None = None) -> _Dim: + """ + Get a new unique dimension name. + + Examples + -------- + >>> _new_unique_dim_name(()) + 'dim_0' + >>> _new_unique_dim_name(("dim_0",)) + 'dim_1' + >>> _new_unique_dim_name(("dim_1", "dim_0")) + 'dim_2' + >>> _new_unique_dim_name(("dim_0", "dim_2")) + 'dim_3' + >>> _new_unique_dim_name(("dim_3", "dim_2")) + 'dim_4' + """ + i = len(dims) if i is None else i + _dim: _Dim = f"dim_{i}" + return _new_unique_dim_name(dims, i=i + 1) if _dim in dims else _dim + + +def _insert_dim(dims: _Dims, dim: _Dim | Default, axis: _Axis) -> _Dims: + if isinstance(dim, Default): + _dim: _Dim = _new_unique_dim_name(dims) + else: + _dim = dim + + d = list(dims) + d.insert(axis, _dim) + return tuple(d) + + +def _filter_next_false( + predicate: Callable[..., bool], iterable: Iterable[_T] +) -> Iterator[_T]: + """ + Make an iterator that filters elements from the iterable returning only those + for which the predicate returns a false value for the second time. + + Variant on itertools.filterfalse but doesn't filter until the 2 second False. + + Examples + -------- + >>> tuple(_filter_next_false(lambda x: x is not None, (1, None, 3, None, 4))) + (1, None, 3, 4) + """ + predicate_has_been_false = False + for x in iterable: + if not predicate(x): + if predicate_has_been_false: + continue + predicate_has_been_false = True + yield x + + +def _replace_ellipsis(key: _IndexKeys, ndim: int) -> _IndexKeysNoEllipsis: + """ + Replace ... with slices, :, : ,: + + >>> _replace_ellipsis((3, Ellipsis, 2), 4) + (3, slice(None, None, None), slice(None, None, None), 2) + + >>> _replace_ellipsis((Ellipsis, None), 2) + (slice(None, None, None), slice(None, None, None), None) + >>> _replace_ellipsis((Ellipsis, None, Ellipsis), 2) + (slice(None, None, None), slice(None, None, None), None) + """ + # https://github.com/dask/dask/blob/569abf8e8048cbfb1d750900468dda0de7c56358/dask/array/slicing.py#L701 + key = tuple(_filter_next_false(lambda x: x is not Ellipsis, key)) + expanded_dims = sum(i is None for i in key) + extra_dimensions = ndim - (len(key) - expanded_dims - 1) + replaced_slices = (slice(None, None, None),) * extra_dimensions + + out: _IndexKeysNoEllipsis = () + for k in key: + if k is Ellipsis: + out += replaced_slices + else: + out += (k,) + return out + + +def _dims_from_tuple_indexing(dims: _Dims, key: _IndexKeys) -> _Dims: + """ + Get the expected dims when using tuples in __getitem__. + + Examples + -------- + >>> _dims_from_tuple_indexing(("x", "y", "z"), ()) + ('x', 'y', 'z') + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0,)) + ('y', 'z') + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0, 0)) + ('z',) + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0, 0, 0)) + () + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0, 0, 0, ...)) + () + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0, ...)) + ('y', 'z') + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0, ..., 0)) + ('y',) + >>> _dims_from_tuple_indexing(("x", "y", "z"), (0, slice(0))) + ('y', 'z') + >>> _dims_from_tuple_indexing(("x", "y"), (None,)) + ('dim_2', 'x', 'y') + >>> _dims_from_tuple_indexing(("x", "y"), (0, None, None, 0)) + ('dim_1', 'dim_2') + >>> _dims_from_tuple_indexing(("x",), (..., 0)) + () + """ + key_no_ellipsis = _replace_ellipsis(key, len(dims)) + _dims = list(dims) + j = 0 + for k in key_no_ellipsis: + if k is None: + # None adds 1 dimension: + _dims.insert(j, _new_unique_dim_name(tuple(_dims))) + j += 1 + elif isinstance(k, int): + # Integer removes 1 dimension: + _dims.pop(j) + elif isinstance(k, slice): + # Slice retains the dimension. + j += 1 + + return tuple(_dims) + + +def _atleast1d_dims(dims: _Dims) -> _Dims: + """ + Set dims atleast 1-dimensional. + + Examples + -------- + >>> _atleast1d_dims(()) + ('dim_0',) + >>> _atleast1d_dims(("x",)) + ('x',) + >>> _atleast1d_dims(("x", "y")) + ('x', 'y') + """ + return (_new_unique_dim_name(dims),) if len(dims) < 1 else dims + + +def _flatten_dims(dims: _Dims) -> _Dims: + """ + Flatten multidimensional dims to 1-dimensional. + + Examples + -------- + >>> _flatten_dims(()) + () + >>> _flatten_dims(("x",)) + ('x',) + >>> _flatten_dims(("x", "y")) + (('x', 'y'),) + """ + return (dims,) if len(dims) > 1 else dims + + +def _move_dims(dims: _Dims, source: _AxisLike, destination: _AxisLike) -> _Dims: + """ + Move dims position in source to the destination position. + + Examples + -------- + >>> _move_dims(("x", "y", "z"), 0, -1) + ('y', 'z', 'x') + >>> _move_dims(("x", "y", "z"), -1, 0) + ('z', 'x', 'y') + >>> _move_dims(("x", "y", "z"), (0, 1), (-1, -2)) + ('z', 'y', 'x') + >>> _move_dims(("x", "y", "z"), (0, 1, 2), (-1, -2, -3)) + ('z', 'y', 'x') + >>> _move_dims(("x", "y", "z"), 0, 1) + ('y', 'x', 'z') + """ + _ndim = len(dims) + _source = _normalize_axis_tuple(source, _ndim) + _destination = _normalize_axis_tuple(destination, _ndim) + order = [n for n in range(_ndim) if n not in _source] + for dest, src in sorted(zip(_destination, _source, strict=True)): + order.insert(dest, src) + + return tuple(dims[i] for i in order) + + +def _reduce_dims(dims: _Dims, *, axis: _AxisLike | None, keepdims: bool) -> _Dims: + """ + Reduce dims according to axis. + + Examples + -------- + >>> _reduce_dims(("x", "y", "z"), axis=None, keepdims=False) + () + >>> _reduce_dims(("x", "y", "z"), axis=1, keepdims=False) + ('x', 'z') + >>> _reduce_dims(("x", "y", "z"), axis=-1, keepdims=False) + ('x', 'y') + + keepdims retains the same dims + + >>> _reduce_dims(("x", "y", "z"), axis=-1, keepdims=True) + ('x', 'y', 'z') + """ + if keepdims: + return dims + + ndim = len(dims) + if axis is None: + _axis = tuple(v for v in range(ndim)) + else: + _axis = _normalize_axis_tuple(axis, ndim) + + k: _IndexKeys = (slice(None),) * ndim + key = list(k) + for v in _axis: + key[v] = 0 + + return _dims_from_tuple_indexing(dims, tuple(key)) + + +def _squeeze_dims(dims: _Dims, shape: _Shape, axis: _AxisLike) -> _Dims: + """ + Squeeze dims. + + Examples + -------- + >>> _squeeze_dims(("x", "y", "z"), (0, 2, 1), (0, 2)) + ('y',) + """ + sizes = dict(zip(dims, shape, strict=True)) + for a in _normalize_axis_tuple(axis, len(dims)): + d = dims[a] + if sizes[d] < 2: + sizes.pop(d) + + return tuple(sizes.keys()) + + +def _raise_if_any_duplicate_dimensions( + dims: _Dims, err_context: str = "This function" +) -> None: + if len(set(dims)) < len(dims): + repeated_dims = {d for d in dims if dims.count(d) > 1} + raise ValueError( + f"{err_context} cannot handle duplicate dimensions, " + f"but dimensions {repeated_dims} appear more than once on this object's dims: {dims}" + ) + + +def _isnone(shape: _Shape) -> tuple[bool, ...]: + # TODO: math.isnan should not be needed for array api, but dask still uses np.nan: + return tuple(v is None and math.isnan(v) for v in shape) + + +def _broadcast_dims(*arrays: NamedArray[Any, Any]) -> tuple[_Dims, _Shape]: + """ + Get the expected broadcasted dims. + + Examples + -------- + >>> import numpy as np + >>> a = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> _broadcast_dims(a) + (('x', 'y', 'z'), (5, 3, 4)) + + Broadcasting 0- and 1-sized dims + + >>> a = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> b = NamedArray(("x", "y", "z"), np.zeros((0, 3, 4))) + >>> _broadcast_dims(a, b) + (('x', 'y', 'z'), (5, 3, 4)) + >>> _broadcast_dims(b, a) + (('x', 'y', 'z'), (5, 3, 4)) + + >>> a = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> b = NamedArray(("x", "y", "z"), np.zeros((1, 3, 4))) + >>> _broadcast_dims(a, b) + (('x', 'y', 'z'), (5, 3, 4)) + + >>> a = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> b = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> _broadcast_dims(a, b) + (('x', 'y', 'z'), (5, 3, 4)) + + Broadcasting different dims + + >>> a = NamedArray(("x",), np.zeros((5,))) + >>> b = NamedArray(("y",), np.zeros((3,))) + >>> _broadcast_dims(a, b) + (('x', 'y'), (5, 3)) + + >>> a = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> b = NamedArray(("y", "z"), np.zeros((3, 4))) + >>> _broadcast_dims(a, b) + (('x', 'y', 'z'), (5, 3, 4)) + >>> _broadcast_dims(b, a) + (('x', 'y', 'z'), (5, 3, 4)) + + + # Errors + + >>> a = NamedArray(("x", "y", "z"), np.zeros((5, 3, 4))) + >>> b = NamedArray(("x", "y", "z"), np.zeros((2, 3, 4))) + >>> _broadcast_dims(a, b) + Traceback (most recent call last): + ... + ValueError: operands could not be broadcast together with dims = (('x', 'y', 'z'), ('x', 'y', 'z')) and shapes = ((5, 3, 4), (2, 3, 4)) + """ + DEFAULT_SIZE = -1 + BROADCASTABLE_SIZES = (0, 1) + BROADCASTABLE_SIZES_OR_DEFAULT = BROADCASTABLE_SIZES + (DEFAULT_SIZE,) + + arrays_dims = tuple(a.dims for a in arrays) + arrays_shapes = tuple(a.shape for a in arrays) + + sizes: dict[Any, Any] = {} + for dims, shape in zip( + zip_longest(*map(reversed, arrays_dims), fillvalue=_default), + zip_longest(*map(reversed, arrays_shapes), fillvalue=DEFAULT_SIZE), + strict=False, + ): + for d, s in zip(reversed(dims), reversed(shape), strict=False): + if isinstance(d, Default): + continue + + if s is None: + raise NotImplementedError("TODO: Handle None in shape, {shapes = }") + + s_prev = sizes.get(d, DEFAULT_SIZE) + if not ( + s == s_prev + or any(v in BROADCASTABLE_SIZES_OR_DEFAULT for v in (s, s_prev)) + ): + raise ValueError( + "operands could not be broadcast together with " + f"dims = {arrays_dims} and shapes = {arrays_shapes}" + ) + + sizes[d] = max(s, s_prev) + + out_dims: _Dims = tuple(reversed(sizes.keys())) + out_shape: _Shape = tuple(reversed(sizes.values())) + return out_dims, out_shape + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 95e7d7adfc3..11c970f83d8 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -2,7 +2,6 @@ import sys from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence -from enum import Enum from types import EllipsisType, ModuleType from typing import ( TYPE_CHECKING, @@ -11,6 +10,7 @@ Literal, Protocol, SupportsIndex, + TypedDict, TypeVar, Union, overload, @@ -21,35 +21,69 @@ try: if sys.version_info >= (3, 11): - from typing import TypeAlias + from typing import Never, TypeAlias else: from typing import TypeAlias + + from typing_extensions import Never except ImportError: if TYPE_CHECKING: raise else: + Never: Any = None Self: Any = None -# Singleton type, as per https://github.com/python/typing/pull/240 -class Default(Enum): - token: Final = 0 +class Default(list[Never]): + """ + Non-Hashable default value. + + A replacement value for Optional None since it is Hashable. + Same idea as https://github.com/python/typing/pull/240 + + Examples + -------- + + Runtime checks: + + >>> _default = Default() + >>> isinstance(_default, Hashable) + False + >>> _default == _default + True + >>> _default is _default + True + Typing usage: -_default = Default.token + >>> x: Hashable | Default = _default + >>> if isinstance(x, Default): + ... y: Default = x + ... else: + ... h: Hashable = x + ... + + TODO: if x is _default does not narrow typing, use isinstance check instead. + """ + + def __repr__(self) -> str: + return f"<{self.__class__.__name__}>" + + +_default: Final[Default] = Default() # https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array _T = TypeVar("_T") _T_co = TypeVar("_T_co", covariant=True) +_ScalarType = TypeVar("_ScalarType", bound=np.generic) +_ScalarType_co = TypeVar("_ScalarType_co", bound=np.generic, covariant=True) + _dtype = np.dtype _DType = TypeVar("_DType", bound=np.dtype[Any]) _DType_co = TypeVar("_DType_co", covariant=True, bound=np.dtype[Any]) # A subset of `npt.DTypeLike` that can be parametrized w.r.t. `np.generic` -_ScalarType = TypeVar("_ScalarType", bound=np.generic) -_ScalarType_co = TypeVar("_ScalarType_co", bound=np.generic, covariant=True) - # A protocol for anything with the dtype attribute @runtime_checkable @@ -58,6 +92,16 @@ class _SupportsDType(Protocol[_DType_co]): def dtype(self) -> _DType_co: ... +class _SupportsReal(Protocol[_T_co]): + @property + def real(self) -> _T_co: ... + + +class _SupportsImag(Protocol[_T_co]): + @property + def imag(self) -> _T_co: ... + + _DTypeLike = Union[ np.dtype[_ScalarType], type[_ScalarType], @@ -70,6 +114,7 @@ def dtype(self) -> _DType_co: ... _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] _ShapeType = TypeVar("_ShapeType", bound=Any) _ShapeType_co = TypeVar("_ShapeType_co", bound=Any, covariant=True) +_Shape1D = tuple[int] _Axis = int _Axes = tuple[_Axis, ...] @@ -85,28 +130,71 @@ def dtype(self) -> _DType_co: ... _Dim = Hashable _Dims = tuple[_Dim, ...] - +_DimsLike2 = Union[_Dim, _Dims] _DimsLike = Union[str, Iterable[_Dim]] # https://data-apis.org/array-api/latest/API_specification/indexing.html # TODO: np.array_api was bugged and didn't allow (None,), but should! # https://github.com/numpy/numpy/pull/25022 # https://github.com/data-apis/array-api/pull/674 -_IndexKey = Union[int, slice, EllipsisType] +_IndexKeyNoEllipsis = Union[int, slice, None] +_IndexKey = Union[_IndexKeyNoEllipsis, EllipsisType] +_IndexKeysNoEllipsis = tuple[_IndexKeyNoEllipsis, ...] _IndexKeys = tuple[_IndexKey, ...] # tuple[Union[_IndexKey, None], ...] _IndexKeyLike = Union[_IndexKey, _IndexKeys] _AttrsLike = Union[Mapping[Any, Any], None] +_ArrayLike = np.typing.ArrayLike -class _SupportsReal(Protocol[_T_co]): - @property - def real(self) -> _T_co: ... +_Device = Any -class _SupportsImag(Protocol[_T_co]): - @property - def imag(self) -> _T_co: ... +class _IInfo(Protocol): + bits: int + max: int + min: int + dtype: _dtype[Any] + + +class _FInfo(Protocol): + bits: int + eps: float + max: float + min: float + smallest_normal: float + dtype: _dtype[Any] + + +_Capabilities = TypedDict( + "_Capabilities", {"boolean indexing": bool, "data-dependent shapes": bool} +) + +_DefaultDataTypes = TypedDict( + "_DefaultDataTypes", + { + "real floating": _dtype[Any], + "complex floating": _dtype[Any], + "integral": _dtype[Any], + "indexing": _dtype[Any], + }, +) + + +class _DataTypes(TypedDict, total=False): + bool: _dtype[Any] + float32: _dtype[Any] + float64: _dtype[Any] + complex64: _dtype[Any] + complex128: _dtype[Any] + int8: _dtype[Any] + int16: _dtype[Any] + int32: _dtype[Any] + int64: _dtype[Any] + uint8: _dtype[Any] + uint16: _dtype[Any] + uint32: _dtype[Any] + uint64: _dtype[Any] @runtime_checkable @@ -210,6 +298,16 @@ def __getitem__( def __array_namespace__(self) -> ModuleType: ... + def to_device( + self, device: _Device, /, stream: None = None + ) -> _arrayapi[_ShapeType_co, _DType_co]: ... + + @property + def device(self) -> _Device: ... + + @property + def mT(self) -> _arrayapi[Any, _DType_co]: ... + # NamedArray can most likely use both __array_function__ and __array_namespace__: _arrayfunction_or_api = (_arrayfunction, _arrayapi) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 683a1266472..e76d76c9678 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -27,6 +27,7 @@ ) from xarray.namedarray._aggregations import NamedArrayAggregations from xarray.namedarray._typing import ( + Default, ErrorOptionsWithWarn, _arrayapi, _arrayfunction_or_api, @@ -39,6 +40,7 @@ _sparsearrayfunction_or_api, _SupportsImag, _SupportsReal, + duckarray, ) from xarray.namedarray.parallelcompat import guess_chunkmanager from xarray.namedarray.pycompat import to_numpy @@ -51,22 +53,24 @@ ) if TYPE_CHECKING: + from enum import IntEnum + from numpy.typing import ArrayLike, NDArray from xarray.core.types import Dims, T_Chunks from xarray.namedarray._typing import ( - Default, _AttrsLike, _Chunks, + _Device, _Dim, _Dims, _DimsLike, _DType, + _IndexKeyLike, _IntOrUnknown, _ScalarType, _Shape, _ShapeType, - duckarray, ) from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint @@ -139,15 +143,15 @@ def _new( attributes you want to store with the array. Will copy the attrs from x by default. """ - dims_ = copy.copy(x._dims) if dims is _default else dims + dims_ = copy.copy(x._dims) if isinstance(dims, Default) else dims attrs_: Mapping[Any, Any] | None - if attrs is _default: + if isinstance(attrs, Default): attrs_ = None if x._attrs is None else x._attrs.copy() else: attrs_ = attrs - if data is _default: + if isinstance(data, Default): return type(x)(dims_, copy.copy(x._data), attrs_) else: cls_ = cast("type[NamedArray[_ShapeType, _DType]]", type(x)) @@ -404,35 +408,418 @@ def copy( """ return self._copy(deep=deep, data=data) - @property - def ndim(self) -> int: + def __len__(self) -> _IntOrUnknown: + try: + return self.shape[0] + except Exception as exc: + raise TypeError("len() of unsized object") from exc + + # < Array api > + + def _maybe_asarray( + self, x: bool | int | float | complex | NamedArray + ) -> NamedArray: """ - Number of array dimensions. + If x is a scalar, use asarray with the same dtype as self. + If it is namedarray already, respect the dtype and return it. - See Also - -------- - numpy.ndarray.ndim + Array API always promotes scalars to the same dtype as the other array. + Arrays are promoted according to result_types. """ - return len(self.shape) + from xarray.namedarray._array_api import asarray - @property - def size(self) -> _IntOrUnknown: + if isinstance(x, NamedArray): + # x is proper array. Respect the chosen dtype. + return x + # x is a scalar. Use the same dtype as self. + # TODO: Is this a good idea? x[Any, int] + 1.4 => int result then. + return asarray(x, dtype=self.dtype) + + # Required methods below: + + def __abs__(self, /) -> Self: + from xarray.namedarray._array_api import abs + + return abs(self) + + def __add__(self, other: int | float | NamedArray, /) -> NamedArray: + from xarray.namedarray._array_api import add + + return add(self, self._maybe_asarray(other)) + + def __and__(self, other: int | bool | NamedArray, /) -> NamedArray: + from xarray.namedarray._array_api import bitwise_and + + return bitwise_and(self, self._maybe_asarray(other)) + + def __array_namespace__(self, /, *, api_version: str | None = None): + if api_version is not None and api_version not in ( + "2021.12", + "2022.12", + "2023.12", + ): + raise ValueError(f"Unrecognized array API version: {api_version!r}") + import xarray.namedarray._array_api as array_api + + return array_api + + def __bool__(self, /) -> bool: + return self._data.__bool__() + + def __complex__(self, /) -> complex: + return self._data.__complex__() + + def __dlpack__( + self, + /, + *, + stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[IntEnum, int] | None = None, + copy: bool | None = None, + ) -> Any: + return self._data.__dlpack__( + stream=stream, max_version=max_version, dl_device=dl_device, copy=copy + ) + + def __dlpack_device__(self, /) -> tuple[IntEnum, int]: + return self._data.__dlpack_device__() + + def __eq__(self, other: int | float | bool | NamedArray, /) -> NamedArray: + from xarray.namedarray._array_api import equal + + return equal(self, self._maybe_asarray(other)) + + def __float__(self, /) -> float: + return self._data.__float__() + + def __floordiv__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import floor_divide + + return floor_divide(self, self._maybe_asarray(other)) + + def __ge__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import greater_equal + + return greater_equal(self, self._maybe_asarray(other)) + + def __getitem__(self, key: _IndexKeyLike | NamedArray) -> NamedArray: """ - Number of elements in the array. + Returns self[key]. - Equal to ``np.prod(a.shape)``, i.e., the product of the array’s dimensions. + Some rules: + * Integers removes the dim. + * Slices and ellipsis maintains same dim. + * None adds a dim. + * tuple follows above but on that specific axis. - See Also + Examples -------- - numpy.ndarray.size + + 1D + + >>> x = NamedArray(("x",), np.array([0, 1, 2])) + >>> key = NamedArray(("x",), np.array([1, 0, 0], dtype=bool)) + >>> xm = x[key] + >>> xm.dims, xm.shape + (('x',), (1,)) + + >>> x = NamedArray(("x",), np.array([0, 1, 2])) + >>> key = NamedArray(("x",), np.array([0, 0, 0], dtype=bool)) + >>> xm = x[key] + >>> xm.dims, xm.shape + (('x',), (0,)) + + Setup a ND array: + + >>> x = NamedArray(("x", "y"), np.arange(3 * 4).reshape((3, 4))) + >>> xm = x[0] + >>> xm.dims, xm.shape + (('y',), (4,)) + >>> xm = x[slice(0)] + >>> xm.dims, xm.shape + (('x', 'y'), (0, 4)) + >>> xm = x[None] + >>> xm.dims, xm.shape + (('dim_2', 'x', 'y'), (1, 3, 4)) + + >>> key = NamedArray(("x", "y"), np.ones((3, 4), dtype=bool)) + >>> xm = x[key] + >>> xm.dims, xm.shape + ((('x', 'y'),), (12,)) + + 0D + + >>> x = NamedArray((), np.array(False, dtype=np.bool)) + >>> key = NamedArray((), np.array(False, dtype=np.bool)) + >>> xm = x[key] + >>> xm.dims, xm.shape + (('dim_0',), (0,)) """ - return math.prod(self.shape) + from xarray.namedarray._array_api._manipulation_functions import ( + _broadcast_arrays, + ) + from xarray.namedarray._array_api._utils import ( + _atleast1d_dims, + _dims_from_tuple_indexing, + _flatten_dims, + ) - def __len__(self) -> _IntOrUnknown: - try: - return self.shape[0] - except Exception as exc: - raise TypeError("len() of unsized object") from exc + if isinstance(key, NamedArray): + self_new, key_new = _broadcast_arrays(self, key) + _data = self_new._data[key_new._data] + _dims = _flatten_dims(_atleast1d_dims(self_new.dims)) + return self._new(_dims, _data) + # elif isinstance(key, int): + # return self._new(self.dims[1:], self._data[key]) + # elif isinstance(key, slice) or key is ...: + # return self._new(self.dims, self._data[key]) + # elif key is None: + # return expand_dims(self) + # elif isinstance(key, tuple): + # _dims = _dims_from_tuple_indexing(self.dims, key) + # return self._new(_dims, self._data[key]) + + elif isinstance(key, int | slice | tuple) or key is None or key is ...: + # TODO: __getitem__ not always available, use expand_dims + _data = self._data[key] + _dims = _dims_from_tuple_indexing( + self.dims, key if isinstance(key, tuple) else (key,) + ) + return self._new(_dims, _data) + else: + raise NotImplementedError(f"{key=} is not supported") + + def __gt__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import greater + + return greater(self, self._maybe_asarray(other)) + + def __index__(self, /) -> int: + return self._data.__index__() + + def __int__(self, /) -> int: + return self._data.__int__() + + def __invert__(self, /): + from xarray.namedarray._array_api import bitwise_invert + + return bitwise_invert(self) + + def __iter__(self: NamedArray, /): + from xarray.namedarray._array_api import asarray + + # TODO: smarter way to retain dims, xarray? + return (asarray(i) for i in self._data) + + def __le__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import less_equal + + return less_equal(self, self._maybe_asarray(other)) + + def __lshift__(self, other: int | NamedArray, /): + from xarray.namedarray._array_api import bitwise_left_shift + + return bitwise_left_shift(self, self._maybe_asarray(other)) + + def __lt__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import less + + return less(self, self._maybe_asarray(other)) + + def __matmul__(self, other: NamedArray, /): + from xarray.namedarray._array_api import matmul + + return matmul(self, self._maybe_asarray(other)) + + def __mod__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import remainder + + return remainder(self, self._maybe_asarray(other)) + + def __mul__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import multiply + + return multiply(self, self._maybe_asarray(other)) + + def __ne__(self, other: int | float | bool | NamedArray, /): + from xarray.namedarray._array_api import not_equal + + return not_equal(self, self._maybe_asarray(other)) + + def __neg__(self, /): + from xarray.namedarray._array_api import negative + + return negative(self) + + def __or__(self, other: int | bool | NamedArray, /): + from xarray.namedarray._array_api import bitwise_or + + return bitwise_or(self, self._maybe_asarray(other)) + + def __pos__(self, /): + from xarray.namedarray._array_api import positive + + return positive(self) + + def __pow__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import pow + + return pow(self, self._maybe_asarray(other)) + + def __rshift__(self, other: int | NamedArray, /): + from xarray.namedarray._array_api import bitwise_right_shift + + return bitwise_right_shift(self, self._maybe_asarray(other)) + + def __setitem__( + self, + key: _IndexKeyLike, + value: int | float | bool | NamedArray, + /, + ) -> None: + if isinstance(key, NamedArray): + key = key._data + self._data.__setitem__(key, self._maybe_asarray(value)._data) + + def __sub__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import subtract + + return subtract(self, self._maybe_asarray(other)) + + def __truediv__(self, other: float | NamedArray, /): + from xarray.namedarray._array_api import divide + + return divide(self, self._maybe_asarray(other)) + + def __xor__(self, other: int | bool | NamedArray, /): + from xarray.namedarray._array_api import bitwise_xor + + return bitwise_xor(self, self._maybe_asarray(other)) + + def __iadd__(self, other: int | float | NamedArray, /): + self._data += self._maybe_asarray(other)._data + return self + + def __radd__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import add + + return add(self._maybe_asarray(other), self) + + def __iand__(self, other: int | bool | NamedArray, /): + self._data &= self._maybe_asarray(other)._data + return self + + def __rand__(self, other: int | bool | NamedArray, /): + from xarray.namedarray._array_api import bitwise_and + + return bitwise_and(self._maybe_asarray(other), self) + + def __ifloordiv__(self, other: int | float | NamedArray, /): + self._data //= self._maybe_asarray(other)._data + return self + + def __rfloordiv__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import floor_divide + + return floor_divide(self._maybe_asarray(other), self) + + def __ilshift__(self, other: int | NamedArray, /): + self._data <<= self._maybe_asarray(other)._data + return self + + def __rlshift__(self, other: int | NamedArray, /): + from xarray.namedarray._array_api import bitwise_left_shift + + return bitwise_left_shift(self._maybe_asarray(other), self) + + def __imatmul__(self, other: NamedArray, /): + self._data @= other._data + return self + + def __rmatmul__(self, other: NamedArray, /): + from xarray.namedarray._array_api import matmul + + return matmul(self._maybe_asarray(other), self) + + def __imod__(self, other: int | float | NamedArray, /): + self._data %= self._maybe_asarray(other)._data + return self + + def __rmod__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import remainder + + return remainder(self._maybe_asarray(other), self) + + def __imul__(self, other: int | float | NamedArray, /): + self._data *= self._maybe_asarray(other)._data + return self + + def __rmul__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import multiply + + return multiply(self._maybe_asarray(other), self) + + def __ior__(self, other: int | bool | NamedArray, /): + self._data |= self._maybe_asarray(other)._data + + return self + + def __ror__(self, other: int | bool | NamedArray, /): + from xarray.namedarray._array_api import bitwise_or + + return bitwise_or(self._maybe_asarray(other), self) + + def __ipow__(self, other: int | float | NamedArray, /): + self._data **= self._maybe_asarray(other)._data + return self + + def __rpow__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import pow + + return pow(self._maybe_asarray(other), self) + + def __irshift__(self, other: int | NamedArray, /): + self._data >>= self._maybe_asarray(other)._data + return self + + def __rrshift__(self, other: int | NamedArray, /): + from xarray.namedarray._array_api import bitwise_right_shift + + return bitwise_right_shift(self._maybe_asarray(other), self) + + def __isub__(self, other: int | float | NamedArray, /): + self._data -= self._maybe_asarray(other)._data + return self + + def __rsub__(self, other: int | float | NamedArray, /): + from xarray.namedarray._array_api import subtract + + return subtract(self._maybe_asarray(other), self) + + def __itruediv__(self, other: float | NamedArray, /): + self._data /= self._maybe_asarray(other)._data + return self + + def __rtruediv__(self, other: float | NamedArray, /): + from xarray.namedarray._array_api import divide + + return divide(self._maybe_asarray(other), self) + + def __ixor__(self, other: int | bool | NamedArray, /): + self._data ^= self._maybe_asarray(other)._data + return self + + def __rxor__(self, other, /): + from xarray.namedarray._array_api import bitwise_xor + + return bitwise_xor(self._maybe_asarray(other), self) + + def to_device(self, device: _Device, /, stream: None = None) -> Self: + if isinstance(self._data, _arrayapi): + return self._replace(data=self._data.to_device(device, stream=stream)) + else: + raise NotImplementedError("Only array api are valid.") @property def dtype(self) -> _DType_co: @@ -446,6 +833,42 @@ def dtype(self) -> _DType_co: """ return self._data.dtype + @property + def device(self) -> _Device: + """ + Device of the array’s elements. + + See Also + -------- + ndarray.device + """ + if isinstance(self._data, _arrayapi): + return self._data.device + else: + raise NotImplementedError("self._data missing device") + + @property + def mT(self) -> NamedArray[Any, _DType_co]: + if isinstance(self._data, _arrayapi): + from xarray.namedarray._array_api._utils import _infer_dims + + _data = self._data.mT + _dims = _infer_dims(_data.shape) + return self._new(_dims, _data) + else: + raise NotImplementedError("self._data missing mT") + + @property + def ndim(self) -> int: + """ + Number of array dimensions. + + See Also + -------- + numpy.ndarray.ndim + """ + return len(self.shape) + @property def shape(self) -> _Shape: """ @@ -462,6 +885,31 @@ def shape(self) -> _Shape: """ return self._data.shape + @property + def size(self) -> _IntOrUnknown: + """ + Number of elements in the array. + + Equal to ``np.prod(a.shape)``, i.e., the product of the array’s dimensions. + + See Also + -------- + numpy.ndarray.size + """ + return math.prod(self.shape) + + @property + def T(self) -> NamedArray[Any, _DType_co]: + """Return a new object with transposed dimensions.""" + if self.ndim != 2: + raise ValueError( + f"x.T requires x to have 2 dimensions, got {self.ndim}. Use x.permute_dims() to permute dimensions." + ) + + return self.permute_dims() + + # + @property def nbytes(self) -> _IntOrUnknown: """ @@ -470,7 +918,7 @@ def nbytes(self) -> _IntOrUnknown: If the underlying data array does not include ``nbytes``, estimates the bytes consumed based on the ``size`` and ``dtype``. """ - from xarray.namedarray._array_api import _get_data_namespace + from xarray.namedarray._array_api._utils import _get_data_namespace if hasattr(self._data, "nbytes"): return self._data.nbytes # type: ignore[no-any-return] @@ -589,6 +1037,19 @@ def real( return real(self) return self._new(data=self._data.real) + @overload + def __array__( + self, dtype: None = ..., /, *, copy: bool | None = ... + ) -> np.ndarray[Any, _DType_co]: ... + @overload + def __array__( + self, dtype: _DType, /, *, copy: bool | None = ... + ) -> np.ndarray[Any, _DType]: ... + def __array__( + self, dtype: _DType | None = ..., /, *, copy: bool | None = ... + ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]: + return np.asarray(self._data) + def __dask_tokenize__(self) -> object: # Use v.data, instead of v._data, in order to cope with the wrappers # around NetCDF and the like @@ -979,12 +1440,12 @@ def _as_sparse( from xarray.namedarray._array_api import astype # TODO: what to do if dask-backended? - if fill_value is _default: + if isinstance(fill_value, Default): dtype, fill_value = dtypes.maybe_promote(self.dtype) else: dtype = dtypes.result_type(self.dtype, fill_value) - if sparse_format is _default: + if isinstance(sparse_format, Default): sparse_format = "coo" try: as_sparse = getattr(sparse, f"as_{sparse_format.lower()}") @@ -1052,16 +1513,6 @@ def permute_dims( return permute_dims(self, axes) - @property - def T(self) -> NamedArray[Any, _DType_co]: - """Return a new object with transposed dimensions.""" - if self.ndim != 2: - raise ValueError( - f"x.T requires x to have 2 dimensions, got {self.ndim}. Use x.permute_dims() to permute dimensions." - ) - - return self.permute_dims() - def broadcast_to( self, dim: Mapping[_Dim, int] | None = None, **dim_kwargs: Any ) -> NamedArray[Any, _DType_co]: @@ -1089,7 +1540,7 @@ def broadcast_to( Examples -------- >>> data = np.asarray([[1.0, 2.0], [3.0, 4.0]]) - >>> array = xr.NamedArray(("x", "y"), data) + >>> array = NamedArray(("x", "y"), data) >>> array.sizes {'x': 2, 'y': 2} @@ -1144,16 +1595,11 @@ def expand_dims( Examples -------- - >>> data = np.asarray([[1.0, 2.0], [3.0, 4.0]]) - >>> array = xr.NamedArray(("x", "y"), data) - - - # expand dimensions by specifying a new dimension name - >>> expanded = array.expand_dims(dim="z") + >>> x = NamedArray(("x", "y"), data) + >>> expanded = x.expand_dims(dim="z") >>> expanded.dims ('z', 'x', 'y') - """ from xarray.namedarray._array_api import expand_dims @@ -1172,3 +1618,9 @@ def _raise_if_any_duplicate_dimensions( raise ValueError( f"{err_context} cannot handle duplicate dimensions, but dimensions {repeated_dims} appear more than once on this object's dims: {dims}" ) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/xarray/tests/namedarray_array_api_skips.txt b/xarray/tests/namedarray_array_api_skips.txt new file mode 100644 index 00000000000..28fb8c836fe --- /dev/null +++ b/xarray/tests/namedarray_array_api_skips.txt @@ -0,0 +1,44 @@ +# finfo(float32).eps returns float32 but should return float +array_api_tests/test_data_type_functions.py::test_finfo[float32] + +# NumPy deviates in some special cases for floordiv +array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity] +array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity] +array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity] +array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity] +array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0] +array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0] +array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity] +array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity] +array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity] +array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity] +array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0] +array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0] +array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity] +array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity] +array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity] +array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity] +array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0] +array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0] + +# https://github.com/numpy/numpy/issues/21213 +array_api_tests/test_special_cases.py::test_binary[__pow__(x1_i is -infinity and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +infinity] +array_api_tests/test_special_cases.py::test_binary[__pow__(x1_i is -0 and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +0] +array_api_tests/test_special_cases.py::test_iop[__ipow__(x1_i is -infinity and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +infinity] +array_api_tests/test_special_cases.py::test_iop[__ipow__(x1_i is -0 and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +0] +array_api_tests/meta/test_hypothesis_helpers.py::test_symmetric_matrices + +# The test suite is incorrectly checking sums that have loss of significance +# (https://github.com/data-apis/array-api-tests/issues/168) +array_api_tests/test_statistical_functions.py::test_sum +array_api_tests/test_statistical_functions.py::test_prod + +# The test suite cannot properly get the signature for vecdot +# https://github.com/numpy/numpy/pull/26237 +array_api_tests/test_signatures.py::test_func_signature[vecdot] +array_api_tests/test_signatures.py::test_extension_func_signature[linalg.vecdot] + +# numpy scalars missing __complex__: +# AttributeError: 'numpy.uint64' object has no attribute '__complex__' +# https://github.com/numpy/numpy/issues/27305 +array_api_tests/test_statistical_functions.py::test_cumulative_sum