Vectorized block diagonal mat inversion

src/porepy/numerics/fv/biot.py

@@ -336,7 +336,7 @@ def discretize(self, sd: pp.Grid, sd_data: dict) -> None:
 
         eta: float = parameter_dictionary.get("mpsa_eta", pp.fvutils.determine_eta(sd))
         inverter: Literal["python", "numba"] = parameter_dictionary.get(
-            "inverter", "numba"
+            "inverter", "python"
         )
 
         alpha: float = parameter_dictionary["biot_alpha"]

src/porepy/numerics/fv/mpfa.py

@@ -86,7 +86,7 @@ def discretize(self, sd: pp.Grid, data: dict) -> None:
             - mpfa_eta (``float``): Optional. Range [0, 1). Location of pressure
                 continuity point. If not given, porepy tries to set an optimal value.
             - mpfa_inverter (``str``): Optional. Inverter to apply for local problems.
-                Can take values 'numba' (default), or 'python'.
+                Can take values 'python' (default), or 'numba'.
             - partition_arguments (``dict``): Optional. Arguments to control the number
                 of subproblems used to discretize the grid. Can be either the target
                 maximal memory use (controlled by keyword 'max_memory' in
@@ -151,7 +151,7 @@ def discretize(self, sd: pp.Grid, data: dict) -> None:
 
         eta: Optional[float] = parameter_dictionary.get("mpfa_eta", None)
         inverter: Literal["numba", "python"] = parameter_dictionary.get(
-            "mpfa_inverter", "numba"
+            "mpfa_inverter", "python"
         )
 
         # Control of the number of subdomanis.
@@ -574,7 +574,7 @@ def _flux_discretization(
         sd: pp.Grid,
         k: pp.SecondOrderTensor,
         bnd: pp.BoundaryCondition,
-        inverter: Literal["python", "numba"],
+        inverter: Optional[Literal["python", "numba"]] = None,
         ambient_dimension: Optional[int] = None,
         eta: Optional[float] = None,
     ) -> tuple[

src/porepy/numerics/fv/mpsa.py

@@ -143,7 +143,7 @@ def discretize(self, sd: pp.Grid, data: dict) -> None:
                 value. This value is set to all subfaces, except the boundary (where,
                 0 is used).
             - inverter (``str``): Optional. Inverter to apply for local problems.
-                Can take values 'numba' (default), or 'python'.
+                Can take values 'python' (default), or 'numba'.
             - partition_arguments (``dict``): Optional. Arguments to control the number
                 of subproblems used to discretize the grid. Can be either the target
                 maximal memory use (controlled by keyword 'max_memory' in
@@ -181,7 +181,7 @@ def discretize(self, sd: pp.Grid, data: dict) -> None:
         hf_eta: Optional[float] = parameter_dictionary.get("reconstruction_eta", None)
 
         inverter: Literal["python", "numba"] = parameter_dictionary.get(
-            "inverter", "numba"
+            "inverter", "python"
         )
 
         # Control of the number of subdomanis.
@@ -514,7 +514,7 @@ def _stress_discretization(
         constit: pp.FourthOrderTensor,
         bound: pp.BoundaryConditionVectorial,
         eta: Optional[float] = None,
-        inverter: Literal["numba", "python"] = "numba",
+        inverter: Optional[Literal["python", "numba"]] = None,
         hf_disp: bool = False,
         hf_eta: Optional[float] = None,
     ) -> tuple[sps.spmatrix, sps.spmatrix, sps.spmatrix, sps.spmatrix]:
@@ -577,7 +577,7 @@ def _stress_discretization(
             eta: Parameter controlling continuity of displacement. If None, a default
                 value will be used, adapted to the grid type.
             inverter: Method to be used for inversion of local systems. Options are
-                'numba' (default) and 'python'.
+                'python' and 'numba'.
             hf_disp: If True, the displacement will be represented by subface values
                 instead of cell values. This is not recommended, but kept for the time
                 being for legacy reasons.
@@ -768,7 +768,7 @@ def _create_inverse_gradient_matrix(
         subcell_topology: pp.fvutils.SubcellTopology,
         bound_exclusion: pp.fvutils.ExcludeBoundaries,
         eta: float,
-        inverter: Literal["python", "numba"],
+        inverter: Optional[Literal["python", "numba"]] = None,
     ) -> tuple[sps.spmatrix, sps.spmatrix, np.ndarray]:
         """
         This is the function where the real discretization takes place. It contains
@@ -1639,7 +1639,7 @@ def _inverse_gradient(
         nno_unique: np.ndarray,
         bound_exclusion: pp.fvutils.ExcludeBoundaries,
         nd: int,
-        inverter: str,
+        inverter: Optional[str] = None,
     ) -> sps.spmatrix:
         """Invert local system to compute the subcell gradient operator.
 

src/porepy/numerics/linalg/matrix_operations.py

@@ -555,9 +555,6 @@ def invert_diagonal_blocks_python(a: sps.spmatrix, sz: np.ndarray) -> np.ndarray
         """
         Invert block diagonal matrix using pure python code.
 
-        The implementation is slow for large matrices, consider to use the
-        numba-accelerated method invert_invert_diagagonal_blocks_numba instead
-
         Parameters
         ----------
         a: sps.crs-matrix, to be inverted
@@ -567,29 +564,58 @@ def invert_diagonal_blocks_python(a: sps.spmatrix, sz: np.ndarray) -> np.ndarray
         -------
         inv_a inverse matrix
         """
-        v = np.zeros(np.sum(np.square(sz)))
-        p1 = 0
-        p2 = 0
-        for b in range(sz.size):
-            n = sz[b]
-            n2 = n * n
-            i = p1 + np.arange(n + 1)
-            # Picking out the sub-matrices here takes a lot of time.
-            v[p2 + np.arange(n2)] = np.linalg.inv(
-                a[i[0] : i[-1], i[0] : i[-1]].A
-            ).ravel()
-            p1 = p1 + n
-            p2 = p2 + n2
+
+        # This function only supports CSR format.
+        if not sps.isspmatrix_csr(a):
+            raise TypeError("Sparse array type not implemented: ", type(a))
+
+        # Retrieve global indices.
+        row, col = a.nonzero()
+
+        # Since blocks should be squared, this statement construct global block indices
+        idx_blocks = np.cumsum([0] + list(sz))
+
+        # Maps block indexation to non-zero indexation
+        idx_nnz = np.searchsorted(row, idx_blocks)
+
+        # Helper function for retrieving blocks
+        def retrieve_block(ib: int) -> np.ndarray:
+            """
+            Parameters
+            ----------
+            id: the block index
+
+            Returns
+            -------
+            dense_block: the dense block
+            """
+
+            # Initialize block
+            dense_block = np.zeros((sz[ib], sz[ib]))
+
+            # Build local from global indexation by subtracting the stride idx_blocks[ib]
+            lr = row[idx_nnz[ib] : idx_nnz[ib + 1]] - idx_blocks[ib]
+            lc = col[idx_nnz[ib] : idx_nnz[ib + 1]] - idx_blocks[ib]
+
+            # Loopless assignment of nonzero entries
+            dense_block[lr, lc] = a.data[idx_nnz[ib] : idx_nnz[ib + 1]]
+            return dense_block
+
+        # Maps retrieve_block to indexation
+        iterator_blocks = map(retrieve_block, range(sz.size))
+
+        # Maps linalg.inv to blocks
+        iterator_block_inverses = map(np.linalg.inv, iterator_blocks)
+
+        # Maps ravel to block_inverses, performs the actual computation and concatenates
+        v = np.concatenate(list(map(np.ravel, iterator_block_inverses)))
         return v
 
     def invert_diagonal_blocks_numba(a: sps.csr_matrix, size: np.ndarray) -> np.ndarray:
         """
         Invert block diagonal matrix by invoking numba acceleration of a simple
         for-loop based algorithm.
 
-        This approach should be more efficient than the related method
-        invert_diagonal_blocks_python for larger problems.
-
         Parameters
         ----------
         a : sps.csr matrix
@@ -603,8 +629,7 @@ def invert_diagonal_blocks_numba(a: sps.csr_matrix, size: np.ndarray) -> np.ndar
             import numba
         except ImportError:
             raise ImportError("Numba not available on the system")
-        # Sort matrix storage before pulling indices and data
-        a.sorted_indices()
+
         ptr = a.indptr
         indices = a.indices
         dat = a.data
@@ -694,15 +719,15 @@ def inv_python(indptr, ind, data, sz):
 
     # Remove blocks of size 0
     s = s[s > 0]
-    # Variable to check if we have tried and failed with numba
-    if method == "numba" or method is None:
+    # Select python vectorized function
+    if method == "python" or method is None:
+        inv_vals = invert_diagonal_blocks_python(mat, s)
+    # Select numba function
+    elif method == "numba":
         try:
             inv_vals = invert_diagonal_blocks_numba(mat, s)
         except np.linalg.LinAlgError:
             raise ValueError("Error in inversion of local linear systems")
-    # Variable to check if we should fall back on python
-    elif method == "python":
-        inv_vals = invert_diagonal_blocks_python(mat, s)
     else:
         raise ValueError(f"Unknown type of block inverter {method}")
     ia = block_diag_matrix(inv_vals, s)