Allow for custom fill-values on conversion to/from scipy.sparse. (#685)

Author: hameerabbasi

sparse/numba_backend/_compressed/compressed.py

@@ -11,9 +11,10 @@
 from .._coo.core import COO
 from .._sparse_array import SparseArray
 from .._utils import (
+    _zero_of_dtype,
     can_store,
     check_compressed_axes,
-    check_zero_fill_value,
+    check_fill_value,
     equivalent,
     normalize_axis,
 )
@@ -137,7 +138,7 @@ def __init__(
         shape=None,
         compressed_axes=None,
         prune=False,
-        fill_value=0,
+        fill_value=None,
         idx_dtype=None,
     ):
         from .._common import _is_scipy_sparse_obj
@@ -176,8 +177,11 @@ def __init__(
 
         self.shape = shape
 
+        if fill_value is None:
+            fill_value = _zero_of_dtype(self.data.dtype)
+
         self._compressed_axes = tuple(compressed_axes) if isinstance(compressed_axes, Iterable) else None
-        self.fill_value = fill_value
+        self.fill_value = self.data.dtype.type(fill_value)
 
         if prune:
             self._prune()
@@ -194,7 +198,7 @@ def copy(self, deep=True):
         return _copy.deepcopy(self) if deep else _copy.copy(self)
 
     @classmethod
-    def from_numpy(cls, x, compressed_axes=None, fill_value=0, idx_dtype=None):
+    def from_numpy(cls, x, compressed_axes=None, fill_value=None, idx_dtype=None):
         coo = COO.from_numpy(x, fill_value=fill_value, idx_dtype=idx_dtype)
         return cls.from_coo(coo, compressed_axes, idx_dtype)
 
@@ -204,12 +208,12 @@ def from_coo(cls, x, compressed_axes=None, idx_dtype=None):
         return cls(arg, shape=shape, compressed_axes=compressed_axes, fill_value=fill_value)
 
     @classmethod
-    def from_scipy_sparse(cls, x):
+    def from_scipy_sparse(cls, x, /, *, fill_value=None):
         if x.format == "csc":
-            return cls((x.data, x.indices, x.indptr), shape=x.shape, compressed_axes=(1,))
+            return cls((x.data, x.indices, x.indptr), shape=x.shape, compressed_axes=(1,), fill_value=fill_value)
 
         x = x.asformat("csr")
-        return cls((x.data, x.indices, x.indptr), shape=x.shape, compressed_axes=(0,))
+        return cls((x.data, x.indices, x.indptr), shape=x.shape, compressed_axes=(0,), fill_value=fill_value)
 
     @classmethod
     def from_iter(cls, x, shape=None, compressed_axes=None, fill_value=None, idx_dtype=None):
@@ -471,13 +475,20 @@ def todok(self):
 
         return DOK.from_coo(self.tocoo())  # probably a temporary solution
 
-    def to_scipy_sparse(self):
+    def to_scipy_sparse(self, accept_fv=None):
         """
         Converts this :obj:`GCXS` object into a :obj:`scipy.sparse.csr_matrix` or `scipy.sparse.csc_matrix`.
+
+        Parameters
+        ----------
+        accept_fv : scalar or list of scalar, optional
+            The list of accepted fill-values. The default accepts only zero.
+
         Returns
         -------
         :obj:`scipy.sparse.csr_matrix` or `scipy.sparse.csc_matrix`
             The converted Scipy sparse matrix.
+
         Raises
         ------
         ValueError
@@ -487,8 +498,7 @@ def to_scipy_sparse(self):
         """
         import scipy.sparse
 
-        check_zero_fill_value(self)
-
+        check_fill_value(self, accept_fv=accept_fv)
         if self.ndim != 2:
             raise ValueError("Can only convert a 2-dimensional array to a Scipy sparse matrix.")
 
@@ -873,9 +883,9 @@ def __init__(self, arg, shape=None, compressed_axes=class_compressed_axes, prune
         super().__init__(arg, shape=shape, compressed_axes=compressed_axes, fill_value=fill_value)
 
     @classmethod
-    def from_scipy_sparse(cls, x):
+    def from_scipy_sparse(cls, x, /, *, fill_value=None):
         x = x.asformat("csr", copy=False)
-        return cls((x.data, x.indices, x.indptr), shape=x.shape)
+        return cls((x.data, x.indices, x.indptr), shape=x.shape, fill_value=fill_value)
 
     def transpose(self, axes: None = None, copy: bool = False) -> Union["CSC", "CSR"]:
         axes = normalize_axis(axes, self.ndim)
@@ -905,9 +915,9 @@ def __init__(self, arg, shape=None, compressed_axes=class_compressed_axes, prune
         super().__init__(arg, shape=shape, compressed_axes=compressed_axes, fill_value=fill_value)
 
     @classmethod
-    def from_scipy_sparse(cls, x):
+    def from_scipy_sparse(cls, x, /, *, fill_value=None):
         x = x.asformat("csc", copy=False)
-        return cls((x.data, x.indices, x.indptr), shape=x.shape)
+        return cls((x.data, x.indices, x.indptr), shape=x.shape, fill_value=fill_value)
 
     def transpose(self, axes: None = None, copy: bool = False) -> Union["CSC", "CSR"]:
         axes = normalize_axis(axes, self.ndim)

sparse/numba_backend/_coo/core.py

@@ -15,6 +15,7 @@
 from .._utils import (
     _zero_of_dtype,
     can_store,
+    check_fill_value,
     check_zero_fill_value,
     equivalent,
     normalize_axis,
@@ -425,14 +426,16 @@ def todense(self):
         return x
 
     @classmethod
-    def from_scipy_sparse(cls, x):
+    def from_scipy_sparse(cls, x, /, *, fill_value=None):
         """
         Construct a :obj:`COO` array from a :obj:`scipy.sparse.spmatrix`
 
         Parameters
         ----------
         x : scipy.sparse.spmatrix
             The sparse matrix to construct the array from.
+        fill_value : scalar
+            The fill-value to use when converting.
 
         Returns
         -------
@@ -456,6 +459,7 @@ def from_scipy_sparse(cls, x):
             shape=x.shape,
             has_duplicates=not x.has_canonical_format,
             sorted=x.has_canonical_format,
+            fill_value=fill_value,
         )
 
     @classmethod
@@ -1155,10 +1159,15 @@ def resize(self, *args, refcheck=True, coords_dtype=np.intp):
         if len(self.data) != len(linear_loc):
             self.data = self.data[:end_idx].copy()
 
-    def to_scipy_sparse(self):
+    def to_scipy_sparse(self, /, *, accept_fv=None):
         """
         Converts this :obj:`COO` object into a :obj:`scipy.sparse.coo_matrix`.
 
+        Parameters
+        ----------
+        accept_fv : scalar or list of scalar, optional
+            The list of accepted fill-values. The default accepts only zero.
+
         Returns
         -------
         :obj:`scipy.sparse.coo_matrix`
@@ -1178,7 +1187,7 @@ def to_scipy_sparse(self):
         """
         import scipy.sparse
 
-        check_zero_fill_value(self)
+        check_fill_value(self, accept_fv=accept_fv)
 
         if self.ndim != 2:
             raise ValueError("Can only convert a 2-dimensional array to a Scipy sparse matrix.")

sparse/numba_backend/_dok.py

@@ -131,14 +131,16 @@ def __init__(self, shape, data=None, dtype=None, fill_value=None):
             raise ValueError("data must be a dict.")
 
     @classmethod
-    def from_scipy_sparse(cls, x):
+    def from_scipy_sparse(cls, x, /, *, fill_value=None):
         """
         Create a :obj:`DOK` array from a :obj:`scipy.sparse.spmatrix`.
 
         Parameters
         ----------
         x : scipy.sparse.spmatrix
             The matrix to convert.
+        fill_value : scalar
+            The fill-value to use when converting.
 
         Returns
         -------
@@ -154,7 +156,7 @@ def from_scipy_sparse(cls, x):
         """
         from sparse import COO
 
-        return COO.from_scipy_sparse(x).asformat(cls)
+        return COO.from_scipy_sparse(x, fill_value=fill_value).asformat(cls)
 
     @classmethod
     def from_coo(cls, x):

sparse/numba_backend/_utils.py

@@ -33,7 +33,7 @@ def assert_eq(x, y, check_nnz=True, compare_dtype=True, **kwargs):
     if isinstance(x, COO) and isinstance(y, COO) and check_nnz:
         assert np.array_equal(x.coords, y.coords)
         assert check_equal(x.data, y.data, **kwargs)
-        assert x.fill_value == y.fill_value
+        assert x.fill_value == y.fill_value or (np.isnan(x.fill_value) and np.isnan(y.fill_value))
         return
 
     if hasattr(x, "todense"):
@@ -528,6 +528,31 @@ def check_compressed_axes(ndim, compressed_axes):
         raise ValueError("axis out of range")
 
 
+def check_fill_value(x, /, *, accept_fv=None) -> None:
+    """Raises on incorrect fill-values.
+
+    Parameters
+    ----------
+    x : SparseArray
+        The array to check
+    accept_fv : scalar or list of scalar, optional
+        The list of accepted fill-values. The default accepts only zero.
+
+    Raises
+    ------
+    ValueError
+        If the fill-value doesn't match.
+    """
+    if accept_fv is None:
+        accept_fv = [0]
+
+    if not isinstance(accept_fv, Iterable):
+        accept_fv = [accept_fv]
+
+    if not any(equivalent(fv, x.fill_value) for fv in accept_fv):
+        raise ValueError(f"{x.fill_value=} but should be in {accept_fv}.")
+
+
 def check_zero_fill_value(*args):
     """
     Checks if all the arguments have zero fill-values.

sparse/numba_backend/tests/test_compressed.py

@@ -1,11 +1,10 @@
 import sparse
 from sparse.numba_backend._compressed import GCXS
-from sparse.numba_backend._utils import assert_eq
+from sparse.numba_backend._utils import assert_eq, equivalent
 
 import pytest
 
 import numpy as np
-import scipy
 
 
 @pytest.fixture(scope="module", params=["f8", "f4", "i8", "i4"])
@@ -177,18 +176,21 @@ def test_tranpose(a, b):
     assert_eq(x.transpose(b), s.transpose(b))
 
 
-def test_to_scipy_sparse():
-    s = sparse.random((3, 5), density=0.5, format="gcxs", compressed_axes=(0,))
-    a = s.to_scipy_sparse()
-    b = scipy.sparse.csr_matrix(s.todense())
+@pytest.mark.parametrize("fill_value_in", [0, np.inf, np.nan, 5, None])
+@pytest.mark.parametrize("fill_value_out", [0, np.inf, np.nan, 5, None])
+@pytest.mark.parametrize("format", [sparse.COO, sparse._compressed.CSR])
+def test_to_scipy_sparse(fill_value_in, fill_value_out, format):
+    s = sparse.random((3, 5), density=0.5, format=format, fill_value=fill_value_in)
 
-    assert_eq(a, b)
+    if not ((fill_value_in in {0, None} and fill_value_out in {0, None}) or equivalent(fill_value_in, fill_value_out)):
+        with pytest.raises(ValueError, match=r"fill_value=.* but should be in .*\."):
+            s.to_scipy_sparse(accept_fv=fill_value_out)
+        return
 
-    s = sparse.random((3, 5), density=0.5, format="gcxs", compressed_axes=(1,))
-    a = s.to_scipy_sparse()
-    b = scipy.sparse.csc_matrix(s.todense())
+    sps_matrix = s.to_scipy_sparse(accept_fv=fill_value_in)
+    s2 = format.from_scipy_sparse(sps_matrix, fill_value=fill_value_out)
 
-    assert_eq(a, b)
+    assert_eq(s, s2)
 
 
 def test_tocoo():

sparse/numba_backend/tests/test_coo.py

@@ -380,14 +380,6 @@ def test_reshape_errors(format):
         s.reshape((3, 5, 1), order="F")
 
 
-def test_to_scipy_sparse():
-    s = sparse.random((3, 5), density=0.5)
-    a = s.to_scipy_sparse()
-    b = scipy.sparse.coo_matrix(s.todense())
-
-    assert_eq(a, b)
-
-
 @pytest.mark.parametrize("a_ndim", [1, 2, 3])
 @pytest.mark.parametrize("b_ndim", [1, 2, 3])
 def test_kron(a_ndim, b_ndim):