more fixes from Thomas (#20)

* address comments and get rest of tests passing

* Force cast to numpy in scipy.ndimage

* fix binary_erosion

* fix sosfilt

* clean up

* address comments

* fix a couple ndimage tests

* integrate latest array-api-compat

* clean up

Author: lithomas1

scipy/_lib/array_api_compat

@@ -21,6 +21,7 @@
 import scipy._lib.array_api_extra as xpx
 from scipy import cluster, interpolate, linalg, optimize, sparse, spatial, stats
 
+skip_xp_backends = pytest.mark.skip_xp_backends
 
 @pytest.mark.slow
 def test__aligned_zeros():
@@ -586,6 +587,7 @@ class TestLazywhere:
 
     @pytest.mark.fail_slow(10)
     @pytest.mark.filterwarnings('ignore::RuntimeWarning')  # overflows, etc.
+    @skip_xp_backends("dask.array", reason="lazywhere doesn't work with dask")
     @given(n_arrays=n_arrays, rng_seed=rng_seed, dtype=dtype, p=p, data=data)
     @pytest.mark.thread_unsafe
     def test_basic(self, n_arrays, rng_seed, dtype, p, data, xp):

scipy/_lib/tests/test__util.py

@@ -8,6 +8,8 @@
 from scipy._lib import array_api_extra as xpx
 from scipy._lib._array_api_no_0d import xp_assert_equal as xp_assert_equal_no_0d
 
+skip_xp_backends = pytest.mark.skip_xp_backends
+
 
 @pytest.mark.skipif(not _GLOBAL_CONFIG["SCIPY_ARRAY_API"],
         reason="Array API test; set environment variable SCIPY_ARRAY_API=1 to run it")
@@ -59,6 +61,10 @@ def test_array_api_extra_hook(self):
         with pytest.raises(TypeError, match=msg):
             xpx.atleast_nd("abc", ndim=0)
 
+    @skip_xp_backends(
+        "dask.array",
+        reason="raw dask.array namespace doesn't ignores copy=True in asarray"
+    )
     def test_copy(self, xp):
         for _xp in [xp, None]:
             x = xp.asarray([1, 2, 3])

scipy/_lib/tests/test_array_api.py

@@ -630,6 +630,10 @@ def single_group(vals):
     if labels is None:
         return single_group(input)
 
+    # manually cast to numpy
+    # since libaries (e.g. dask) that implement __array_function__
+    # will not return a numpy array from broadcast_arrays
+    labels = np.asarray(labels)
     # ensure input and labels match sizes
     input, labels = np.broadcast_arrays(input, labels)
 
@@ -944,6 +948,10 @@ def single_group(vals, positions):
     if labels is None:
         return single_group(input, positions)
 
+    # manually cast to numpy
+    # since libaries (e.g. dask) that implement __array_function__
+    # will not return a numpy array from broadcast_arrays
+    labels = np.asarray(labels)
     # ensure input and labels match sizes
     input, labels = np.broadcast_arrays(input, labels)
 

scipy/ndimage/_measurements.py

@@ -36,6 +36,8 @@
 from . import _nd_image
 from . import _filters
 
+from scipy._lib._array_api import is_dask, array_namespace
+
 __all__ = ['iterate_structure', 'generate_binary_structure', 'binary_erosion',
            'binary_dilation', 'binary_opening', 'binary_closing',
            'binary_hit_or_miss', 'binary_propagation', 'binary_fill_holes',
@@ -220,7 +222,14 @@ def _binary_erosion(input, structure, iterations, mask, output,
     except TypeError as e:
         raise TypeError('iterations parameter should be an integer') from e
 
-    input = np.asarray(input)
+    if is_dask(array_namespace(input)):
+        # Note: If you create an dask array with ones
+        # it does a stride trick where it makes an array
+        # (with stride 0) using a scalar
+        # this messes up the C ndimage iteration code
+        input = np.asarray(input, order="C")
+    else:
+        input = np.asarray(input)
     ndim = input.ndim
     if np.iscomplexobj(input):
         raise TypeError('Complex type not supported')

scipy/ndimage/_morphology.py

@@ -191,7 +191,7 @@ def test_correlate01(self, xp):
 
     @xfail_xp_backends('cupy', reason="Differs by a factor of two?")
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
-    @skip_xp_backends("dask.array", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="wrong answer")
     def test_correlate01_overlap(self, xp):
         array = xp.reshape(xp.arange(256), (16, 16))
         weights = xp.asarray([2])
@@ -881,7 +881,7 @@ def test_gauss06(self, xp):
         assert_array_almost_equal(output1, output2)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
-    @skip_xp_backends("dask.array", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="wrong result")
     def test_gauss_memory_overlap(self, xp):
         input = xp.arange(100 * 100, dtype=xp.float32)
         input = xp.reshape(input, (100, 100))
@@ -2647,7 +2647,6 @@ def test_gaussian_radius_invalid(xp):
 
 
 @skip_xp_backends("jax.numpy", reason="output array is read-only")
-@skip_xp_backends("dask.array", reason="output array is read-only.")
 class TestThreading:
     def check_func_thread(self, n, fun, args, out):
         from threading import Thread

scipy/ndimage/tests/test_filters.py

@@ -149,14 +149,14 @@ def test_basic(self, xp):
 
 
 def test_label01(xp):
-    data = xp.ones([])
+    data = xp.ones(())
     out, n = ndimage.label(data)
     assert out == 1
     assert n == 1
 
 
 def test_label02(xp):
-    data = xp.zeros([])
+    data = xp.zeros(())
     out, n = ndimage.label(data)
     assert out == 0
     assert n == 0
@@ -557,6 +557,8 @@ def test_value_indices03(xp):
 
         trueKeys = xp.unique_values(a)
         vi = ndimage.value_indices(a)
+        # TODO: list(trueKeys) needs len of trueKeys
+        # (which is unknown for dask since it is the result of an unique call)
         assert list(vi.keys()) == list(trueKeys)
         for k in [int(x) for x in trueKeys]:
             trueNdx = xp.nonzero(a == k)
@@ -665,7 +667,6 @@ def test_sum11(xp):
         assert_almost_equal(output, xp.asarray(6.0), check_0d=False)
 
 
-@skip_xp_backends("dask.array", reason="data-dependent output shapes")
 def test_sum12(xp):
     labels = xp.asarray([[1, 2], [2, 4]], dtype=xp.int8)
     for type in types:
@@ -675,7 +676,6 @@ def test_sum12(xp):
         assert_array_almost_equal(output, xp.asarray([4.0, 0.0, 5.0]))
 
 
-@skip_xp_backends("dask.array", reason="data-dependent output shapes")
 def test_sum_labels(xp):
     labels = xp.asarray([[1, 2], [2, 4]], dtype=xp.int8)
     for type in types:
@@ -688,7 +688,6 @@ def test_sum_labels(xp):
         assert xp.all(output_sum == output_labels)
         assert_array_almost_equal(output_labels, xp.asarray([4.0, 0.0, 5.0]))
 
-@xfail_xp_backends("dask.array", reason="dask outputs wrong results here")
 def test_mean01(xp):
     labels = np.asarray([1, 0], dtype=bool)
     labels = xp.asarray(labels)
@@ -699,7 +698,6 @@ def test_mean01(xp):
         assert_almost_equal(output, xp.asarray(2.0), check_0d=False)
 
 
-@xfail_xp_backends("dask.array", reason="dask outputs wrong results here")
 def test_mean02(xp):
     labels = np.asarray([1, 0], dtype=bool)
     input = np.asarray([[1, 2], [3, 4]], dtype=bool)
@@ -710,7 +708,6 @@ def test_mean02(xp):
     assert_almost_equal(output, xp.asarray(1.0), check_0d=False)
 
 
-@xfail_xp_backends("dask.array", reason="dask outputs wrong results here")
 def test_mean03(xp):
     labels = xp.asarray([1, 2])
     for type in types:
@@ -721,7 +718,6 @@ def test_mean03(xp):
         assert_almost_equal(output, xp.asarray(3.0), check_0d=False)
 
 
-@xfail_xp_backends("dask.array", reason="dask outputs wrong results here")
 def test_mean04(xp):
     labels = xp.asarray([[1, 2], [2, 4]], dtype=xp.int8)
     with np.errstate(all='ignore'):
@@ -768,7 +764,6 @@ def test_minimum03(xp):
         assert_almost_equal(output, xp.asarray(2.0), check_0d=False)
 
 
-@skip_xp_backends('dask.array', reason="no argsort in Dask")
 def test_minimum04(xp):
     labels = xp.asarray([[1, 2], [2, 3]])
     for type in types:
@@ -808,7 +803,6 @@ def test_maximum03(xp):
         assert_almost_equal(output, xp.asarray(4.0), check_0d=False)
 
 
-@skip_xp_backends('dask.array', reason="no argsort in Dask")
 def test_maximum04(xp):
     labels = xp.asarray([[1, 2], [2, 3]])
     for type in types:
@@ -848,8 +842,6 @@ def test_median02(xp):
     assert_almost_equal(output, xp.asarray(1.0), check_0d=False)
 
 
-@skip_xp_backends("dask.array",
-                  reason="dask.array.median only implemented for along an axis.")
 def test_median03(xp):
     a = xp.asarray([[1, 2, 0, 1],
                     [5, 3, 0, 4],
@@ -863,15 +855,13 @@ def test_median03(xp):
     assert_almost_equal(output, xp.asarray(3.0), check_0d=False)
 
 
-@xfail_xp_backends("dask.array", reason="Crash inside dask searchsorted")
 def test_median_gh12836_bool(xp):
     # test boolean addition fix on example from gh-12836
     a = np.asarray([1, 1], dtype=bool)
     a = xp.asarray(a)
     output = ndimage.median(a, labels=xp.ones((2,)), index=xp.asarray([1]))
     assert_array_almost_equal(output, xp.asarray([1.0]))
 
-@xfail_xp_backends("dask.array", reason="Crash inside dask searchsorted")
 def test_median_no_int_overflow(xp):
     # test integer overflow fix on example from gh-12836
     a = xp.asarray([65, 70], dtype=xp.int8)
@@ -912,10 +902,6 @@ def test_variance04(xp):
     output = ndimage.variance(input)
     assert_almost_equal(output, xp.asarray(0.25), check_0d=False)
 
-# dask.array is maybe due to failed conversion to numpy?
-# array-api-strict should've caught use of non array API functions I think
-@skip_xp_backends("dask.array",
-                  reason="conjugate called on dask.array which doesn't exist")
 def test_variance05(xp):
     labels = xp.asarray([2, 2, 3])
     for type in types:
@@ -925,7 +911,6 @@ def test_variance05(xp):
         output = ndimage.variance(input, labels, 2)
         assert_almost_equal(output, xp.asarray(1.0), check_0d=False)
 
-@skip_xp_backends("dask.array", reason="Data-dependent output shapes")
 def test_variance06(xp):
     labels = xp.asarray([2, 2, 3, 3, 4])
     with np.errstate(all='ignore'):
@@ -970,10 +955,6 @@ def test_standard_deviation04(xp):
     assert_almost_equal(output, xp.asarray(0.5), check_0d=False)
 
 
-# dask.array is maybe due to failed conversion to numpy?
-# array-api-strict should've caught use of non array API functions I think
-@skip_xp_backends("dask.array",
-                  reason="conjugate called on dask.array which doesn't exist")
 def test_standard_deviation05(xp):
     labels = xp.asarray([2, 2, 3])
     for type in types:
@@ -983,7 +964,6 @@ def test_standard_deviation05(xp):
         assert_almost_equal(output, xp.asarray(1.0), check_0d=False)
 
 
-@skip_xp_backends("dask.array", reason="data-dependent output shapes")
 def test_standard_deviation06(xp):
     labels = xp.asarray([2, 2, 3, 3, 4])
     with np.errstate(all='ignore'):
@@ -996,7 +976,6 @@ def test_standard_deviation06(xp):
             assert_array_almost_equal(output, xp.asarray([1.0, 1.0, 0.0]))
 
 
-@skip_xp_backends("dask.array", reason="data-dependent output shapes")
 def test_standard_deviation07(xp):
     labels = xp.asarray([1])
     with np.errstate(all='ignore'):
@@ -1070,7 +1049,6 @@ def test_minimum_position06(xp):
         assert output == (0, 1)
 
 
-@skip_xp_backends('dask.array', reason="no argsort in Dask")
 def test_minimum_position07(xp):
     labels = xp.asarray([1, 2, 3, 4])
     for type in types:
@@ -1136,7 +1114,6 @@ def test_maximum_position05(xp):
         assert output == (0, 0)
 
 
-@skip_xp_backends('dask.array', reason="no argsort in Dask")
 def test_maximum_position06(xp):
     labels = xp.asarray([1, 2, 0, 4])
     for type in types:
@@ -1149,7 +1126,6 @@ def test_maximum_position06(xp):
         assert output[0] == (0, 0)
         assert output[1] == (1, 1)
 
-@xfail_xp_backends("dask.array", reason="crash in dask.array searchsorted")
 @xfail_xp_backends("torch", reason="output[1] is wrong on pytorch")
 def test_maximum_position07(xp):
     # Test float labels
@@ -1165,7 +1141,6 @@ def test_maximum_position07(xp):
         assert output[1] == (0, 3)
 
 
-@xfail_xp_backends("dask.array", reason="dask wrong answer")
 def test_extrema01(xp):
     labels = np.asarray([1, 0], dtype=bool)
     labels = xp.asarray(labels)
@@ -1182,7 +1157,6 @@ def test_extrema01(xp):
         assert output1 == (output2, output3, output4, output5)
 
 
-@xfail_xp_backends("dask.array", reason="dask wrong answer")
 def test_extrema02(xp):
     labels = xp.asarray([1, 2])
     for type in types:
@@ -1201,7 +1175,6 @@ def test_extrema02(xp):
         assert output1 == (output2, output3, output4, output5)
 
 
-@skip_xp_backends('dask.array', reason="no argsort in Dask")
 def test_extrema03(xp):
     labels = xp.asarray([[1, 2], [2, 3]])
     for type in types:
@@ -1230,7 +1203,6 @@ def test_extrema03(xp):
         assert output1[3] == output5
 
 
-@skip_xp_backends('dask.array', reason="no argsort in Dask")
 def test_extrema04(xp):
     labels = xp.asarray([1, 2, 0, 4])
     for type in types:
@@ -1307,7 +1279,6 @@ def test_center_of_mass06(xp):
     assert output == expected
 
 
-@xfail_xp_backends("dask.array", reason="wrong output shape")
 def test_center_of_mass07(xp):
     labels = xp.asarray([1, 0])
     expected = (0.5, 0.0)
@@ -1317,7 +1288,6 @@ def test_center_of_mass07(xp):
     assert output == expected
 
 
-@xfail_xp_backends("dask.array", reason="wrong output shape")
 def test_center_of_mass08(xp):
     labels = xp.asarray([1, 2])
     expected = (0.5, 1.0)
@@ -1327,7 +1297,6 @@ def test_center_of_mass08(xp):
     assert output == expected
 
 
-@skip_xp_backends("dask.array", reason="data-dependent output shapes")
 def test_center_of_mass09(xp):
     labels = xp.asarray((1, 2))
     expected = xp.asarray([(0.5, 0.0), (0.5, 1.0)], dtype=xp.float64)
@@ -1365,7 +1334,6 @@ def test_histogram03(xp):
     assert_array_almost_equal(output[1], expected2)
 
 
-@skip_xp_backends("dask.array", reason="data-dependent output shapes")
 def test_stat_funcs_2d(xp):
     a = xp.asarray([[5, 6, 0, 0, 0], [8, 9, 0, 0, 0], [0, 0, 0, 3, 5]])
     lbl = xp.asarray([[1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [0, 0, 0, 2, 2]])

scipy/ndimage/tests/test_measurements.py

@@ -15,11 +15,6 @@
 pytestmark = [skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy'])]
 
 
-@xfail_xp_backends('dask.array',
-    reason="Dask.array gets wrong results here. "
-           "Some tests can pass when creating input array from list of ones"
-           "instead of xp.ones, so maybe something is getting corrupted here."
-)
 class TestNdimageMorphology:
 
     @xfail_xp_backends('cupy', reason='CuPy does not have distance_transform_bf.')

scipy/ndimage/tests/test_morphology.py

@@ -727,6 +727,11 @@ def group_delay(system, w=512, whole=False, fs=2*pi):
 
 def _validate_sos(sos):
     """Helper to validate a SOS input"""
+    # manually cast to numpy array
+    # since libs like dask implement __array_function__
+    # (and will return a dask array instead of casting to
+    # ndarray in atleast_2d)
+    sos = np.asarray(sos)
     sos = np.atleast_2d(sos)
     if sos.ndim != 2:
         raise ValueError('sos array must be 2D')