diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 5a10f9d..4d5a78f 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -59,6 +59,7 @@ import numba import numpy as np +import numpy.typing as npt __all__ = [ "DEFAULT_RECIPE", @@ -872,7 +873,7 @@ def toarray(self) -> np.ndarray: ) return out - def to_scipy_sparse(self) -> Any: + def to_scipy_sparse(self, dtype: npt.DTypeLike = np.float64) -> Any: """The uncentered, scaled sparse ``Delta`` term, as a real scipy sparse array. Same sparsity pattern as the underlying raw array (a ``csr_array`` @@ -881,11 +882,23 @@ def to_scipy_sparse(self) -> Any: this view to code (such as ``parafac2``'s CuPy/MLX GPU backends) that only knows how to move a plain NumPy/SciPy array onto a device, rather than this view's own ``__matmul__``/``__rmatmul__``. + + Parameters + ---------- + dtype : numpy dtype-like, default ``np.float64`` + Dtype of the returned array's ``data`` (values only -- indices + stay at their existing width). Pass e.g. ``np.float32`` to + materialize directly in a lower-precision dtype instead of + materializing at float64 and downcasting afterwards, which + briefly holds both the float64 and downcast copies in memory at + once for no benefit when the caller only ever wanted the + smaller dtype (e.g. to match a GPU backend's own float32 + working precision). """ import scipy.sparse as sp arr = self._arr - data = np.empty(arr.nnz, dtype=np.float64) + data = np.empty(arr.nnz, dtype=dtype) if self._format == "csc": _materialize_delta_major_is_col( arr.major_ptr, diff --git a/tests/test_vcs_norm.py b/tests/test_vcs_norm.py index 16858c3..5885fe1 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -189,6 +189,20 @@ def test_to_scipy_sparse_matches_format(dense, vcls): assert sparse.dtype == np.float64 +def test_to_scipy_sparse_dtype_argument(dense, vcls): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + default = nv.to_scipy_sparse() + f32 = nv.to_scipy_sparse(dtype=np.float32) + assert f32.dtype == np.float32 + assert f32.nnz == default.nnz + np.testing.assert_allclose( + f32.toarray(), default.toarray().astype(np.float32), rtol=1e-5, atol=1e-5 + ) + + def test_means_property_matches_c_times_s(dense, vcls): if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0")