Merge pull request #1326 from keileg/sparse_diags

Faster construction of sparse block matrices

pyproject.toml

@@ -44,7 +44,8 @@ development = [
     "isort",
     "ruff",
     "mypy",
-    "mypy-extensions"
+    "mypy-extensions",
+    "traitlets"
 ]
 testing = [
     "pytest >= 4.6",

src/porepy/grids/mortar_grid.py

@@ -762,7 +762,7 @@ def mortar_to_secondary_avg(self, nd: int = 1) -> sps.spmatrix:
         """
         return sparse_kronecker_product(self._mortar_to_secondary_avg, nd)
 
-    def sign_of_mortar_sides(self, nd: int = 1) -> sps.spmatrix:
+    def sign_of_mortar_sides(self, nd: int = 1) -> sps.dia_matrix:
         """Assign positive or negative weight to the two sides of a mortar grid.
 
         This is needed e.g. to make projection operators into signed projections, for

src/porepy/models/geometry.py

@@ -452,7 +452,9 @@ def local_coordinates(self, subdomains: list[pp.Grid]) -> pp.ad.SparseArray:
                 ].project_tangential_normal(sd.num_cells)
                 for sd in subdomains
             ]
-            local_coord_proj = sps.block_diag(local_coord_proj_list, format="csr")
+            local_coord_proj = pp.matrix_operations.csc_matrix_from_sparse_blocks(
+                local_coord_proj_list
+            )
         else:
             # Also treat no subdomains.
             local_coord_proj = sps.csr_matrix((0, 0))
@@ -620,7 +622,9 @@ def internal_boundary_normal_to_outwards(
                 matrices.append(switcher_int)
 
             # Construct the block diagonal matrix.
-            sign_flipper = pp.ad.SparseArray(sps.block_diag(matrices).tocsr())
+            sign_flipper = pp.ad.SparseArray(
+                pp.matrix_operations.sparse_dia_from_sparse_blocks(matrices)
+            )
         sign_flipper.set_name("Flip_normal_vectors")
         return sign_flipper
 

src/porepy/numerics/ad/_ad_utils.py

@@ -631,4 +631,19 @@ def parse(self, mdg: pp.MixedDimensionalGrid) -> sps.spmatrix:
 
         else:
             # This is a standard discretization; wrap it in a diagonal sparse matrix.
-            return sps.block_diag(mat, format="csr")
+
+            if all([m.format == "dia" for m in mat]):
+                # If all matrices are of dia-format, we can try to use a special method
+                # for forming a sparse dia matrix from the blocks. This is more
+                # efficient than the below csr-based method. However, the matrices can
+                # only be nonzero along their main diagonal. This condition does not
+                # hold true for all dia-matrices, so there is a chance a ValueError may
+                # be raised. If so, we let it pass and fall back to the csr-based
+                # method.
+                try:
+                    return pp.matrix_operations.sparse_dia_from_sparse_blocks(mat)
+                except ValueError:
+                    # Use the csr-based method below.
+                    pass
+
+            return pp.matrix_operations.csr_matrix_from_sparse_blocks(mat)

src/porepy/numerics/ad/grid_operators.py

@@ -351,7 +351,10 @@ def sign_of_mortar_sides(self) -> SparseArray:
                 assert isinstance(intf, pp.MortarGrid)  # Appease mypy
                 mats.append(intf.sign_of_mortar_sides(self.dim))
 
-            block_mat = SparseArray(sps.block_diag(mats), name="SignOfMortarSides")
+            block_mat = SparseArray(
+                pp.matrix_operations.sparse_dia_from_sparse_blocks(mats),
+                name="SignOfMortarSides",
+            )
 
         # Store the matrix for later use and return.A
         self._sign_of_mortar_sides = block_mat
@@ -944,7 +947,7 @@ def parse(self, mdg: pp.MixedDimensionalGrid) -> sps.spmatrix:
 
         """
         mat = [sd.divergence(dim=self.dim) for sd in self.subdomains]
-        matrix = sps.block_diag(mat)
+        matrix = pp.matrix_operations.csr_matrix_from_sparse_blocks(mat)
         return matrix
 
 

src/porepy/numerics/fv/biot.py

@@ -952,7 +952,9 @@ def component_wise_ordering(
             # component-based. This is to build a block diagonal sparse matrix
             # compatible with igrad * rhs_units, that is first all x-component, then y,
             # and z
-            return sps.block_diag([mat[:, ind[i]] for i in range(nd)], format="csr")
+            return pp.matrix_operations.csr_matrix_from_sparse_blocks(
+                [mat[:, ind[i]] for i in range(nd)]
+            )
 
         # Reshape nAlpha component-wise
         nAlpha_grad = component_wise_ordering(nAlpha_grad, nd, sub_cell_index)

src/porepy/numerics/fv/mpfa.py

@@ -427,7 +427,7 @@ def discretize(self, sd: pp.Grid, data: dict) -> None:
             # however this scales poorly with the number of blocks. Instead, use the
             # existing workaround to create a csr matrix based on R, and then pick out
             # the right parts of that one.
-            full_rot_mat = pp.matrix_operations.csr_matrix_from_blocks(
+            full_rot_mat = pp.matrix_operations.csr_matrix_from_dense_blocks(
                 # Replicate R with the right ordering of data elements
                 np.tile(R.ravel(), (1, sd.num_cells)).ravel(),
                 # size of the blocks - this will always be the dimension of the rotation

src/porepy/numerics/fv/tpfa.py

@@ -363,7 +363,7 @@ def _block_diagonal_grid_property_matrix(
         self,
         domains: list[pp.Grid],
         grid_property_getter: Callable[[pp.Grid], Any],
-    ) -> sps.spmatrix:
+    ) -> sps.csr_matrix:
         """Construct mapping matrix for the connectivity between two grids entities.
 
         The mapping matrix is a block diagonal matrix where each block contains 1 where
@@ -385,10 +385,9 @@ def _block_diagonal_grid_property_matrix(
             mat_loc = grid_property_getter(g)
             blocks.append(mat_loc)
 
-        block_matrix = pp.matrix_operations.optimized_compressed_storage(
-            sps.block_diag(blocks)
-        )
-        return block_matrix
+        # All of the calling methods have grid_property_getter return a csr_matrix, so
+        # we use csr_matrix_from_sparse_blocks to construct the block diagonal matrix.
+        return pp.matrix_operations.csr_matrix_from_sparse_blocks(blocks)
 
     def half_face_map(
         self,
@@ -397,7 +396,7 @@ def half_face_map(
         to_entity: Literal["cells", "faces", "half_faces"] = "half_faces",
         dimensions: tuple[int, int] = (1, 1),
         with_sign: bool = False,
-    ) -> sps.spmatrix:
+    ) -> sps.csr_matrix:
         """Mapping between half-faces and cells or faces.
 
         Parameters:
@@ -473,10 +472,7 @@ def get_matrix(g: pp.Grid) -> sps.csr_matrix:
             )
             return mat
 
-        return self._block_diagonal_grid_property_matrix(
-            subdomains,
-            get_matrix,
-        )
+        return self._block_diagonal_grid_property_matrix(subdomains, get_matrix)
 
     def _cell_face_vectors(self, subdomains: list[pp.Grid]) -> sps.spmatrix:
         """Distance between face centers and cell centers.
@@ -526,13 +522,12 @@ def get_c_f_vec_matrix(g: pp.Grid) -> sps.csr_matrix:
             return mat
 
         dist_vec = self._block_diagonal_grid_property_matrix(
-            subdomains,
-            get_c_f_vec_matrix,
+            subdomains, get_c_f_vec_matrix
         )
 
         return dist_vec
 
-    def _normal_vectors(self, subdomains: list[pp.Grid]) -> sps.spmatrix:
+    def _normal_vectors(self, subdomains: list[pp.Grid]) -> sps.csr_matrix:
         """Normal vectors on half-faces, repeated for each dimension.
 
         Parameters:
@@ -589,10 +584,7 @@ def get_matrix(g: pp.Grid) -> sps.csr_matrix:
             )
             return mat
 
-        return self._block_diagonal_grid_property_matrix(
-            subdomains,
-            get_matrix,
-        )
+        return self._block_diagonal_grid_property_matrix(subdomains, get_matrix)
 
     def _cell_face_distances(self, subdomains: list[pp.Grid]) -> np.ndarray:
         """Scalar distance between face centers and cell centers for each half face.
@@ -748,7 +740,4 @@ def get_matrix(g: pp.Grid) -> sps.csr_matrix:
             )
             return mat
 
-        return self._block_diagonal_grid_property_matrix(
-            subdomains,
-            get_matrix,
-        )
+        return self._block_diagonal_grid_property_matrix(subdomains, get_matrix)

src/porepy/numerics/fv/tpsa.py

@@ -780,7 +780,7 @@ def discretize(self, sd: Grid, data: dict) -> None:
             z = np.zeros(nf)
             Rn_data = np.array([[z, n[2], -n[1]], [-n[2], z, n[0]], [n[1], -n[0], z]])
 
-            Rn_hat = pp.matrix_operations.csr_matrix_from_blocks(
+            Rn_hat = pp.matrix_operations.csr_matrix_from_dense_blocks(
                 Rn_data.ravel("F"), nd, nf
             )
             Rn_bar = Rn_hat