Repository: numpy/numpy
Pull request: ENH: isclose/allclose: support array_like `atol`/`rtol`
Supersedes gh-14343
Closes gh-14320

gh-14320 reported that `np.allclose` accepted vector `atol` but failed when a NaN was present.
gh-14343 began to add full support for array `atol`/`rtol` in `allclose`/`isclose`, but became inactive.
@seberg [asked](https://github.com/numpy/numpy/pull/14343#issuecomment-1551072477) if anyone would like to finish it, hence this PR. 

I preserved the original commits and made the requested changes. I also fixed the case of non-array array_like `rtol`/`atol` and strengthened the tests.

Problem statement:
allclose() vector `atol` parameter wrong shape if NaN in a or b.
<!-- Please describe the issue in detail here, and fill in the fields below -->

```python
>>>np.allclose([1,2,float('nan')], [1,2,float('nan')], equal_nan=True)
True
>>>np.allclose([1,2,float('nan')], [1,2,float('nan')], atol=[1,1], equal_nan=True)
True
>>>np.allclose([1,2,3], [1,2,3)], atol=[1,1,1], equal_nan=True)
True
>>>np.allclose([1,2,float('nan')], [1,2,float('nan')], atol=[1,1,1], equal_nan=True)
ValueError: operands could not be broadcast together with shapes (3,) (2,)
```

It appears that `allclose` is basically removing entries where `a` or `b` is NaN before passing them on to `isclose` with the same `atol` provided to `allclose`. This of course doesn't work if I expect there to be no NaNs some of the time and need to provide a longer `atol` because my `a` and `b` are different scales in different positions.

What should happen is `allclose` removes the same positions from `atol` before passing it on, so the dimensions always agree.

Hints from issue discussion:
I'm not sure support for a vector in atol was intended, but adding support would be fine by me
> I'm not sure support for a vector in atol was intended, but adding support would be fine by me

The unintended support for vectors came from broadcasting inside [`within_tol`](https://github.com/numpy/numpy/blob/ccfbcc1cd9a4035a467f2e982a565ab27de25b6b/numpy/core/numeric.py#L2352) in `np.isclose()`. 
 
As a support for vector `atol` and `rtol`, will the following addition below the [line](https://github.com/numpy/numpy/blob/ccfbcc1cd9a4035a467f2e982a565ab27de25b6b/numpy/core/numeric.py#L2376) be correct @eric-wieser?
```
atol = full(x.shape, atol)
rtol = full(x.shape, rtol)
```
These two lines will raise `ValueError: could not broadcast input array from shape (p) into shape (q)` if `atol` or `rtol` is not a scaler and cannot also be casted to the same shape as `x`.  If desired,  the error can be caught and re-raised with a more informative message like `'atol must be a scalar or of shape %s' %x.shape`.

The line 
```
 cond[finite] = within_tol(x[finite], y[finite], atol, rtol)
```
 can now then be written as 
```
 cond[finite] = within_tol(x[finite], y[finite], atol[finite], rtol[finite])
```

If `atol` and `rtol` must strictly be scalars, then I will suggest that a check is added to make that more explicit. 
I am not sure I understand the use case for non-scalar `rtol`/`atol`. I closed #14343, and I think a PR to make sure they are scalars would be more helpful since that is what is in the docs, with perhaps a deprecation cycle.
> I think a PR to make sure they are scalars would be more helpful since that is what is in the docs, with perhaps a deprecation cycle.

Any news on this? Honestly I also believe that this behavior is very inconsistent with the docs, and somehow in version 1.24.0, `np.allclose` allows passing in `atol` as a vector, while `np.testing.assert_allclose` only raises error because of a print command...

```shell-session
1585     return np.core.numeric.isclose(x, y, rtol=rtol, atol=atol,
   1586                                    equal_nan=equal_nan)
   1588 actual, desired = np.asanyarray(actual), np.asanyarray(desired)
-> 1589 header = f'Not equal to tolerance rtol={rtol:g}, atol={atol:g}'
   1590 assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
   1591                      verbose=verbose, header=header, equal_nan=equal_nan)

TypeError: unsupported format string passed to list.__format__
```

> I am not sure I understand the use case for non-scalar `rtol`/`atol`. I closed #14343, and 

Meanwhile, what is your recommended way to compare two 3D vectors `x` and `y` where the first two axis should have absolute tolerance of 0.1 while the third axis should have absolute tolerance of 1? I believe this is the use case that many of us run into, and it is quite natural to write 

```python3
import numpy as np
x = [1, 2, 3]
y = [1.1, 2.1, 5]
np.allclose(x, y, atol=[0.1, 0.1, 1])
```
A nonscalar `atol` makes sense if you want to ensure position 1 across two vectors is within `a` of each other, position 2 is within `b` of each other, etc., in only a single call.

Commit message: MAINT: isclose/allclose: simplify support for array-like tolerances

Patch to review:
diff --git a/doc/release/upcoming_changes/24878.improvement.rst b/doc/release/upcoming_changes/24878.improvement.rst
new file mode 100644
index 000000000000..68cb0d7a1534
--- /dev/null
+++ b/doc/release/upcoming_changes/24878.improvement.rst
@@ -0,0 +1,5 @@
+Support for array ``atol`` and ``rtol`` in ``np.isclose`` and ``np.allclose``
+-----------------------------------------------------------------------------
+The keywords ``atol`` and ``rtol`` in ``np.isclose`` and ``np.allclose``
+now accept both scalars and arrays. An array, if given, must broadcast
+to the shapes of the arguments `a` and `b`.
\ No newline at end of file
diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py
index 50229151d6a4..b1ddffd671de 100644
--- a/numpy/_core/numeric.py
+++ b/numpy/_core/numeric.py
@@ -2174,7 +2174,7 @@ def identity(n, dtype=None, *, like=None):
 
 
 def _allclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None):
-    return (a, b)
+    return (a, b, rtol, atol)
 
 
 @array_function_dispatch(_allclose_dispatcher)
@@ -2195,9 +2195,9 @@ def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
     ----------
     a, b : array_like
         Input arrays to compare.
-    rtol : float
+    rtol : array_like
         The relative tolerance parameter (see Notes).
-    atol : float
+    atol : array_like
         The absolute tolerance parameter (see Notes).
     equal_nan : bool
         Whether to compare NaN's as equal.  If True, NaN's in `a` will be
@@ -2253,7 +2253,7 @@ def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
 
 
 def _isclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None):
-    return (a, b)
+    return (a, b, rtol, atol)
 
 
 @array_function_dispatch(_isclose_dispatcher)
@@ -2274,9 +2274,9 @@ def isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
     ----------
     a, b : array_like
         Input arrays to compare.
-    rtol : float
+    rtol : array_like
         The relative tolerance parameter (see Notes).
-    atol : float
+    atol : array_like
         The absolute tolerance parameter (see Notes).
     equal_nan : bool
         Whether to compare NaN's as equal.  If True, NaN's in `a` will be
@@ -2341,8 +2341,7 @@ def within_tol(x, y, atol, rtol):
         with errstate(invalid='ignore'), _no_nep50_warning():
             return less_equal(abs(x-y), atol + rtol * abs(y))
 
-    x = asanyarray(a)
-    y = asanyarray(b)
+    x, y, atol, rtol = np.broadcast_arrays(a, b, atol, rtol, subok=True)
 
     # Make sure y is an inexact type to avoid bad behavior on abs(MIN_INT).
     # This will cause casting of x later. Also, make sure to allow subclasses
@@ -2362,13 +2361,9 @@ def within_tol(x, y, atol, rtol):
     else:
         finite = xfin & yfin
         cond = zeros_like(finite, subok=True)
-        # Because we're using boolean indexing, x & y must be the same shape.
-        # Ideally, we'd just do x, y = broadcast_arrays(x, y). It's in
-        # lib.stride_tricks, though, so we can't import it here.
-        x = x * ones_like(cond)
-        y = y * ones_like(cond)
         # Avoid subtraction with infinite/nan values...
-        cond[finite] = within_tol(x[finite], y[finite], atol, rtol)
+        cond[finite] = within_tol(x[finite], y[finite],
+                                  atol[finite], rtol[finite])
         # Check for equality of infinite values...
         cond[~finite] = (x[~finite] == y[~finite])
         if equal_nan:
diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py
index 63c0e853b8c2..b40927f4db93 100644
--- a/numpy/_core/tests/test_numeric.py
+++ b/numpy/_core/tests/test_numeric.py
@@ -2926,6 +2926,22 @@ def test_ip_isclose(self):
         for (x, y), result in zip(tests, results):
             assert_array_equal(np.isclose(x, y), result)
 
+        x = np.array([2.1, 2.1, 2.1, 2.1, 5, np.nan])
+        y = np.array([2, 2, 2, 2, np.nan, 5])
+        atol = [0.11, 0.09, 1e-8, 1e-8, 1, 1]
+        rtol = [1e-8, 1e-8, 0.06, 0.04, 1, 1]
+        expected = np.array([True, False, True, False, False, False])
+        assert_array_equal(np.isclose(x, y, rtol=rtol, atol=atol), expected)
+
+        message = "operands could not be broadcast together..."
+        atol = np.array([1e-8, 1e-8])
+        with assert_raises(ValueError, msg=message):
+            np.isclose(x, y, atol=atol)
+
+        rtol = np.array([1e-5, 1e-5])
+        with assert_raises(ValueError, msg=message):
+            np.isclose(x, y, rtol=rtol)
+
     def tst_all_isclose(self, x, y):
         assert_(np.all(np.isclose(x, y)), "%s and %s not close" % (x, y))
 
@@ -2946,6 +2962,17 @@ def test_ip_all_isclose(self):
         for (x, y) in self.all_close_tests:
             self.tst_all_isclose(x, y)
 
+        x = np.array([2.3, 3.6, 4.4, np.nan])
+        y = np.array([2, 3, 4, np.nan])
+        atol = [0.31, 0, 0, 1]
+        rtol = [0, 0.21, 0.11, 1]
+        assert np.allclose(x, y, atol=atol, rtol=rtol, equal_nan=True)
+        assert not np.allclose(x, y, atol=0.1, rtol=0.1, equal_nan=True)
+
+        # Show that gh-14330 is resolved
+        assert np.allclose([1, 2, float('nan')], [1, 2, float('nan')],
+                           atol=[1, 1, 1], equal_nan=True)
+
     def test_ip_none_isclose(self):
         self._setup()
         for (x, y) in self.none_close_tests:


Code review comments:
diff_hunk: @@ -2341,8 +2341,7 @@ def within_tol(x, y, atol, rtol):
         with errstate(invalid='ignore'), _no_nep50_warning():
             return less_equal(abs(x-y), atol + rtol * abs(y))
 
-    x = asanyarray(a)
-    y = asanyarray(b)
+    x, y, atol, rtol = np.broadcast_arrays(a, b, atol, rtol, subok=True)
path: numpy/_core/numeric.py
text: @user1:
Hmmmpf, this is one of the cases similar to gh-24957.  So, `atol` and `rtol` stop being Python scalars (with NEP 50) or 0-d objects (without NEP 50) and that changes promotion. 

We need a solution for that, right now, we could skip it add logic for them being `type(atol) in (int, float, complex)`.  The question remains a bit whether we want a way to say `mark_was_weak_scalar(array)`, to turn on weak promotion for an *array*, but... I don't really want such a mark to survive for long and it is unclear what operations should retain it.

@user2:
I would think that if NumPy is hitting this, then other libraries and users will hit it as well. Maybe we should rethink the transition and how we provide (or not) backward incompatibility shims.

@user3:
I guess the casting could be done earlier - something like 
```
dtype=np.result_type(a, b, atol, rtol)
atol = np.asanyarray(atol, dtype)
rtol = np.asanyarray(rtol, dtype)
```
and then broadcast?

If broadcast is a common culprit, we might add a new keyword argument that does this implicitly?

@author:
> If broadcast is a common culprit, we might add a new keyword argument that does this implicitly?

~~Yes, please!~~ This is something Sebastian and I have discussed before.

Oops, but I've wanted the promotion to follow NEP50, so it would be nice if the user could choose which rules to follow.

@user3:
Yes, that would be the suggestion: to allow the dtype of arrays created from python scalars to follow NEP 50 promotion rules.

@user1:
There are different approaches:
* just make the current stuff a bit more convenient.
* Try to tag on the info to an array.
* Just deal with it here, the `isinstance` check and special code paths are not that terrible (and also useful for micro optimization probably).

To me no solution is obvious, unfortunately.  Tagging the info on to the array or dtype, means that you have a very tricky situation of how to drag the info around and mainly about when to preserve and when to drop it.
In this case for example, you are calling `arr[boolean_idx]` creating a copy, should that drop to prevent accidentally returning such an array?  Assignments at least could be prevent by making it readonly, but...  Also the special path to make `arr > -1` work for such an array is another tricky thing.
I can somewhat think of ways to do this, but I am not sure they will be very pleasant.

Moving resolution to earlier (maybe with a helper) can help.  But in some cases might also not be quite what we want because it makes sense to ask `isclose(2 meters, 1.9 meters, atol=0.1, rtol=0.1 meters)`.  But it is less clear that the: `dtype = result_type(meters, float)` is reasonable (for the `atol`)?