Sparse arrays (#1373)

* MAINT: Add sps.sparray to isinstance checks

* TST: sparse_array support in ad test

* MAINT: Applied suggestions from the review

---------

Co-authored-by: Ingrid Kristine Jacobsen <[email protected]>

Author: 2 people

src/porepy/models/constitutive_laws.py

@@ -430,7 +430,7 @@ def aperture(self, subdomains: list[pp.Grid]) -> pp.ad.Operator:
                     parent_cells_to_intersection_cells
                 )
 
-                assert isinstance(weight_value, sps.spmatrix)  # for mypy
+                assert isinstance(weight_value, (sps.spmatrix, sps.sparray))  # for mypy
                 average_weights = np.ravel(weight_value.sum(axis=1))
                 nonzero = average_weights > 0
                 average_weights[nonzero] = 1 / average_weights[nonzero]

src/porepy/numerics/ad/_ad_parser.py

@@ -139,7 +139,7 @@ def evaluate(
                     result_list[index] = pp.ad.AdArray(
                         res, sps.csr_matrix((res.shape[0], equation_system.num_dofs()))
                     )
-                elif isinstance(res, (sps.spmatrix, np.ndarray)):
+                elif isinstance(res, (sps.spmatrix, sps.sparray, np.ndarray)):
                     # This will cover numpy arrays of higher dimensions (> 1) and sparse
                     # matrices.
                     #
@@ -420,7 +420,7 @@ def _get_error_message(
         msg += "The second argument represents the expression:\n " + msg_1 + nl
 
         # Finally some information on sizes
-        if isinstance(results[0], sps.spmatrix):
+        if isinstance(results[0], (sps.spmatrix, sps.sparray)):
             msg += f"First argument is a sparse matrix of size {results[0].shape}\n"
         elif isinstance(results[0], pp.ad.AdArray):
             msg += (
@@ -430,7 +430,7 @@ def _get_error_message(
         elif isinstance(results[0], np.ndarray):
             msg += f"First argument is a numpy array of size {results[0].size}\n"
 
-        if isinstance(results[1], sps.spmatrix):
+        if isinstance(results[1], (sps.spmatrix, sps.sparray)):
             msg += f"Second argument is a sparse matrix of size {results[1].shape}\n"
         elif isinstance(results[1], pp.ad.AdArray):
             msg += (

src/porepy/numerics/ad/forward_mode.py

@@ -175,7 +175,7 @@ def __add__(self, other: AdType) -> AdArray:
                 raise ValueError("Only 1d numpy arrays can be added to AdArrays")
             return AdArray(self.val + other, self.jac)
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             raise ValueError("Sparse matrices cannot be added to AdArrays")
 
         elif isinstance(other, pp.ad.AdArray):
@@ -268,7 +268,7 @@ def __mul__(self, other: AdType) -> AdArray:
             new_jac = self._diagvec_mul_jac(other)
             return AdArray(new_val, new_jac)
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             raise ValueError(
                 """Sparse matrices cannot be multiplied with  AdArrays elementwise.
                 Did you mean to use the @ operator?
@@ -315,7 +315,7 @@ def __rmul__(self, other: AdType) -> AdArray:
 
         """
 
-        if isinstance(other, (float, sps.spmatrix, np.ndarray, int)):
+        if isinstance(other, (float, sps.spmatrix, sps.sparray, np.ndarray, int)):
             # In these cases, there is no difference between left and right
             # multiplication, so we simply invoke the standard __mul__ function.
             return self.__mul__(other)
@@ -370,7 +370,7 @@ def __pow__(self, other: AdType) -> AdArray:
             new_jac = self._diagvec_mul_jac(other * (self.val ** (other - 1)))
             return AdArray(new_val, new_jac)
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             raise ValueError("Cannot raise AdArrays to power of sparse matrices.")
 
         elif isinstance(other, pp.matrix_operations.ArraySlicer):
@@ -440,7 +440,7 @@ def __rpow__(self, other: AdType) -> AdArray:
             new_jac = self._diagvec_mul_jac((other**self.val) * np.log(other))
             return AdArray(new_val, new_jac)
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             raise ValueError("Cannot raise sparse matrices to the power of Ad arrays.")
 
         elif isinstance(other, pp.ad.AdArray):
@@ -482,7 +482,7 @@ def __truediv__(self, other: AdType) -> AdArray:
             new_jac = self._diagvec_mul_jac(other.astype(float) ** (-1.0))
             return AdArray(new_val, new_jac)
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             raise ValueError("AdArrays cannot be divided by sparse matrices.")
 
         elif isinstance(other, pp.matrix_operations.ArraySlicer):
@@ -510,7 +510,7 @@ def __rtruediv__(self, other: AdType) -> AdArray:
 
         """
 
-        if isinstance(other, (float, int, np.ndarray, sps.spmatrix)):
+        if isinstance(other, (float, int, np.ndarray, sps.spmatrix, sps.sparray)):
             # Divide a float or a numpy array by self is the same as raising self to
             # the power of -1 and multiplying by the float. The multiplication will
             # end upcalling self.__mul__, which will do the right checks for numpy
@@ -545,7 +545,7 @@ class documentation for restrictions on admissible types for this
                 f""" {type(other)}."""
             )
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             # This goes against the way equations should be formulated in the AD
             # framework, variables should not be right-multiplied by anything. Raise
             # a value error to make sure this is not done.
@@ -574,7 +574,7 @@ class documentation for restrictions on admissible types for this
                 f""" {type(other)}."""
             )
 
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             # This is the standard matrix-vector multiplication
             if self.jac.shape[0] != other.shape[1]:
                 raise ValueError(

src/porepy/numerics/ad/functions.py

@@ -416,7 +416,7 @@ def maximum(var_0: FloatType, var_1: FloatType) -> FloatType:
     # Start from var_0, then change entries corresponding to inds.
     max_jac = jacs[0].copy()
 
-    if isinstance(max_jac, sps.spmatrix):
+    if isinstance(max_jac, (sps.spmatrix, sps.sparray)):
         # Enforce csr format, unless the matrix is csc, in which case we keep it.
         if not max_jac.getformat() == "csc":
             max_jac = max_jac.tocsr()

src/porepy/numerics/ad/operators.py

@@ -573,7 +573,7 @@ def value_and_jacobian(
             return AdArray(
                 ad, sps.csr_matrix((ad.shape[0], equation_system.num_dofs()))
             )
-        elif isinstance(ad, (sps.spmatrix, np.ndarray)):
+        elif isinstance(ad, (sps.spmatrix, sps.sparray, np.ndarray)):
             # this case coverse both, dense and sparse matrices returned from
             # discretizations f.e.
             raise NotImplementedError(
@@ -901,7 +901,7 @@ def _parse_other(self, other):
             return [self, Scalar(other)]
         elif isinstance(other, np.ndarray):
             return [self, DenseArray(other)]
-        elif isinstance(other, sps.spmatrix):
+        elif isinstance(other, (sps.spmatrix, sps.sparray)):
             return [self, SparseArray(other)]
         elif isinstance(other, AdArray):
             # This may happen when using nested pp.ad.Function.

src/porepy/numerics/linalg/matrix_operations.py

@@ -626,7 +626,7 @@ def __matmul__(
         # Slice matrix, vector, or AdArray by calling relevant helper methods.
         if isinstance(x, np.ndarray):
             sliced = self._slice_vector(x)
-        elif isinstance(x, sps.spmatrix):
+        elif isinstance(x, (sps.spmatrix, sps.sparray)):
             sliced = self._slice_matrix(x)
         elif isinstance(x, pp.ad.AdArray):
             val = self._slice_vector(x.val)

tests/numerics/ad/test_operators.py

@@ -1068,11 +1068,15 @@ def _get_dense_array(wrapped: bool) -> np.ndarray | pp.ad.DenseArray:
         return array
 
 
-def _get_sparse_array(wrapped: bool) -> sps.spmatrix | pp.ad.SparseArray:
+def _get_sparse_array(
+    wrapped: bool, use_csr_matrix: bool
+) -> sps.spmatrix | sps.sparray | pp.ad.SparseArray:
     """Helper to set a sparse array (scipy sparse array). Expected values in the test
     are hardcoded with respect to this value. The array is either returned as-is, or
     wrapped as an Ad SparseArray."""
-    mat = sps.csr_matrix(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])).astype(float)
+    inner = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    mat = sps.csr_matrix(inner) if use_csr_matrix else sps.csr_array(inner)
+    mat = mat.astype(float)
     if wrapped:
         return pp.ad.SparseArray(mat)
     else:
@@ -1172,7 +1176,7 @@ def _expected_value(
         except ValueError:
             assert op in ["@"]
             return False
-    elif isinstance(var_1, float) and isinstance(var_2, sps.spmatrix):
+    elif isinstance(var_1, float) and isinstance(var_2, (sps.spmatrix, sps.sparray)):
         try:
             # This should fail for all operations expect from multiplication.
             val = eval(f"var_1 {op} var_2")
@@ -1188,13 +1192,15 @@ def _expected_value(
             return False
     elif isinstance(var_1, np.ndarray) and isinstance(var_2, np.ndarray):
         return eval(f"var_1 {op} var_2")
-    elif isinstance(var_1, np.ndarray) and isinstance(var_2, sps.spmatrix):
+    elif isinstance(var_1, np.ndarray) and isinstance(
+        var_2, (sps.spmatrix, sps.sparray)
+    ):
         try:
             return eval(f"var_1 {op} var_2")
         except TypeError:
             assert op in ["/", "**"]
             return False
-    elif isinstance(var_1, sps.spmatrix) and isinstance(var_2, float):
+    elif isinstance(var_1, (sps.spmatrix, sps.sparray)) and isinstance(var_2, float):
         if op == "**":
             # SciPy has implemented a limited version matrix powers to scalars, but not
             # with a satisfactory flexibility. If we try to evaluate the expression, it
@@ -1209,7 +1215,9 @@ def _expected_value(
             return val
         except (ValueError, NotImplementedError):
             return False
-    elif isinstance(var_1, sps.spmatrix) and isinstance(var_2, np.ndarray):
+    elif isinstance(var_1, (sps.spmatrix, sps.sparray)) and isinstance(
+        var_2, np.ndarray
+    ):
         if op == "**":
             # SciPy has implemented a limited version matrix powers to numpy arrays, but
             # not with a satisfactory flexibility. If we try to evaluate the expression,
@@ -1222,10 +1230,12 @@ def _expected_value(
             assert op in ["**"]
             return False
 
-    elif isinstance(var_1, sps.spmatrix) and isinstance(var_2, sps.spmatrix):
+    elif isinstance(var_1, (sps.spmatrix, sps.sparray)) and isinstance(
+        var_2, (sps.spmatrix, sps.sparray)
+    ):
         try:
             return eval(f"var_1 {op} var_2")
-        except (ValueError, TypeError):
+        except (ValueError, TypeError, NotImplementedError):
             assert op in ["**"]
             return False
 
@@ -1434,7 +1444,9 @@ def _expected_value(
             )
             return pp.ad.AdArray(val, jac)
 
-    elif isinstance(var_1, pp.ad.AdArray) and isinstance(var_2, sps.spmatrix):
+    elif isinstance(var_1, pp.ad.AdArray) and isinstance(
+        var_2, (sps.spmatrix, sps.sparray)
+    ):
         return False
     elif isinstance(var_1, sps.spmatrix) and isinstance(var_2, pp.ad.AdArray):
         # This combination is only allowed for matrix-vector products (op = "@")
@@ -1444,6 +1456,14 @@ def _expected_value(
             return pp.ad.AdArray(val, jac)
         else:
             return False
+    elif isinstance(var_1, sps.sparray) and isinstance(var_2, pp.ad.AdArray):
+        # This combination is only allowed for matrix-vector products (op = "@")
+        if op == "@":
+            val = var_1 @ var_2.val
+            jac = var_1 @ var_2.jac
+            return pp.ad.AdArray(val, jac)
+        else:
+            return False
 
     elif isinstance(var_1, pp.ad.AdArray) and isinstance(var_2, pp.ad.AdArray):
         # For this case, var_2 was modified manually to be twice var_1, see comments in
@@ -1525,10 +1545,16 @@ def _expected_value(
             return pp.ad.AdArray(val, jac)
         elif op == "@":
             return False
+    else:
+        raise ValueError(f"Unknown classes: {type(var_1)}, {type(var_2)}.")
 
 
-@pytest.mark.parametrize("var_1", ["scalar", "dense", "sparse", "ad"])
-@pytest.mark.parametrize("var_2", ["scalar", "dense", "sparse", "ad"])
+@pytest.mark.parametrize(
+    "var_1", ["scalar", "dense", "sparse_matrix", "sparse_array", "ad"]
+)
+@pytest.mark.parametrize(
+    "var_2", ["scalar", "dense", "sparse_matrix", "sparse_array", "ad"]
+)
 @pytest.mark.parametrize("op", ["+", "-", "*", "/", "**", "@"])
 @pytest.mark.parametrize("wrapped", [True, False])
 def test_arithmetic_operations_on_ad_objects(
@@ -1572,8 +1598,10 @@ def _var_from_string(v, do_wrap: bool):
             return _get_scalar(do_wrap)
         elif v == "dense":
             return _get_dense_array(do_wrap)
-        elif v == "sparse":
-            return _get_sparse_array(do_wrap)
+        elif v == "sparse_matrix":
+            return _get_sparse_array(do_wrap, use_csr_matrix=True)
+        elif v == "sparse_array":
+            return _get_sparse_array(do_wrap, use_csr_matrix=False)
         elif v == "ad":
             return _get_ad_array(do_wrap)
         else:
@@ -1624,7 +1652,7 @@ def _compare(v1, v2):
             assert np.isclose(v1, v2)
         elif isinstance(v1, np.ndarray):
             assert np.allclose(v1, v2)
-        elif isinstance(v1, sps.spmatrix):
+        elif isinstance(v1, (sps.spmatrix, sps.sparray)):
             assert np.allclose(v1.toarray(), v2.toarray())
         elif isinstance(v1, pp.ad.AdArray):
             assert np.allclose(v1.val, v2.val)