diff --git a/docs/sphinx/api-manipulation.md b/docs/sphinx/api-manipulation.md index b8875fd5..a778786b 100644 --- a/docs/sphinx/api-manipulation.md +++ b/docs/sphinx/api-manipulation.md @@ -10,4 +10,5 @@ broadcast_shapes expand_dims pad + replace ``` diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index c4e5d42a..d108ca4d 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -9,7 +9,7 @@ from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index from ._lazy import lazy_apply from ._linalg import kron -from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad +from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad, replace from ._searching import searchsorted from ._set import isin, nunique, setdiff1d, union1d from ._sorting import argpartition, partition @@ -45,6 +45,7 @@ "pad", "partition", "rad2deg", + "replace", "searchsorted", "setdiff1d", "sinc", diff --git a/src/array_api_extra/_agnostic/_manipulation.py b/src/array_api_extra/_agnostic/_manipulation.py index fb28faac..0d2770ef 100644 --- a/src/array_api_extra/_agnostic/_manipulation.py +++ b/src/array_api_extra/_agnostic/_manipulation.py @@ -1,5 +1,6 @@ """Array-agnostic implementations for manipulation functions.""" +import cmath import math import typing from collections.abc import Sequence @@ -8,7 +9,7 @@ from .._lib import _compat, _helpers from .._lib._typing import Array, ArrayNamespace -__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad"] +__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad", "replace"] def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace) -> Array: @@ -97,3 +98,11 @@ def pad( device=_compat.device(x), ) return at(padded, tuple(slices)).set(x) + + +def replace( # numpydoc ignore=PR01,RT01 + x: Array, old: complex, new: complex, /, *, xp: ArrayNamespace +) -> Array: + """See docstring in `array_api_extra._manipulation`.""" + mask = xp.isnan(x) if cmath.isnan(old) else x == old + return xp.where(mask, new, x) diff --git a/src/array_api_extra/_manipulation.py b/src/array_api_extra/_manipulation.py index fa3eae9b..9bd2c071 100644 --- a/src/array_api_extra/_manipulation.py +++ b/src/array_api_extra/_manipulation.py @@ -7,7 +7,7 @@ from ._lib import _compat, _helpers from ._lib._typing import Array, ArrayNamespace -__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad"] +__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad", "replace"] def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace | None = None) -> Array: @@ -282,3 +282,72 @@ def pad( return _agnostic._manipulation.pad( x, pad_width, constant_values=constant_values, xp=xp ) + + +def replace( + x: Array, + old: complex, + new: complex, + /, + *, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Replace each element of an array equal to a scalar with another scalar. + + Parameters + ---------- + x : array + Input array. + old : python scalar + Value to replace. If NaN, every NaN element of `x` is replaced. + new : python scalar + Replacement value. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + An array with the same shape and data type as `x`, in which each element + equal to `old` is replaced by `new` and all other elements are equal to the + corresponding elements of `x`. `x` is not modified. + + Notes + ----- + `old` and `new` should have types and values compatible with the data type of + `x` (see :ref:`array-api:mixing-scalars-and-arrays`). Otherwise, the behaviour + depends on the backend. + + For sparse arrays, elements which are not explicitly stored are replaced too: + if `old` equals the fill value, the fill value of the result is `new`. + The result keeps the sparse format of `x` and is never densified. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([1, 2, 1, 3]) + >>> xpx.replace(x, 1, 0) + Array([0, 2, 0, 3], dtype=array_api_strict.int64) + + NaN elements are replaced when `old` is NaN: + + >>> x = xp.asarray([1.0, xp.nan, 3.0]) + >>> xpx.replace(x, xp.nan, 0.0) + Array([1., 0., 3.], dtype=array_api_strict.float64) + """ + xp = _compat.array_namespace(x) if xp is None else xp + + if _compat.is_pydata_sparse_namespace(xp): + # `elemwise` applies the replacement to both the stored elements and the + # fill value, so implicit elements are replaced without densifying + def replace_dense(data: Array) -> Array: # numpydoc ignore=PR01,RT01 + """Replace in the dense arrays underlying a sparse array.""" + return _agnostic._manipulation.replace( + data, old, new, xp=_compat.array_namespace(data) + ) + + return xp.elemwise(replace_dense, x) + + return _agnostic._manipulation.replace(x, old, new, xp=xp) diff --git a/tests/main/test_manipulation.py b/tests/main/test_manipulation.py index f6e06d04..81c99a11 100644 --- a/tests/main/test_manipulation.py +++ b/tests/main/test_manipulation.py @@ -1,9 +1,10 @@ import math +from typing import Any, cast import numpy as np import pytest -from array_api_extra import atleast_nd, broadcast_shapes, expand_dims, pad +from array_api_extra import atleast_nd, broadcast_shapes, expand_dims, pad, replace from array_api_extra._lib._backends import Backend from array_api_extra._lib._compat import device as get_device from array_api_extra._lib._typing import ArrayNamespace, Device @@ -13,6 +14,7 @@ lazy_xp_function(broadcast_shapes) lazy_xp_function(expand_dims) lazy_xp_function(pad) +lazy_xp_function(replace) class TestAtLeastND: @@ -371,3 +373,86 @@ def test_sequence_of_tuples_width(self, xp: ArrayNamespace): assert padded.shape == (4, 6) padded = pad(a, ((1, 0), (0, 0))) assert padded.shape == (4, 4) + + +class TestReplace: + def test_simple(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 1, 3]) + assert_equal(replace(x, 1, 5), xp.asarray([5, 2, 5, 3])) + + def test_no_match(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3]) + assert_equal(replace(x, 4, 5), x) + + def test_zero(self, xp: ArrayNamespace): + # Zero is the fill value of sparse arrays + x = xp.asarray([0, 1, 0, 2]) + assert_equal(replace(x, 0, 5), xp.asarray([5, 1, 5, 2])) + assert_equal(replace(x, 1, 0), xp.asarray([0, 0, 0, 2])) + + def test_nan(self, xp: ArrayNamespace): + x = xp.asarray([1.0, xp.nan, 2.0, xp.nan]) + assert_equal(replace(x, xp.nan, 0.0), xp.asarray([1.0, 0.0, 2.0, 0.0])) + assert_equal(replace(x, 1.0, xp.nan), xp.asarray([xp.nan, xp.nan, 2.0, xp.nan])) + + def test_complex(self, xp: ArrayNamespace): + x = xp.asarray([1 + 1j, 2 + 0j, complex(xp.nan, 1)]) + expect = xp.asarray([0j, 2 + 0j, complex(xp.nan, 1)]) + assert_equal(replace(x, 1 + 1j, 0j), expect) + # Complex elements with a NaN component are NaN + expect = xp.asarray([1 + 1j, 2 + 0j, 3j]) + assert_equal(replace(x, complex(xp.nan, 0), 3j), expect) + + def test_bool(self, xp: ArrayNamespace): + x = xp.asarray([True, False, True]) + assert_equal(replace(x, True, False), xp.asarray([False, False, False])) + + @pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "float32", "float64"]) + def test_dtype(self, xp: ArrayNamespace, dtype: str): + x = xp.asarray([1, 2, 1], dtype=getattr(xp, dtype)) + expect = xp.asarray([3, 2, 3], dtype=getattr(xp, dtype)) + assert_equal(replace(x, 1, 3), expect) + + def test_0d(self, xp: ArrayNamespace): + assert_equal(replace(xp.asarray(1.0), 1.0, 2.0), xp.asarray(2.0)) + assert_equal(replace(xp.asarray(1.0), 3.0, 2.0), xp.asarray(1.0)) + + def test_ndim(self, xp: ArrayNamespace): + x = xp.asarray(np.reshape(np.arange(2 * 3 * 4) % 3, (2, 3, 4))) + expect = xp.asarray(np.reshape(np.arange(2 * 3 * 4) % 3 * 4 % 5, (2, 3, 4))) + actual = replace(x, 1, 4) + assert_equal(replace(actual, 2, 3), expect) + + def test_input_not_modified(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 1]) + _ = replace(x, 1, 5) + assert_equal(x, xp.asarray([1, 2, 1])) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([1.0, 2.0], device=device) + assert get_device(replace(x, 1.0, 3.0)) == device + + def test_xp(self, xp: ArrayNamespace): + x = xp.asarray([1, 2]) + assert_equal(replace(x, 1, 3, xp=xp), xp.asarray([3, 2])) + + @pytest.mark.parametrize("fmt", ["coo", "gcxs", "dok"]) + def test_sparse(self, fmt: str): + sparse = pytest.importorskip("sparse") + x: Any = sparse.asarray(np.asarray([0.0, 1.0, 0.0, np.nan]), format=fmt) + + # Replacing the fill value changes the fill value, without densifying + actual = cast(Any, replace(x, 0.0, 5.0)) + assert type(actual) is type(x) + assert actual.fill_value == 5.0 + assert actual.nnz == x.nnz + np.testing.assert_array_equal(actual.todense(), [5.0, 1.0, 5.0, np.nan]) + + actual = cast(Any, replace(x, np.nan, 2.0)) + assert type(actual) is type(x) + assert actual.fill_value == 0.0 + np.testing.assert_array_equal(actual.todense(), [0.0, 1.0, 0.0, 2.0]) + + # Input is not modified + assert x.fill_value == 0.0 + np.testing.assert_array_equal(x.todense(), [0.0, 1.0, 0.0, np.nan])