scipy 1.15.3__cp312-cp312-macosx_14_0_arm64.whl → 1.16.0rc2__cp312-cp312-macosx_14_0_arm64.whl

scipy/cluster/tests/test_vq.py

@@ -1,27 +1,25 @@
-import warnings
+import math
 import sys
 from copy import deepcopy
 from threading import Lock
 
 import numpy as np
-from numpy.testing import (
-    assert_array_equal, assert_equal, assert_, suppress_warnings
-)
+from numpy.testing import assert_array_equal, suppress_warnings
 import pytest
 from pytest import raises as assert_raises
 
 from scipy.cluster.vq import (kmeans, kmeans2, py_vq, vq, whiten,
                               ClusterError, _krandinit)
 from scipy.cluster import _vq
-from scipy.conftest import array_api_compatible
 from scipy.sparse._sputils import matrix
 
 from scipy._lib import array_api_extra as xpx
 from scipy._lib._array_api import (
-    SCIPY_ARRAY_API, array_namespace, xp_copy, xp_assert_close, xp_assert_equal
+    SCIPY_ARRAY_API, eager_warns, is_lazy_array, make_xp_test_case,
+    xp_copy, xp_assert_close, xp_assert_equal
 )
 
-pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
+xfail_xp_backends = pytest.mark.xfail_xp_backends
 skip_xp_backends = pytest.mark.skip_xp_backends
 
 TESTDATA_2D = np.array([
@@ -80,6 +78,7 @@ CODET2 = np.array([[11.0/3, 8.0/3],
 LABEL1 = np.array([0, 1, 2, 2, 2, 2, 1, 2, 1, 1, 1])
 
 
+@make_xp_test_case(whiten)
 class TestWhiten:
 
     def test_whiten(self, xp):
@@ -96,13 +95,7 @@ class TestWhiten:
                           [0.45067590, 0.45464607]])
         xp_assert_close(whiten(obs), desired, rtol=1e-5)
 
-    @pytest.fixture
-    def whiten_lock(self):
-        return Lock()
-
-    @skip_xp_backends('jax.numpy',
-                      reason='jax arrays do not support item assignment')
-    def test_whiten_zero_std(self, xp, whiten_lock):
+    def test_whiten_zero_std(self, xp):
         desired = xp.asarray([[0., 1.0, 2.86666544],
                               [0., 1.0, 1.32460034],
                               [0., 1.0, 3.74382172]])
@@ -111,27 +104,32 @@ class TestWhiten:
                           [0., 1., 0.34243798],
                           [0., 1., 0.96785929]])
 
-        with whiten_lock:
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter('always')
-
-                xp_assert_close(whiten(obs), desired, rtol=1e-5)
+        with eager_warns(obs, RuntimeWarning, match="standard deviation zero"):
+            actual = whiten(obs)
+        xp_assert_close(actual, desired, rtol=1e-5)
 
-                assert_equal(len(w), 1)
-                assert_(issubclass(w[-1].category, RuntimeWarning))
+    @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask")
+    @pytest.mark.parametrize("bad_value", [math.nan, math.inf, -math.inf])
+    def test_whiten_not_finite(self, bad_value, xp):
+        obs = xp.asarray([[0.98744510, bad_value],
+                          [0.62093317, 0.19406729],
+                          [0.87545741, 0.00735733],
+                          [0.85124403, 0.26499712],
+                          [0.45067590, 0.45464607]])
 
-    def test_whiten_not_finite(self, xp):
-        for bad_value in xp.nan, xp.inf, -xp.inf:
-            obs = xp.asarray([[0.98744510, bad_value],
-                              [0.62093317, 0.19406729],
-                              [0.87545741, 0.00735733],
-                              [0.85124403, 0.26499712],
-                              [0.45067590, 0.45464607]])
+        if is_lazy_array(obs):
+            desired = xp.asarray([[5.08738849, math.nan],
+                                  [3.19909255, math.nan],
+                                  [4.51041982, math.nan],
+                                  [4.38567074, math.nan],
+                                  [2.32191480, math.nan]])
+            xp_assert_close(whiten(obs), desired, rtol=1e-5)
+        else:
             assert_raises(ValueError, whiten, obs)
 
     @pytest.mark.skipif(SCIPY_ARRAY_API,
                         reason='`np.matrix` unsupported in array API mode')
-    def test_whiten_not_finite_matrix(self, xp):
+    def test_whiten_not_finite_matrix(self):
         for bad_value in np.nan, np.inf, -np.inf:
             obs = matrix([[0.98744510, bad_value],
                           [0.62093317, 0.19406729],
@@ -141,9 +139,9 @@ class TestWhiten:
             assert_raises(ValueError, whiten, obs)
 
 
+@make_xp_test_case(vq)
 class TestVq:
 
-    @skip_xp_backends(cpu_only=True)
     def test_py_vq(self, xp):
         initc = np.concatenate([[X[0]], [X[1]], [X[2]]])
         # label1.dtype varies between int32 and int64 over platforms
@@ -153,28 +151,26 @@ class TestVq:
 
     @pytest.mark.skipif(SCIPY_ARRAY_API,
                         reason='`np.matrix` unsupported in array API mode')
-    def test_py_vq_matrix(self, xp):
+    def test_py_vq_matrix(self):
         initc = np.concatenate([[X[0]], [X[1]], [X[2]]])
         # label1.dtype varies between int32 and int64 over platforms
         label1 = py_vq(matrix(X), matrix(initc))[0]
         assert_array_equal(label1, LABEL1)
 
-    @skip_xp_backends(np_only=True, reason='`_vq` only supports NumPy backend')
     def test_vq(self, xp):
         initc = np.concatenate([[X[0]], [X[1]], [X[2]]])
-        label1, _ = _vq.vq(xp.asarray(X), xp.asarray(initc))
+        label1, _ = _vq.vq(X, initc)
         assert_array_equal(label1, LABEL1)
         _, _ = vq(xp.asarray(X), xp.asarray(initc))
 
     @pytest.mark.skipif(SCIPY_ARRAY_API,
                         reason='`np.matrix` unsupported in array API mode')
-    def test_vq_matrix(self, xp):
+    def test_vq_matrix(self):
         initc = np.concatenate([[X[0]], [X[1]], [X[2]]])
         label1, _ = _vq.vq(matrix(X), matrix(initc))
         assert_array_equal(label1, LABEL1)
         _, _ = vq(matrix(X), matrix(initc))
 
-    @skip_xp_backends(cpu_only=True)
     def test_vq_1d(self, xp):
         # Test special rank 1 vq algo, python implementation.
         data = X[:, 0]
@@ -187,18 +183,15 @@ class TestVq:
         xp_assert_equal(ta, xp.asarray(a, dtype=xp.int64), check_dtype=False)
         xp_assert_equal(tb, xp.asarray(b))
 
-    @skip_xp_backends(np_only=True, reason='`_vq` only supports NumPy backend')
-    def test__vq_sametype(self, xp):
-        a = xp.asarray([1.0, 2.0], dtype=xp.float64)
-        b = a.astype(xp.float32)
+    def test__vq_sametype(self):
+        a = np.asarray([1.0, 2.0])
+        b = a.astype(np.float32)
         assert_raises(TypeError, _vq.vq, a, b)
 
-    @skip_xp_backends(np_only=True, reason='`_vq` only supports NumPy backend')
-    def test__vq_invalid_type(self, xp):
-        a = xp.asarray([1, 2], dtype=int)
+    def test__vq_invalid_type(self):
+        a = np.asarray([1, 2], dtype=int)
         assert_raises(TypeError, _vq.vq, a, a)
 
-    @skip_xp_backends(cpu_only=True)
     def test_vq_large_nfeat(self, xp):
         X = np.random.rand(20, 20)
         code_book = np.random.rand(3, 20)
@@ -222,7 +215,6 @@ class TestVq:
         # codes1.dtype varies between int32 and int64 over platforms
         xp_assert_equal(codes1, xp.asarray(codes0, dtype=xp.int64), check_dtype=False)
 
-    @skip_xp_backends(cpu_only=True)
     def test_vq_large_features(self, xp):
         X = np.random.rand(10, 5) * 1000000
         code_book = np.random.rand(2, 5) * 1000000
@@ -238,8 +230,8 @@ class TestVq:
 
 # Whole class skipped on GPU for now;
 # once pdist/cdist are hooked up for CuPy, more tests will work
-@skip_xp_backends(cpu_only=True)
-class TestKMean:
+@make_xp_test_case(kmeans, kmeans2)
+class TestKMeans:
 
     def test_large_features(self, xp):
         # Generate a data set with large values, and run kmeans on it to
@@ -267,7 +259,7 @@ class TestKMean:
 
     @pytest.mark.skipif(SCIPY_ARRAY_API,
                         reason='`np.matrix` unsupported in array API mode')
-    def test_kmeans_simple_matrix(self, xp):
+    def test_kmeans_simple_matrix(self):
         rng = np.random.default_rng(54321)
         initc = np.concatenate([[X[0]], [X[1]], [X[2]]])
         code1 = kmeans(matrix(X), matrix(initc), iter=1, rng=rng)[0]
@@ -302,9 +294,9 @@ class TestKMean:
 
     @pytest.mark.skipif(SCIPY_ARRAY_API,
                         reason='`np.matrix` unsupported in array API mode')
-    def test_kmeans2_simple_matrix(self, xp):
+    def test_kmeans2_simple_matrix(self):
         rng = np.random.default_rng(12345678)
-        initc = xp.asarray(np.concatenate([[X[0]], [X[1]], [X[2]]]))
+        initc = np.concatenate([[X[0]], [X[1]], [X[2]]])
         code1 = kmeans2(matrix(X), matrix(initc), iter=1, rng=rng)[0]
         code2 = kmeans2(matrix(X), matrix(initc), iter=2, rng=rng)[0]
 
@@ -334,8 +326,6 @@ class TestKMean:
         data = xp.reshape(data, (20, 20))[:10, :]
         kmeans2(data, 2)
 
-    @skip_xp_backends('jax.numpy',
-                      reason='jax arrays do not support item assignment')
     def test_kmeans2_init(self, xp):
         rng = np.random.default_rng(12345678)
         data = xp.asarray(TESTDATA_2D)
@@ -357,6 +347,7 @@ class TestKMean:
     def krand_lock(self):
         return Lock()
 
+    @xfail_xp_backends('dask.array', reason="Wrong answer")
     @pytest.mark.skipif(sys.platform == 'win32',
                         reason='Fails with MemoryError in Wine.')
     def test_krandinit(self, xp, krand_lock):
@@ -364,13 +355,12 @@ class TestKMean:
         datas = [xp.reshape(data, (200, 2)),
                  xp.reshape(data, (20, 20))[:10, :]]
         k = int(1e6)
-        xp_test = array_namespace(data)
         with krand_lock:
             for data in datas:
                 rng = np.random.default_rng(1234)
-                init = _krandinit(data, k, rng, xp_test)
-                orig_cov = xpx.cov(data.T, xp=xp_test)
-                init_cov = xpx.cov(init.T, xp=xp_test)
+                init = _krandinit(data, k, rng, xp)
+                orig_cov = xpx.cov(data.T, xp=xp)
+                init_cov = xpx.cov(init.T, xp=xp)
                 xp_assert_close(orig_cov, init_cov, atol=1.1e-2)
 
     def test_kmeans2_empty(self, xp):
@@ -390,8 +380,6 @@ class TestKMean:
         xp_assert_close(res[0], xp.asarray([4.], dtype=xp.float64))
         xp_assert_close(res[1], xp.asarray(2.3999999999999999, dtype=xp.float64)[()])
 
-    @skip_xp_backends('jax.numpy',
-                      reason='jax arrays do not support item assignment')
     def test_kmeans2_kpp_low_dim(self, xp):
         # Regression test for gh-11462
         rng = np.random.default_rng(2358792345678234568)
@@ -401,8 +389,6 @@ class TestKMean:
         xp_assert_close(res, prev_res)
 
     @pytest.mark.thread_unsafe
-    @skip_xp_backends('jax.numpy',
-                      reason='jax arrays do not support item assignment')
     def test_kmeans2_kpp_high_dim(self, xp):
         # Regression test for gh-11462
         rng = np.random.default_rng(23587923456834568)
@@ -427,8 +413,6 @@ class TestKMean:
         xp_assert_close(res[0], xp.asarray([-0.4,  8.], dtype=xp.float64))
         xp_assert_close(res[1], xp.asarray(1.0666666666666667, dtype=xp.float64)[()])
 
-    @skip_xp_backends('jax.numpy',
-                      reason='jax arrays do not support item assignment')
     def test_kmeans_and_kmeans2_random_seed(self, xp):
 
         seed_list = [

scipy/cluster/vq.py

@@ -67,9 +67,8 @@ code book.
 import warnings
 import numpy as np
 from collections import deque
-from scipy._lib._array_api import (
-    _asarray, array_namespace, xp_size, xp_copy
-)
+from scipy._lib._array_api import (_asarray, array_namespace, is_lazy_array,
+                                   xp_capabilities, xp_copy, xp_size)
 from scipy._lib._util import (check_random_state, rng_integers,
                               _transition_to_rng)
 from scipy._lib import array_api_extra as xpx
@@ -86,7 +85,8 @@ class ClusterError(Exception):
     pass
 
 
-def whiten(obs, check_finite=True):
+@xp_capabilities()
+def whiten(obs, check_finite=None):
     """
     Normalize a group of observations on a per feature basis.
 
@@ -100,19 +100,19 @@ def whiten(obs, check_finite=True):
     ----------
     obs : ndarray
         Each row of the array is an observation.  The
-        columns are the features seen during each observation.
+        columns are the features seen during each observation::
 
-        >>> #         f0    f1    f2
-        >>> obs = [[  1.,   1.,   1.],  #o0
-        ...        [  2.,   2.,   2.],  #o1
-        ...        [  3.,   3.,   3.],  #o2
-        ...        [  4.,   4.,   4.]]  #o3
+            #        f0  f1  f2
+            obs = [[ 1., 1., 1.],  #o0
+                   [ 2., 2., 2.],  #o1
+                   [ 3., 3., 3.],  #o2
+                   [ 4., 4., 4.]]  #o3
 
     check_finite : bool, optional
         Whether to check that the input matrices contain only finite numbers.
         Disabling may give a performance gain, but may result in problems
         (crashes, non-termination) if the inputs do contain infinities or NaNs.
-        Default: True
+        Default: True for eager backends and False for lazy ones.
 
     Returns
     -------
@@ -134,17 +134,21 @@ def whiten(obs, check_finite=True):
 
     """
     xp = array_namespace(obs)
+    if check_finite is None:
+        check_finite = not is_lazy_array(obs)
     obs = _asarray(obs, check_finite=check_finite, xp=xp)
     std_dev = xp.std(obs, axis=0)
     zero_std_mask = std_dev == 0
-    if xp.any(zero_std_mask):
-        std_dev[zero_std_mask] = 1.0
+    std_dev = xpx.at(std_dev, zero_std_mask).set(1.0)
+    if check_finite and xp.any(zero_std_mask):
         warnings.warn("Some columns have standard deviation zero. "
                       "The values of these columns will not change.",
                       RuntimeWarning, stacklevel=2)
     return obs / std_dev
 
 
+@xp_capabilities(cpu_only=True, reason="uses spatial.distance.cdist",
+                 jax_jit=False, allow_dask_compute=True)
 def vq(obs, code_book, check_finite=True):
     """
     Assign codes from a code book to observations.
@@ -168,13 +172,12 @@ def vq(obs, code_book, check_finite=True):
     code_book : ndarray
         The code book is usually generated using the k-means algorithm.
         Each row of the array holds a different code, and the columns are
-        the features of the code.
+        the features of the code::
 
-         >>> #              f0    f1    f2   f3
-         >>> code_book = [
-         ...             [  1.,   2.,   3.,   4.],  #c0
-         ...             [  1.,   2.,   3.,   4.],  #c1
-         ...             [  1.,   2.,   3.,   4.]]  #c2
+            #              f0  f1  f2  f3
+            code_book = [[ 1., 2., 3., 4.],  #c0
+                         [ 1., 2., 3., 4.],  #c1
+                         [ 1., 2., 3., 4.]]  #c2
 
     check_finite : bool, optional
         Whether to check that the input matrices contain only finite numbers.
@@ -208,10 +211,9 @@ def vq(obs, code_book, check_finite=True):
     code_book = _asarray(code_book, xp=xp, check_finite=check_finite)
     ct = xp.result_type(obs, code_book)
 
-    c_obs = xp.astype(obs, ct, copy=False)
-    c_code_book = xp.astype(code_book, ct, copy=False)
-
     if xp.isdtype(ct, kind='real floating'):
+        c_obs = xp.astype(obs, ct, copy=False)
+        c_code_book = xp.astype(code_book, ct, copy=False)
         c_obs = np.asarray(c_obs)
         c_code_book = np.asarray(c_code_book)
         result = _vq.vq(c_obs, c_code_book)
@@ -310,25 +312,24 @@ def _kmeans(obs, guess, thresh=1e-5, xp=None):
     code_book = guess
     diff = xp.inf
     prev_avg_dists = deque([diff], maxlen=2)
+
+    np_obs = np.asarray(obs)
     while diff > thresh:
         # compute membership and distances between obs and code_book
         obs_code, distort = vq(obs, code_book, check_finite=False)
         prev_avg_dists.append(xp.mean(distort, axis=-1))
         # recalc code_book as centroids of associated obs
-        obs = np.asarray(obs)
         obs_code = np.asarray(obs_code)
-        code_book, has_members = _vq.update_cluster_means(obs, obs_code,
+        code_book, has_members = _vq.update_cluster_means(np_obs, obs_code,
                                                           code_book.shape[0])
-        obs = xp.asarray(obs)
-        obs_code = xp.asarray(obs_code)
-        code_book = xp.asarray(code_book)
-        has_members = xp.asarray(has_members)
         code_book = code_book[has_members]
+        code_book = xp.asarray(code_book)
         diff = xp.abs(prev_avg_dists[0] - prev_avg_dists[1])
 
     return code_book, prev_avg_dists[1]
 
 
+@xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True)
 @_transition_to_rng("seed")
 def kmeans(obs, k_or_guess, iter=20, thresh=1e-5, check_finite=True,
            *, rng=None):
@@ -480,7 +481,7 @@ def kmeans(obs, k_or_guess, iter=20, thresh=1e-5, check_finite=True,
     if k != guess:
         raise ValueError("If k_or_guess is a scalar, it must be an integer.")
     if k < 1:
-        raise ValueError("Asked for %d clusters." % k)
+        raise ValueError(f"Asked for {k} clusters.")
 
     rng = check_random_state(rng)
 
@@ -609,15 +610,16 @@ def _kpp(data, k, rng, xp):
 
     for i in range(k):
         if i == 0:
-            init[i, :] = data[rng_integers(rng, data.shape[0]), :]
-
+            data_idx = rng_integers(rng, data.shape[0])
         else:
             D2 = cdist(init[:i,:], data, metric='sqeuclidean').min(axis=0)
             probs = D2/D2.sum()
             cumprobs = probs.cumsum()
             r = rng.uniform()
             cumprobs = np.asarray(cumprobs)
-            init[i, :] = data[np.searchsorted(cumprobs, r), :]
+            data_idx = int(np.searchsorted(cumprobs, r))
+
+        init = xpx.at(init)[i, :].set(data[data_idx, :])
 
     if ndim == 1:
         init = init[:, 0]
@@ -643,6 +645,7 @@ def _missing_raise():
 _valid_miss_meth = {'warn': _missing_warn, 'raise': _missing_raise}
 
 
+@xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True)
 @_transition_to_rng("seed")
 def kmeans2(data, k, iter=10, thresh=1e-5, minit='random',
             missing='warn', check_finite=True, *, rng=None):
@@ -799,8 +802,9 @@ def kmeans2(data, k, iter=10, thresh=1e-5, minit='random',
         nc = int(code_book)
 
         if nc < 1:
-            raise ValueError("Cannot ask kmeans2 for %d clusters"
-                             " (k was %s)" % (nc, code_book))
+            raise ValueError(
+                f"Cannot ask kmeans2 for {nc} clusters (k was {code_book})"
+            )
         elif nc != code_book:
             warnings.warn("k was not an integer, was converted.", stacklevel=2)
 
@@ -814,7 +818,7 @@ def kmeans2(data, k, iter=10, thresh=1e-5, minit='random',
 
     data = np.asarray(data)
     code_book = np.asarray(code_book)
-    for i in range(iter):
+    for _ in range(iter):
         # Compute the nearest neighbor for each obs using the current code book
         label = vq(data, code_book, check_finite=check_finite)[0]
         # Update the code book by computing centroids