612 lines
22 KiB
Python
612 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
import sys
|
|
from abc import abstractmethod
|
|
from collections.abc import Mapping
|
|
from typing import TYPE_CHECKING, Any, Generic, cast, overload
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from packaging.version import Version
|
|
|
|
from xarray.core.indexing import ExplicitlyIndexed
|
|
from xarray.namedarray._typing import (
|
|
_arrayfunction_or_api,
|
|
_default,
|
|
_DType_co,
|
|
_ShapeType_co,
|
|
)
|
|
from xarray.namedarray.core import NamedArray, from_array
|
|
|
|
if TYPE_CHECKING:
|
|
from types import ModuleType
|
|
|
|
from numpy.typing import ArrayLike, DTypeLike, NDArray
|
|
|
|
from xarray.namedarray._typing import (
|
|
Default,
|
|
_AttrsLike,
|
|
_Dim,
|
|
_DimsLike,
|
|
_DType,
|
|
_IndexKeyLike,
|
|
_IntOrUnknown,
|
|
_Shape,
|
|
_ShapeLike,
|
|
duckarray,
|
|
)
|
|
|
|
|
|
class CustomArrayBase(Generic[_ShapeType_co, _DType_co]):
|
|
def __init__(self, array: duckarray[Any, _DType_co]) -> None:
|
|
self.array: duckarray[Any, _DType_co] = array
|
|
|
|
@property
|
|
def dtype(self) -> _DType_co:
|
|
return self.array.dtype
|
|
|
|
@property
|
|
def shape(self) -> _Shape:
|
|
return self.array.shape
|
|
|
|
|
|
class CustomArray(
|
|
CustomArrayBase[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co]
|
|
):
|
|
def __array__(
|
|
self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None
|
|
) -> np.ndarray[Any, np.dtype[np.generic]]:
|
|
if Version(np.__version__) >= Version("2.0.0"):
|
|
return np.asarray(self.array, dtype=dtype, copy=copy)
|
|
else:
|
|
return np.asarray(self.array, dtype=dtype)
|
|
|
|
|
|
class CustomArrayIndexable(
|
|
CustomArrayBase[_ShapeType_co, _DType_co],
|
|
ExplicitlyIndexed,
|
|
Generic[_ShapeType_co, _DType_co],
|
|
):
|
|
def __getitem__(
|
|
self, key: _IndexKeyLike | CustomArrayIndexable[Any, Any], /
|
|
) -> CustomArrayIndexable[Any, _DType_co]:
|
|
if isinstance(key, CustomArrayIndexable):
|
|
if isinstance(key.array, type(self.array)):
|
|
# TODO: key.array is duckarray here, can it be narrowed down further?
|
|
# an _arrayapi cannot be used on a _arrayfunction for example.
|
|
return type(self)(array=self.array[key.array]) # type: ignore[index]
|
|
else:
|
|
raise TypeError("key must have the same array type as self")
|
|
else:
|
|
return type(self)(array=self.array[key])
|
|
|
|
def __array_namespace__(self) -> ModuleType:
|
|
return np
|
|
|
|
|
|
def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]:
|
|
# Mypy checks a is valid:
|
|
b: duckarray[Any, _DType] = a
|
|
|
|
# Runtime check if valid:
|
|
if isinstance(b, _arrayfunction_or_api):
|
|
return b
|
|
else:
|
|
missing_attrs = ""
|
|
actual_attrs = set(dir(b))
|
|
for t in _arrayfunction_or_api:
|
|
if sys.version_info >= (3, 13):
|
|
# https://github.com/python/cpython/issues/104873
|
|
from typing import get_protocol_members
|
|
|
|
expected_attrs = get_protocol_members(t)
|
|
elif sys.version_info >= (3, 12):
|
|
expected_attrs = t.__protocol_attrs__
|
|
else:
|
|
from typing import _get_protocol_attrs # type: ignore[attr-defined]
|
|
|
|
expected_attrs = _get_protocol_attrs(t)
|
|
|
|
missing_attrs_ = expected_attrs - actual_attrs
|
|
if missing_attrs_:
|
|
missing_attrs += f"{t.__name__} - {missing_attrs_}\n"
|
|
raise TypeError(
|
|
f"a ({type(a)}) is not a valid _arrayfunction or _arrayapi. "
|
|
"Missing following attrs:\n"
|
|
f"{missing_attrs}"
|
|
)
|
|
|
|
|
|
class NamedArraySubclassobjects:
|
|
@pytest.fixture
|
|
def target(self, data: np.ndarray[Any, Any]) -> Any:
|
|
"""Fixture that needs to be overridden"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def cls(self, *args: Any, **kwargs: Any) -> Any:
|
|
"""Method that needs to be overridden"""
|
|
raise NotImplementedError
|
|
|
|
@pytest.fixture
|
|
def data(self) -> np.ndarray[Any, np.dtype[Any]]:
|
|
return 0.5 * np.arange(10).reshape(2, 5)
|
|
|
|
@pytest.fixture
|
|
def random_inputs(self) -> np.ndarray[Any, np.dtype[np.float32]]:
|
|
return np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5))
|
|
|
|
def test_properties(self, target: Any, data: Any) -> None:
|
|
assert target.dims == ("x", "y")
|
|
assert np.array_equal(target.data, data)
|
|
assert target.dtype == float
|
|
assert target.shape == (2, 5)
|
|
assert target.ndim == 2
|
|
assert target.sizes == {"x": 2, "y": 5}
|
|
assert target.size == 10
|
|
assert target.nbytes == 80
|
|
assert len(target) == 2
|
|
|
|
def test_attrs(self, target: Any) -> None:
|
|
assert target.attrs == {}
|
|
attrs = {"foo": "bar"}
|
|
target.attrs = attrs
|
|
assert target.attrs == attrs
|
|
assert isinstance(target.attrs, dict)
|
|
target.attrs["foo"] = "baz"
|
|
assert target.attrs["foo"] == "baz"
|
|
|
|
@pytest.mark.parametrize(
|
|
"expected", [np.array([1, 2], dtype=np.dtype(np.int8)), [1, 2]]
|
|
)
|
|
def test_init(self, expected: Any) -> None:
|
|
actual = self.cls(("x",), expected)
|
|
assert np.array_equal(np.asarray(actual.data), expected)
|
|
|
|
actual = self.cls(("x",), expected)
|
|
assert np.array_equal(np.asarray(actual.data), expected)
|
|
|
|
def test_data(self, random_inputs: Any) -> None:
|
|
expected = self.cls(["x", "y", "z"], random_inputs)
|
|
assert np.array_equal(np.asarray(expected.data), random_inputs)
|
|
with pytest.raises(ValueError):
|
|
expected.data = np.random.random((3, 4)).astype(np.float64)
|
|
d2 = np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5))
|
|
expected.data = d2
|
|
assert np.array_equal(np.asarray(expected.data), d2)
|
|
|
|
|
|
class TestNamedArray(NamedArraySubclassobjects):
|
|
def cls(self, *args: Any, **kwargs: Any) -> NamedArray[Any, Any]:
|
|
return NamedArray(*args, **kwargs)
|
|
|
|
@pytest.fixture
|
|
def target(self, data: np.ndarray[Any, Any]) -> NamedArray[Any, Any]:
|
|
return NamedArray(["x", "y"], data)
|
|
|
|
@pytest.mark.parametrize(
|
|
"expected",
|
|
[
|
|
np.array([1, 2], dtype=np.dtype(np.int8)),
|
|
pytest.param(
|
|
[1, 2],
|
|
marks=pytest.mark.xfail(
|
|
reason="NamedArray only supports array-like objects"
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_init(self, expected: Any) -> None:
|
|
super().test_init(expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"dims, data, expected, raise_error",
|
|
[
|
|
(("x",), [1, 2, 3], np.array([1, 2, 3]), False),
|
|
((1,), np.array([4, 5, 6]), np.array([4, 5, 6]), False),
|
|
((), 2, np.array(2), False),
|
|
# Fail:
|
|
(
|
|
("x",),
|
|
NamedArray("time", np.array([1, 2, 3], dtype=np.dtype(np.int64))),
|
|
np.array([1, 2, 3]),
|
|
True,
|
|
),
|
|
],
|
|
)
|
|
def test_from_array(
|
|
self,
|
|
dims: _DimsLike,
|
|
data: ArrayLike,
|
|
expected: np.ndarray[Any, Any],
|
|
raise_error: bool,
|
|
) -> None:
|
|
actual: NamedArray[Any, Any]
|
|
if raise_error:
|
|
with pytest.raises(TypeError, match="already a Named array"):
|
|
actual = from_array(dims, data)
|
|
|
|
# Named arrays are not allowed:
|
|
from_array(actual) # type: ignore[call-overload]
|
|
else:
|
|
actual = from_array(dims, data)
|
|
|
|
assert np.array_equal(np.asarray(actual.data), expected)
|
|
|
|
def test_from_array_with_masked_array(self) -> None:
|
|
masked_array: np.ndarray[Any, np.dtype[np.generic]]
|
|
masked_array = np.ma.array([1, 2, 3], mask=[False, True, False]) # type: ignore[no-untyped-call]
|
|
with pytest.raises(NotImplementedError):
|
|
from_array(("x",), masked_array)
|
|
|
|
def test_from_array_with_0d_object(self) -> None:
|
|
data = np.empty((), dtype=object)
|
|
data[()] = (10, 12, 12)
|
|
narr = from_array((), data)
|
|
np.array_equal(np.asarray(narr.data), data)
|
|
|
|
# TODO: Make xr.core.indexing.ExplicitlyIndexed pass as a subclass of_arrayfunction_or_api
|
|
# and remove this test.
|
|
def test_from_array_with_explicitly_indexed(
|
|
self, random_inputs: np.ndarray[Any, Any]
|
|
) -> None:
|
|
array: CustomArray[Any, Any]
|
|
array = CustomArray(random_inputs)
|
|
output: NamedArray[Any, Any]
|
|
output = from_array(("x", "y", "z"), array)
|
|
assert isinstance(output.data, np.ndarray)
|
|
|
|
array2: CustomArrayIndexable[Any, Any]
|
|
array2 = CustomArrayIndexable(random_inputs)
|
|
output2: NamedArray[Any, Any]
|
|
output2 = from_array(("x", "y", "z"), array2)
|
|
assert isinstance(output2.data, CustomArrayIndexable)
|
|
|
|
def test_real_and_imag(self) -> None:
|
|
expected_real: np.ndarray[Any, np.dtype[np.float64]]
|
|
expected_real = np.arange(3, dtype=np.float64)
|
|
|
|
expected_imag: np.ndarray[Any, np.dtype[np.float64]]
|
|
expected_imag = -np.arange(3, dtype=np.float64)
|
|
|
|
arr: np.ndarray[Any, np.dtype[np.complex128]]
|
|
arr = expected_real + 1j * expected_imag
|
|
|
|
named_array: NamedArray[Any, np.dtype[np.complex128]]
|
|
named_array = NamedArray(["x"], arr)
|
|
|
|
actual_real: duckarray[Any, np.dtype[np.float64]] = named_array.real.data
|
|
assert np.array_equal(np.asarray(actual_real), expected_real)
|
|
assert actual_real.dtype == expected_real.dtype
|
|
|
|
actual_imag: duckarray[Any, np.dtype[np.float64]] = named_array.imag.data
|
|
assert np.array_equal(np.asarray(actual_imag), expected_imag)
|
|
assert actual_imag.dtype == expected_imag.dtype
|
|
|
|
# Additional tests as per your original class-based code
|
|
@pytest.mark.parametrize(
|
|
"data, dtype",
|
|
[
|
|
("foo", np.dtype("U3")),
|
|
(b"foo", np.dtype("S3")),
|
|
],
|
|
)
|
|
def test_from_array_0d_string(self, data: Any, dtype: DTypeLike) -> None:
|
|
named_array: NamedArray[Any, Any]
|
|
named_array = from_array([], data)
|
|
assert named_array.data == data
|
|
assert named_array.dims == ()
|
|
assert named_array.sizes == {}
|
|
assert named_array.attrs == {}
|
|
assert named_array.ndim == 0
|
|
assert named_array.size == 1
|
|
assert named_array.dtype == dtype
|
|
|
|
def test_from_array_0d_object(self) -> None:
|
|
named_array: NamedArray[Any, Any]
|
|
named_array = from_array([], (10, 12, 12))
|
|
expected_data = np.empty((), dtype=object)
|
|
expected_data[()] = (10, 12, 12)
|
|
assert np.array_equal(np.asarray(named_array.data), expected_data)
|
|
|
|
assert named_array.dims == ()
|
|
assert named_array.sizes == {}
|
|
assert named_array.attrs == {}
|
|
assert named_array.ndim == 0
|
|
assert named_array.size == 1
|
|
assert named_array.dtype == np.dtype("O")
|
|
|
|
def test_from_array_0d_datetime(self) -> None:
|
|
named_array: NamedArray[Any, Any]
|
|
named_array = from_array([], np.datetime64("2000-01-01"))
|
|
assert named_array.dtype == np.dtype("datetime64[D]")
|
|
|
|
@pytest.mark.parametrize(
|
|
"timedelta, expected_dtype",
|
|
[
|
|
(np.timedelta64(1, "D"), np.dtype("timedelta64[D]")),
|
|
(np.timedelta64(1, "s"), np.dtype("timedelta64[s]")),
|
|
(np.timedelta64(1, "m"), np.dtype("timedelta64[m]")),
|
|
(np.timedelta64(1, "h"), np.dtype("timedelta64[h]")),
|
|
(np.timedelta64(1, "us"), np.dtype("timedelta64[us]")),
|
|
(np.timedelta64(1, "ns"), np.dtype("timedelta64[ns]")),
|
|
(np.timedelta64(1, "ps"), np.dtype("timedelta64[ps]")),
|
|
(np.timedelta64(1, "fs"), np.dtype("timedelta64[fs]")),
|
|
(np.timedelta64(1, "as"), np.dtype("timedelta64[as]")),
|
|
],
|
|
)
|
|
def test_from_array_0d_timedelta(
|
|
self, timedelta: np.timedelta64, expected_dtype: np.dtype[np.timedelta64]
|
|
) -> None:
|
|
named_array: NamedArray[Any, Any]
|
|
named_array = from_array([], timedelta)
|
|
assert named_array.dtype == expected_dtype
|
|
assert named_array.data == timedelta
|
|
|
|
@pytest.mark.parametrize(
|
|
"dims, data_shape, new_dims, raises",
|
|
[
|
|
(["x", "y", "z"], (2, 3, 4), ["a", "b", "c"], False),
|
|
(["x", "y", "z"], (2, 3, 4), ["a", "b"], True),
|
|
(["x", "y", "z"], (2, 4, 5), ["a", "b", "c", "d"], True),
|
|
([], [], (), False),
|
|
([], [], ("x",), True),
|
|
],
|
|
)
|
|
def test_dims_setter(
|
|
self, dims: Any, data_shape: Any, new_dims: Any, raises: bool
|
|
) -> None:
|
|
named_array: NamedArray[Any, Any]
|
|
named_array = NamedArray(dims, np.asarray(np.random.random(data_shape)))
|
|
assert named_array.dims == tuple(dims)
|
|
if raises:
|
|
with pytest.raises(ValueError):
|
|
named_array.dims = new_dims
|
|
else:
|
|
named_array.dims = new_dims
|
|
assert named_array.dims == tuple(new_dims)
|
|
|
|
def test_duck_array_class(self) -> None:
|
|
numpy_a: NDArray[np.int64]
|
|
numpy_a = np.array([2.1, 4], dtype=np.dtype(np.int64))
|
|
check_duck_array_typevar(numpy_a)
|
|
|
|
masked_a: np.ma.MaskedArray[Any, np.dtype[np.int64]]
|
|
masked_a = np.ma.asarray([2.1, 4], dtype=np.dtype(np.int64)) # type: ignore[no-untyped-call]
|
|
check_duck_array_typevar(masked_a)
|
|
|
|
custom_a: CustomArrayIndexable[Any, np.dtype[np.int64]]
|
|
custom_a = CustomArrayIndexable(numpy_a)
|
|
check_duck_array_typevar(custom_a)
|
|
|
|
def test_duck_array_class_array_api(self) -> None:
|
|
# Test numpy's array api:
|
|
nxp = pytest.importorskip("array_api_strict", minversion="1.0")
|
|
|
|
# TODO: nxp doesn't use dtype typevars, so can only use Any for the moment:
|
|
arrayapi_a: duckarray[Any, Any] # duckarray[Any, np.dtype[np.int64]]
|
|
arrayapi_a = nxp.asarray([2.1, 4], dtype=nxp.int64)
|
|
check_duck_array_typevar(arrayapi_a)
|
|
|
|
def test_new_namedarray(self) -> None:
|
|
dtype_float = np.dtype(np.float32)
|
|
narr_float: NamedArray[Any, np.dtype[np.float32]]
|
|
narr_float = NamedArray(("x",), np.array([1.5, 3.2], dtype=dtype_float))
|
|
assert narr_float.dtype == dtype_float
|
|
|
|
dtype_int = np.dtype(np.int8)
|
|
narr_int: NamedArray[Any, np.dtype[np.int8]]
|
|
narr_int = narr_float._new(("x",), np.array([1, 3], dtype=dtype_int))
|
|
assert narr_int.dtype == dtype_int
|
|
|
|
class Variable(
|
|
NamedArray[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co]
|
|
):
|
|
@overload
|
|
def _new(
|
|
self,
|
|
dims: _DimsLike | Default = ...,
|
|
data: duckarray[Any, _DType] = ...,
|
|
attrs: _AttrsLike | Default = ...,
|
|
) -> Variable[Any, _DType]: ...
|
|
|
|
@overload
|
|
def _new(
|
|
self,
|
|
dims: _DimsLike | Default = ...,
|
|
data: Default = ...,
|
|
attrs: _AttrsLike | Default = ...,
|
|
) -> Variable[_ShapeType_co, _DType_co]: ...
|
|
|
|
def _new(
|
|
self,
|
|
dims: _DimsLike | Default = _default,
|
|
data: duckarray[Any, _DType] | Default = _default,
|
|
attrs: _AttrsLike | Default = _default,
|
|
) -> Variable[Any, _DType] | Variable[_ShapeType_co, _DType_co]:
|
|
dims_ = copy.copy(self._dims) if dims is _default else dims
|
|
|
|
attrs_: Mapping[Any, Any] | None
|
|
if attrs is _default:
|
|
attrs_ = None if self._attrs is None else self._attrs.copy()
|
|
else:
|
|
attrs_ = attrs
|
|
|
|
if data is _default:
|
|
return type(self)(dims_, copy.copy(self._data), attrs_)
|
|
cls_ = cast("type[Variable[Any, _DType]]", type(self))
|
|
return cls_(dims_, data, attrs_)
|
|
|
|
var_float: Variable[Any, np.dtype[np.float32]]
|
|
var_float = Variable(("x",), np.array([1.5, 3.2], dtype=dtype_float))
|
|
assert var_float.dtype == dtype_float
|
|
|
|
var_int: Variable[Any, np.dtype[np.int8]]
|
|
var_int = var_float._new(("x",), np.array([1, 3], dtype=dtype_int))
|
|
assert var_int.dtype == dtype_int
|
|
|
|
def test_replace_namedarray(self) -> None:
|
|
dtype_float = np.dtype(np.float32)
|
|
np_val: np.ndarray[Any, np.dtype[np.float32]]
|
|
np_val = np.array([1.5, 3.2], dtype=dtype_float)
|
|
np_val2: np.ndarray[Any, np.dtype[np.float32]]
|
|
np_val2 = 2 * np_val
|
|
|
|
narr_float: NamedArray[Any, np.dtype[np.float32]]
|
|
narr_float = NamedArray(("x",), np_val)
|
|
assert narr_float.dtype == dtype_float
|
|
|
|
narr_float2: NamedArray[Any, np.dtype[np.float32]]
|
|
narr_float2 = NamedArray(("x",), np_val2)
|
|
assert narr_float2.dtype == dtype_float
|
|
|
|
class Variable(
|
|
NamedArray[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co]
|
|
):
|
|
@overload
|
|
def _new(
|
|
self,
|
|
dims: _DimsLike | Default = ...,
|
|
data: duckarray[Any, _DType] = ...,
|
|
attrs: _AttrsLike | Default = ...,
|
|
) -> Variable[Any, _DType]: ...
|
|
|
|
@overload
|
|
def _new(
|
|
self,
|
|
dims: _DimsLike | Default = ...,
|
|
data: Default = ...,
|
|
attrs: _AttrsLike | Default = ...,
|
|
) -> Variable[_ShapeType_co, _DType_co]: ...
|
|
|
|
def _new(
|
|
self,
|
|
dims: _DimsLike | Default = _default,
|
|
data: duckarray[Any, _DType] | Default = _default,
|
|
attrs: _AttrsLike | Default = _default,
|
|
) -> Variable[Any, _DType] | Variable[_ShapeType_co, _DType_co]:
|
|
dims_ = copy.copy(self._dims) if dims is _default else dims
|
|
|
|
attrs_: Mapping[Any, Any] | None
|
|
if attrs is _default:
|
|
attrs_ = None if self._attrs is None else self._attrs.copy()
|
|
else:
|
|
attrs_ = attrs
|
|
|
|
if data is _default:
|
|
return type(self)(dims_, copy.copy(self._data), attrs_)
|
|
cls_ = cast("type[Variable[Any, _DType]]", type(self))
|
|
return cls_(dims_, data, attrs_)
|
|
|
|
var_float: Variable[Any, np.dtype[np.float32]]
|
|
var_float = Variable(("x",), np_val)
|
|
assert var_float.dtype == dtype_float
|
|
|
|
var_float2: Variable[Any, np.dtype[np.float32]]
|
|
var_float2 = var_float._replace(("x",), np_val2)
|
|
assert var_float2.dtype == dtype_float
|
|
|
|
@pytest.mark.parametrize(
|
|
"dim,expected_ndim,expected_shape,expected_dims",
|
|
[
|
|
(None, 3, (1, 2, 5), (None, "x", "y")),
|
|
(_default, 3, (1, 2, 5), ("dim_2", "x", "y")),
|
|
("z", 3, (1, 2, 5), ("z", "x", "y")),
|
|
],
|
|
)
|
|
def test_expand_dims(
|
|
self,
|
|
target: NamedArray[Any, np.dtype[np.float32]],
|
|
dim: _Dim | Default,
|
|
expected_ndim: int,
|
|
expected_shape: _ShapeLike,
|
|
expected_dims: _DimsLike,
|
|
) -> None:
|
|
result = target.expand_dims(dim=dim)
|
|
assert result.ndim == expected_ndim
|
|
assert result.shape == expected_shape
|
|
assert result.dims == expected_dims
|
|
|
|
@pytest.mark.parametrize(
|
|
"dims, expected_sizes",
|
|
[
|
|
((), {"y": 5, "x": 2}),
|
|
(["y", "x"], {"y": 5, "x": 2}),
|
|
(["y", ...], {"y": 5, "x": 2}),
|
|
],
|
|
)
|
|
def test_permute_dims(
|
|
self,
|
|
target: NamedArray[Any, np.dtype[np.float32]],
|
|
dims: _DimsLike,
|
|
expected_sizes: dict[_Dim, _IntOrUnknown],
|
|
) -> None:
|
|
actual = target.permute_dims(*dims)
|
|
assert actual.sizes == expected_sizes
|
|
|
|
def test_permute_dims_errors(
|
|
self,
|
|
target: NamedArray[Any, np.dtype[np.float32]],
|
|
) -> None:
|
|
with pytest.raises(ValueError, match=r"'y'.*permuted list"):
|
|
dims = ["y"]
|
|
target.permute_dims(*dims)
|
|
|
|
@pytest.mark.parametrize(
|
|
"broadcast_dims,expected_ndim",
|
|
[
|
|
({"x": 2, "y": 5}, 2),
|
|
({"x": 2, "y": 5, "z": 2}, 3),
|
|
({"w": 1, "x": 2, "y": 5}, 3),
|
|
],
|
|
)
|
|
def test_broadcast_to(
|
|
self,
|
|
target: NamedArray[Any, np.dtype[np.float32]],
|
|
broadcast_dims: Mapping[_Dim, int],
|
|
expected_ndim: int,
|
|
) -> None:
|
|
expand_dims = set(broadcast_dims.keys()) - set(target.dims)
|
|
# loop over expand_dims and call .expand_dims(dim=dim) in a loop
|
|
for dim in expand_dims:
|
|
target = target.expand_dims(dim=dim)
|
|
result = target.broadcast_to(broadcast_dims)
|
|
assert result.ndim == expected_ndim
|
|
assert result.sizes == broadcast_dims
|
|
|
|
def test_broadcast_to_errors(
|
|
self, target: NamedArray[Any, np.dtype[np.float32]]
|
|
) -> None:
|
|
with pytest.raises(
|
|
ValueError,
|
|
match=r"operands could not be broadcast together with remapped shapes",
|
|
):
|
|
target.broadcast_to({"x": 2, "y": 2})
|
|
|
|
with pytest.raises(ValueError, match=r"Cannot add new dimensions"):
|
|
target.broadcast_to({"x": 2, "y": 2, "z": 2})
|
|
|
|
def test_warn_on_repeated_dimension_names(self) -> None:
|
|
with pytest.warns(UserWarning, match="Duplicate dimension names"):
|
|
NamedArray(("x", "x"), np.arange(4).reshape(2, 2))
|
|
|
|
def test_aggregation(self) -> None:
|
|
x: NamedArray[Any, np.dtype[np.int64]]
|
|
x = NamedArray(("x", "y"), np.arange(4).reshape(2, 2))
|
|
|
|
result = x.sum()
|
|
assert isinstance(result.data, np.ndarray)
|
|
|
|
|
|
def test_repr() -> None:
|
|
x: NamedArray[Any, np.dtype[np.uint64]]
|
|
x = NamedArray(("x",), np.array([0], dtype=np.uint64))
|
|
|
|
# Reprs should not crash:
|
|
r = x.__repr__()
|
|
x._repr_html_()
|
|
|
|
# Basic comparison:
|
|
assert r == "<xarray.NamedArray (x: 1)> Size: 8B\narray([0], dtype=uint64)"
|