✨ linalg: support sparse arrays in clarkson_woodruff_transform

scipy-stubs/linalg/_sketches.pyi

@@ -1,54 +1,114 @@
-from typing import overload
+from typing import Any, Never, TypeAlias, TypeVar, overload
 
 import numpy as np
+import optype as op
 import optype.numpy as onp
+import optype.numpy.compat as npc
 
-from scipy.sparse import csc_matrix
+from scipy.sparse import csc_matrix, sparray, spmatrix
+from scipy.sparse._base import _spbase
 
 __all__ = ["clarkson_woodruff_transform"]
 
+_ScalarT = TypeVar("_ScalarT", bound=npc.number | np.bool_)
+
+_ToInt: TypeAlias = int | npc.integer
+_ToSparse: TypeAlias = _spbase[_ScalarT] | sparray[_ScalarT] | spmatrix[_ScalarT]
+
 ###
 
-def cwt_matrix(n_rows: onp.ToInt, n_columns: onp.ToInt, rng: onp.random.ToRNG | None = None) -> csc_matrix[np.int_]: ...
+def cwt_matrix(n_rows: _ToInt, n_columns: _ToInt, rng: onp.random.ToRNG | None = None) -> csc_matrix[np.int_]: ...
 
 #
 @overload
 def clarkson_woodruff_transform(
-    input_matrix: onp.ToIntND,
-    sketch_size: onp.ToInt,
+    input_matrix: _ToSparse[Never],
+    sketch_size: _ToInt,
+    rng: onp.random.ToRNG | None = None,
+    *,
+    seed: onp.random.ToRNG | None = None,
+) -> csc_matrix[Any]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: _ToSparse[npc.integer | np.bool_],
+    sketch_size: _ToInt,
+    rng: onp.random.ToRNG | None = None,
+    *,
+    seed: onp.random.ToRNG | None = None,
+) -> csc_matrix[np.int_]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: _ToSparse[np.float32 | np.float64],
+    sketch_size: _ToInt,
+    rng: onp.random.ToRNG | None = None,
+    *,
+    seed: onp.random.ToRNG | None = None,
+) -> csc_matrix[np.float64]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: _ToSparse[np.longdouble],
+    sketch_size: _ToInt,
     rng: onp.random.ToRNG | None = None,
     *,
     seed: onp.random.ToRNG | None = None,
+) -> csc_matrix[np.longdouble]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: _ToSparse[np.complex64 | np.complex128],
+    sketch_size: _ToInt,
+    rng: onp.random.ToRNG | None = None,
+    *,
+    seed: onp.random.ToRNG | None = None,
+) -> csc_matrix[np.complex128]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: _ToSparse[np.clongdouble],
+    sketch_size: _ToInt,
+    rng: onp.random.ToRNG | None = None,
+    *,
+    seed: onp.random.ToRNG | None = None,
+) -> csc_matrix[np.clongdouble]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: onp.CanArrayND[Never],
+    sketch_size: _ToInt,
+    rng: onp.random.ToRNG | None = None,
+    *,
+    seed: onp.random.ToRNG | None = None,
+) -> onp.ArrayND[Any]: ...
+@overload
+def clarkson_woodruff_transform(
+    input_matrix: onp.ToIntND, sketch_size: _ToInt, rng: onp.random.ToRNG | None = None, *, seed: onp.random.ToRNG | None = None
 ) -> onp.ArrayND[np.int_]: ...
 @overload
 def clarkson_woodruff_transform(
-    input_matrix: onp.ToJustFloat64_ND,
-    sketch_size: onp.ToInt,
+    input_matrix: onp.ToArrayND[op.JustFloat, np.float16 | np.float32 | np.float64],
+    sketch_size: _ToInt,
     rng: onp.random.ToRNG | None = None,
     *,
     seed: onp.random.ToRNG | None = None,
 ) -> onp.ArrayND[np.float64]: ...
 @overload
 def clarkson_woodruff_transform(
-    input_matrix: onp.ToJustFloatND,
-    sketch_size: onp.ToInt,
+    input_matrix: onp.ToArrayND[np.longdouble, np.longdouble],
+    sketch_size: _ToInt,
     rng: onp.random.ToRNG | None = None,
     *,
     seed: onp.random.ToRNG | None = None,
-) -> onp.ArrayND[np.float64 | np.longdouble]: ...
+) -> onp.ArrayND[np.longdouble]: ...
 @overload
 def clarkson_woodruff_transform(
-    input_matrix: onp.ToJustComplex128_ND,
-    sketch_size: onp.ToInt,
+    input_matrix: onp.ToArrayND[op.JustComplex, np.complex64 | np.complex128],
+    sketch_size: _ToInt,
     rng: onp.random.ToRNG | None = None,
     *,
     seed: onp.random.ToRNG | None = None,
 ) -> onp.ArrayND[np.complex128]: ...
 @overload
 def clarkson_woodruff_transform(
-    input_matrix: onp.ToJustComplexND,
-    sketch_size: onp.ToInt,
+    input_matrix: onp.ToArrayND[np.clongdouble, np.clongdouble],
+    sketch_size: _ToInt,
     rng: onp.random.ToRNG | None = None,
     *,
     seed: onp.random.ToRNG | None = None,
-) -> onp.ArrayND[np.complex128 | np.clongdouble]: ...
+) -> onp.ArrayND[np.clongdouble]: ...

tests/linalg/test__sketches.pyi

@@ -0,0 +1,47 @@
+from typing import Any, assert_type
+
+import numpy as np
+import numpy.typing as npt
+
+from scipy.linalg import clarkson_woodruff_transform
+from scipy.sparse import csc_matrix, sparray, spmatrix
+
+like_bool_2d: list[list[bool]]
+like_i64_2d: list[list[int]]
+like_f64_2d: list[list[float]]
+like_c128_2d: list[list[complex]]
+
+arr_any: npt.NDArray[Any]
+arr_i8: npt.NDArray[np.int8]
+arr_f32: npt.NDArray[np.float32]
+arr_f80: npt.NDArray[np.longdouble]
+arr_c64: npt.NDArray[np.complex64]
+arr_c160: npt.NDArray[np.clongdouble]
+
+sparse_any: spmatrix[Any] | sparray[Any]
+sparse_i8: spmatrix[np.int8] | sparray[np.int8]
+sparse_f32: spmatrix[np.float32] | sparray[np.float32]
+sparse_f80: spmatrix[np.longdouble] | sparray[np.longdouble]
+sparse_c64: spmatrix[np.complex64] | sparray[np.complex64]
+sparse_c160: spmatrix[np.clongdouble] | sparray[np.clongdouble]
+
+###
+
+assert_type(clarkson_woodruff_transform(like_bool_2d, 2), npt.NDArray[np.int_])
+assert_type(clarkson_woodruff_transform(like_i64_2d, 2), npt.NDArray[np.int_])
+assert_type(clarkson_woodruff_transform(like_f64_2d, 2), npt.NDArray[np.float64])
+assert_type(clarkson_woodruff_transform(like_c128_2d, 2), npt.NDArray[np.complex128])
+
+assert_type(clarkson_woodruff_transform(arr_any, 2), npt.NDArray[Any])  # type: ignore[assert-type]
+assert_type(clarkson_woodruff_transform(arr_i8, 2), npt.NDArray[np.int_])
+assert_type(clarkson_woodruff_transform(arr_f32, 2), npt.NDArray[np.float64])
+assert_type(clarkson_woodruff_transform(arr_c64, 2), npt.NDArray[np.complex128])
+assert_type(clarkson_woodruff_transform(arr_f80, 2), npt.NDArray[np.longdouble])
+assert_type(clarkson_woodruff_transform(arr_c160, 2), npt.NDArray[np.clongdouble])
+
+assert_type(clarkson_woodruff_transform(sparse_any, 2), csc_matrix[Any])
+assert_type(clarkson_woodruff_transform(sparse_i8, 2), csc_matrix[np.int_])
+assert_type(clarkson_woodruff_transform(sparse_f32, 2), csc_matrix[np.float64])
+assert_type(clarkson_woodruff_transform(sparse_c64, 2), csc_matrix[np.complex128])
+assert_type(clarkson_woodruff_transform(sparse_f80, 2), csc_matrix[np.longdouble])
+assert_type(clarkson_woodruff_transform(sparse_c160, 2), csc_matrix[np.clongdouble])