ENH: new functions isclose and allclose

Author: crusaderky

docs/api-reference.md

@@ -6,11 +6,13 @@
     :nosignatures:
     :toctree: generated
 
+    allclose
     at
     atleast_nd
     cov
     create_diagonal
     expand_dims
+    isclose
     kron
     nunique
     pad

pixi.lock

@@ -293,6 +293,7 @@ checks = [
     "all",  # report on all checks, except the below
     "EX01", # most docstrings do not need an example
     "SA01", # data-apis/array-api-extra#87
+    "SA04", # Missing description for See Also cross-reference
     "ES01", # most docstrings do not need an extended summary
 ]
 exclude = [ # don't report on objects that match any of these regex

pyproject.toml

@@ -1,6 +1,6 @@
 """Extra array functions built on top of the array API standard."""
 
-from ._delegation import pad
+from ._delegation import allclose, isclose, pad
 from ._lib._at import at
 from ._lib._funcs import (
     atleast_nd,
@@ -18,11 +18,13 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "allclose",
     "at",
     "atleast_nd",
     "cov",
     "create_diagonal",
     "expand_dims",
+    "isclose",
     "kron",
     "nunique",
     "pad",

src/array_api_extra/__init__.py

@@ -7,7 +7,7 @@
 from ._lib._utils._compat import array_namespace
 from ._lib._utils._typing import Array
 
-__all__ = ["pad"]
+__all__ = ["allclose", "isclose", "pad"]
 
 
 def _delegate(xp: ModuleType, *backends: Backend) -> bool:
@@ -29,6 +29,144 @@ def _delegate(xp: ModuleType, *backends: Backend) -> bool:
     return any(backend.is_namespace(xp) for backend in backends)
 
 
+def allclose(
+    a: Array,
+    b: Array,
+    *,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Return True if two arrays are element-wise equal within a tolerance.
+
+    This is a simple convenience reduction around `isclose`.
+
+    Parameters
+    ----------
+    a, b : Array
+        Input arrays to compare.
+    rtol : array_like, optional
+        The relative tolerance parameter (see Notes).
+    atol : array_like, optional
+        The absolute tolerance parameter (see Notes).
+    equal_nan : bool, optional
+        Whether to compare NaN's as equal. If True, NaN's in `a` will be considered
+        equal to NaN's in `b` in the output array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a` and `b`. Default: infer.
+
+    Returns
+    -------
+    Array
+        A 0-dimensional boolean array, containing `True` if `a` is close to `b`, and
+        `False` otherwise.
+
+    See Also
+    --------
+    isclose
+    math.isclose
+
+    Notes
+    -----
+    If `xp` is a lazy backend (e.g. Dask, JAX), you may not be able to test the result
+    contents with `bool(allclose(a, b))` or `if allclose(a, b): ...`.
+    """
+    xp = array_namespace(a, b) if xp is None else xp
+    return xp.all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp))
+
+
+def isclose(
+    a: Array,
+    b: Array,
+    *,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Return a boolean array where two arrays are element-wise equal within a tolerance.
+
+    The tolerance values are positive, typically very small numbers. The relative
+    difference (rtol * abs(b)) and the absolute difference atol are added together to
+    compare against the absolute difference between a and b.
+
+    NaNs are treated as equal if they are in the same place and if equal_nan=True. Infs
+    are treated as equal if they are in the same place and of the same sign in both
+    arrays.
+
+    Parameters
+    ----------
+    a, b : Array
+        Input arrays to compare.
+    rtol : array_like, optional
+        The relative tolerance parameter (see Notes).
+    atol : array_like, optional
+        The absolute tolerance parameter (see Notes).
+    equal_nan : bool, optional
+        Whether to compare NaN's as equal. If True, NaN's in `a` will be considered
+        equal to NaN's in `b` in the output array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a` and `b`. Default: infer.
+
+    Returns
+    -------
+    Array
+        A boolean array of shape broadcasted from `a` and `b`, containing `True` where
+        ``a`` is close to ``b``, and `False` otherwise.
+
+    Warnings
+    --------
+    The default atol is not appropriate for comparing numbers with magnitudes much
+    smaller than one ) (see notes).
+
+    See Also
+    --------
+    allclose
+    math.isclose
+
+    Notes
+    -----
+    For finite values, `isclose` uses the following equation to test whether two
+    floating point values are equivalent::
+
+        absolute(a - b) <= (atol + rtol * absolute(b))
+
+    Unlike the built-in `math.isclose`, the above equation is not symmetric in a and b,
+    so that `isclose(a, b)` might be different from `isclose(b, a)` in some rare
+    cases.
+
+    The default value of `atol` is not appropriate when the reference value `b` has
+    magnitude smaller than one. For example, it is unlikely that `a = 1e-9` and
+    `b = 2e-9` should be considered "close", yet `isclose(1e-9, 2e-9)` is `True` with
+    default settings. Be sure to select atol for the use case at hand, especially for
+    defining the threshold below which a non-zero value in `a` will be considered
+    "close" to a very small or zero value in `b`.
+
+    The comparison of `a` and `b` uses standard broadcasting, which means that `a` and
+    `b` need not have the same shape in order for `isclose(a, b)` to evaluate to
+    `True`.
+
+    `isclose` is not defined for non-numeric data types. `bool` is considered a numeric
+    data-type for this purpose.
+    """
+    xp = array_namespace(a, b) if xp is None else xp
+
+    if _delegate(
+        xp,
+        Backend.NUMPY,
+        Backend.CUPY,
+        Backend.DASK,
+        Backend.JAX,
+        Backend.TORCH,
+    ):
+        return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+    return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+
 def pad(
     x: Array,
     pad_width: int | tuple[int, int] | list[tuple[int, int]],

src/array_api_extra/_delegation.py

@@ -304,6 +304,36 @@ def expand_dims(
     return a
 
 
+def isclose(
+    a: Array,
+    b: Array,
+    *,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+    xp: ModuleType | None = None,
+) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    xp = array_namespace(a, b) if xp is None else xp
+
+    a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating"))
+    b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating"))
+    if a_inexact or b_inexact:
+        # FIXME: use scipy's lazywhere to suppress warnings on inf
+        out = xp.abs(a - b) <= (atol + rtol * xp.abs(b))
+        out = xp.where(xp.isinf(a) & xp.isinf(b), xp.sign(a) == xp.sign(b), out)
+        if equal_nan:
+            out = xp.where(xp.isnan(a) & xp.isnan(b), xp.asarray(True), out)
+        return out
+
+    # integer types
+    atol = int(atol)
+    if rtol == 0:
+        return xp.abs(a - b) <= atol
+    nrtol = int(1.0 / rtol)
+    return xp.abs(a - b) <= (atol + xp.abs(b) // nrtol)
+
+
 def kron(a: Array, b: Array, /, *, xp: ModuleType | None = None) -> Array:
     """
     Kronecker product of two arrays.

src/array_api_extra/_lib/_funcs.py

@@ -6,11 +6,13 @@
 import pytest
 
 from array_api_extra import (
+    allclose,
     at,
     atleast_nd,
     cov,
     create_diagonal,
     expand_dims,
+    isclose,
     kron,
     nunique,
     pad,
@@ -23,7 +25,7 @@
 from array_api_extra._lib._utils._typing import Array, Device
 
 # some xp backends are untyped
-# mypy: disable-error-code=no-untyped-usage
+# mypy: disable-error-code=no-untyped-def
 
 
 @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no expand_dims")
@@ -252,6 +254,116 @@ def test_xp(self, xp: ModuleType):
         assert y.shape == (1, 1, 1, 3)
 
 
+@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no isdtype")
+class TestIsClose:
+    # FIXME use lazywhere to avoid warnings on inf
+    @pytest.mark.filterwarnings("ignore:invalid value encountered")
+    @pytest.mark.parametrize(
+        ("a", "b"),
+        [
+            (0.0, 0.0),
+            (1.0, 1.0),
+            (1.0, 2.0),
+            (1.0, -1.0),
+            (100.0, 101.0),
+            (0, 0),
+            (1, 1),
+            (1, 2),
+            (1, -1),
+            (1.0 + 1j, 1.0 + 1j),
+            (1.0 + 1j, 1.0 - 1j),
+            (float("inf"), float("inf")),
+            (float("inf"), 100.0),
+            (float("inf"), float("-inf")),
+            (float("nan"), float("nan")),
+            (float("nan"), 0.0),
+            (0.0, float("nan")),
+            (1e6, 1e6 + 1),  # True - within rtol
+            (1e6, 1e6 + 100),  # False - outside rtol
+            (1e-6, 1.1e-6),  # False - outside atol
+            (1e-7, 1.1e-7),  # True - outside atol
+            (1e6 + 0j, 1e6 + 1j),  # True - within rtol
+            (1e6 + 0j, 1e6 + 100j),  # False - outside rtol
+        ],
+    )
+    def test_basic(self, a: float, b: float, xp: ModuleType):
+        a_xp = xp.asarray(a)
+        b_xp = xp.asarray(b)
+
+        xp_assert_equal(isclose(a_xp, b_xp), xp.asarray(np.isclose(a, b)))
+        xp_assert_equal(allclose(a_xp, b_xp), xp.asarray(np.allclose(a, b)))
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            r_xp = xp.asarray(np.arange(10), dtype=a_xp.dtype)
+            ar_xp = a_xp * r_xp
+            br_xp = b_xp * r_xp
+            ar_np = a * np.arange(10)
+            br_np = b * np.arange(10)
+
+        xp_assert_equal(isclose(ar_xp, br_xp), xp.asarray(np.isclose(ar_np, br_np)))
+
+    @pytest.mark.parametrize("dtype", ["float32", "int32"])
+    def test_broadcast(self, dtype: str, xp: ModuleType):
+        dtype = getattr(xp, dtype)
+        a = xp.asarray([1, 2, 3], dtype=dtype)
+        b = xp.asarray([[1], [5]], dtype=dtype)
+        actual = isclose(a, b)
+        expect = xp.asarray(
+            [[True, False, False], [False, False, False]], dtype=xp.bool
+        )
+
+        xp_assert_equal(actual, expect)
+
+    # FIXME use lazywhere to avoid warnings on inf
+    @pytest.mark.filterwarnings("ignore:invalid value encountered")
+    def test_some_inf(self, xp: ModuleType):
+        a = xp.asarray([0.0, 1.0, float("inf"), float("inf"), float("inf")])
+        b = xp.asarray([1e-9, 1.0, float("inf"), float("-inf"), 2.0])
+        actual = isclose(a, b)
+        xp_assert_equal(actual, xp.asarray([True, True, True, False, False]))
+
+    def test_equal_nan(self, xp: ModuleType):
+        a = xp.asarray([float("nan"), float("nan"), 1.0])
+        b = xp.asarray([float("nan"), 1.0, float("nan")])
+        xp_assert_equal(isclose(a, b), xp.asarray([False, False, False]))
+        xp_assert_equal(isclose(a, b, equal_nan=True), xp.asarray([True, False, False]))
+        xp_assert_equal(allclose(a[:1], b[:1]), xp.asarray(False))
+        xp_assert_equal(allclose(a[:1], b[:1], equal_nan=True), xp.asarray(True))
+
+    @pytest.mark.parametrize("dtype", ["float32", "complex64", "int32"])
+    def test_tolerance(self, dtype: str, xp: ModuleType):
+        dtype = getattr(xp, dtype)
+        a = xp.asarray([100, 100], dtype=dtype)
+        b = xp.asarray([101, 102], dtype=dtype)
+        xp_assert_equal(isclose(a, b), xp.asarray([False, False]))
+        xp_assert_equal(isclose(a, b, atol=1), xp.asarray([True, False]))
+        xp_assert_equal(isclose(a, b, rtol=0.01), xp.asarray([True, False]))
+        xp_assert_equal(allclose(a[:1], b[:1]), xp.asarray(False))
+        xp_assert_equal(allclose(a[:1], b[:1], atol=1), xp.asarray(True))
+        xp_assert_equal(allclose(a[:1], b[:1], rtol=0.01), xp.asarray(True))
+
+        # Attempt to trigger division by 0 in rtol on int dtype
+        xp_assert_equal(isclose(a, b, rtol=0), xp.asarray([False, False]))
+        xp_assert_equal(isclose(a, b, atol=1, rtol=0), xp.asarray([True, False]))
+        xp_assert_equal(allclose(a[:1], b[:1], rtol=0), xp.asarray(False))
+        xp_assert_equal(allclose(a[:1], b[:1], atol=1, rtol=0), xp.asarray(True))
+
+    def test_very_small_numbers(self, xp: ModuleType):
+        a = xp.asarray([1e-9, 1e-9])
+        b = xp.asarray([1.0001e-9, 1.00001e-9])
+        # Difference is below default atol
+        xp_assert_equal(isclose(a, b), xp.asarray([True, True]))
+        # Use only rtol
+        xp_assert_equal(isclose(a, b, atol=0), xp.asarray([False, True]))
+        xp_assert_equal(isclose(a, b, atol=0, rtol=0), xp.asarray([False, False]))
+
+    def test_xp(self, xp: ModuleType):
+        a = xp.asarray([0.0, 0.0])
+        b = xp.asarray([1e-9, 1e-4])
+        xp_assert_equal(isclose(a, b, xp=xp), xp.asarray([True, False]))
+
+
 @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no expand_dims")
 class TestKron:
     def test_basic(self, xp: ModuleType):