Improvements in petsc usage

src/darsia/measure/wasserstein.py

@@ -320,15 +320,15 @@ def _setup_linear_solver(self) -> None:
         ], f"Linear solver {self.linear_solver_type} not supported."
         assert self.formulation in [
             "full",
-            "flux-reduced",
+            "flux_reduced",
             "pressure",
         ], f"Formulation {self.formulation} not supported."
 
-        if self.linear_solver_type == "ksp":
-            if self.formulation == "flux-reduced":
-                raise ValueError(
-                    "KSP solver only supports for full and pressure formulation."
-                )
+        # if self.linear_solver_type == "ksp":
+        #     if self.formulation == "flux_reduced":
+        #         raise ValueError(
+        #             "KSP solver only supports for full and pressure formulation."
+        #         )
 
         # Setup inrastructure for Schur complement reduction
         if self.formulation == "flux_reduced":
@@ -475,7 +475,9 @@ def setup_ksp_solver(
         """
         # Define CG solver
         self.linear_solver = darsia.linalg.KSP(
-            matrix, field_ises=field_ises, nullspace=nullspace
+            matrix,
+            field_ises=field_ises,
+            nullspace=nullspace,
         )
 
         # Define solver options
@@ -493,15 +495,29 @@ def setup_ksp_solver(
                     "pc_factor_mat_solver_type": "mumps",
                 }
             else:
-                prec = linear_solver_options.get("pc_type", "hypre")
                 self.solver_options = {
                     "ksp_type": approach,
                     # "ksp_monitor_true_residual": None,
                     "ksp_rtol": rtol,
                     "ksp_atol": atol,
                     "ksp_max_it": maxiter,
-                    "pc_type": prec,
+                    "pc_type": "hypre",
                 }
+                if self.grid.dim == 3:
+                    self.solver_options.update(
+                        {
+                            # tuning parameters for the multigrid
+                            # https://mooseframework.inl.gov/releases/moose/2021-09-15/application_development/hypre.html
+                            "pc_hypre_type": "boomeramg",
+                            "pc_hypre_boomeramg_strong_threshold": 0.7,
+                            "pc_hypre_boomeramg_max_iter": 1,
+                            "pc_hypre_boomeramg_agg_nl": 2,
+                            "pc_hypre_boomeramg_interp_type": "ext+i",  # "classic" or "ext+i"
+                        }
+                    )
+                # Include other PETSc options passed from the user
+                other_petsc_options = linear_solver_options.get("petsc_options", {})
+                self.solver_options.update(other_petsc_options)
         else:
             if approach == "direct":
                 self.solver_options = {
@@ -992,7 +1008,9 @@ def _compute_face_weight(self, flat_flux: np.ndarray) -> np.ndarray:
 
             # Determine the l2 norm of the fluxes on the faces
             weighted_face_flux = self._product(harm_avg_face_weights, full_face_flux)
-            norm_weighted_face_flux = np.linalg.norm(weighted_face_flux, 2, axis=1)
+            norm_weighted_face_flux = np.maximum(
+                np.linalg.norm(weighted_face_flux, 2, axis=1), self.regularization
+            )
 
             # Combine weights**2 / |weight * flux| on faces
             face_weights = harm_avg_face_weights**2 / norm_weighted_face_flux
@@ -1028,6 +1046,48 @@ def optimality_conditions(
         )
 
     # ! ---- Solver methods ----
+    def setup_petsc_variables(self, weight: np.ndarray = None):
+        """
+        Instantiate the Petsc variables for the Schur complement
+        """
+        # Setup Petsc operators
+        self.div_petsc = darsia.linalg.numpy_to_petsc(self.div)
+        n = self.grid.num_faces
+
+        # THE FOLLOWING DOES NOT WORK in the matrix-matrix multiplication
+        # PETSc.Mat().createDiagonal(self.inv_weight_laplacian_vec)
+        if weight is None:
+            ones = np.ones(n)
+            scipy_sparse = sps.diags(ones)
+            self.weight_laplacian_vec = darsia.linalg.numpy_to_petsc(ones)
+        else:
+            scipy_sparse = sps.diags(weight)
+            self.weight_laplacian_vec = darsia.linalg.numpy_to_petsc(weight)
+        self.weight_laplacian_matrix = darsia.linalg.numpy_to_petsc(scipy_sparse)
+
+        # We need to create a matrix and set the diagonal
+
+        self.weight_laplacian_matrix.setDiagonal(self.weight_laplacian_vec)
+
+        # We store also the grad in order to use the matrix-matrix-matrix multiplication
+        self.grad_petsc = self.div_petsc.copy()
+        self.grad_petsc.transpose()
+
+        # assign
+        self.laplacian_matrix = self.div_petsc.matMatMult(
+            self.weight_laplacian_matrix, self.grad_petsc
+        )
+
+    def assemble_schur_complement(self, weights: np.ndarray):
+        """
+        Assemble the Schur complement matrix = D^T * W * D
+        in the Petsc matrix self.weight_laplacian_matrix
+        """
+        self.weight_laplacian_vec.setArray(weights)
+        self.weight_laplacian_matrix.setDiagonal(self.weight_laplacian_vec)
+        self.div_petsc.matMatMult(
+            self.weight_laplacian_matrix, self.grad_petsc, result=self.laplacian_matrix
+        )
 
     def linear_solve(
         self,
@@ -1115,30 +1175,66 @@ def linear_solve(
             # Allocate memory for solution
             solution = np.zeros_like(rhs)
 
-            # 1. Reduce flux block
-            (
-                self.reduced_matrix,
-                self.reduced_rhs,
-                self.matrix_flux_inv,
-            ) = self.eliminate_flux(matrix, rhs)
+            if setup_linear_solver:
+                # 1. Reduce flux block
+                if self.linear_solver_type == "ksp":
+                    # Update the Schur complement matrix B^T * A^{-1} * B
+                    weight = 1.0 / matrix.diagonal()[self.flux_slice]
+                    self.matrix_flux_inv = sps.diags(weight)
+                    if not hasattr(self, "div_petsc"):
+                        # setup for the first time
+                        self.setup_petsc_variables(weight=weight)
+                    else:
+                        self.assemble_schur_complement(weights=weight)
+
+                else:
+                    # Compute the inverse of the flux block
+                    self.matrix_flux_inv = self.compute_invA(matrix)
+                    # form the schur complement
+                    self.reduced_matrix = self.compute_Schur_complement(
+                        self.matrix_flux_inv
+                    )
+
+            # create the reduced rhs
+            if self.linear_solver_type == "ksp":
+                self.reduced_rhs = rhs[self.pressure_slice].copy()
+                self.reduced_rhs -= self.div.dot(
+                    self.matrix_flux_inv.dot(rhs[self.flux_slice])
+                )
+            else:
+                self.reduced_rhs = self.compute_reduced_rhs(self.matrix_flux_inv, rhs)
 
-            # 2. Build linear solver for reduced system
             if setup_linear_solver:
+                # 2. Build linear solver for reduced system
                 if self.linear_solver_type == "direct":
                     self.setup_direct_solver(self.reduced_matrix)
                 elif self.linear_solver_type == "amg":
                     self.setup_amg_solver(self.reduced_matrix)
                 elif self.linear_solver_type == "cg":
                     self.setup_cg_solver(self.reduced_matrix)
+                elif self.linear_solver_type == "ksp":
+                    # Setup KSP solver for the first time
+                    if not hasattr(self, "linear_solver"):
+                        kernel = np.ones(self.grid.num_cells)
+                        kernel /= np.linalg.norm(kernel)
+                        self.setup_ksp_solver(self.laplacian_matrix, nullspace=[kernel])
+                    else:
+                        # Just update the matrix
+                        self.linear_solver.ksp.setOperators(self.laplacian_matrix)
 
             # Stop timer to measure setup time
             time_setup = time.time() - tic
 
             # 3. Solve for the pressure and lagrange multiplier
             tic = time.time()
-            solution[self.reduced_system_slice] = self.linear_solver.solve(
-                self.reduced_rhs, **self.solver_options
-            )
+            if self.linear_solver_type == "ksp":
+                pot = self.linear_solver.solve(self.reduced_rhs, **self.solver_options)
+                solution[self.pressure_slice] = pot
+                solution[self.lagrange_multiplier_slice] = 0.0
+            else:
+                solution[self.reduced_system_slice] = self.linear_solver.solve(
+                    self.reduced_rhs, **self.solver_options
+                )
 
             # 4. Compute flux update
             solution[self.flux_slice] = self.compute_flux_update(solution, rhs)
@@ -1174,24 +1270,27 @@ def linear_solve(
                     "Implementation requires solution satisfy the constraint."
                 )
 
-            # 1. Reduce flux block
-            (
-                self.reduced_matrix,
-                self.reduced_rhs,
-                self.matrix_flux_inv,
-            ) = self.eliminate_flux(matrix, rhs)
+            if setup_linear_solver:
+                # 1. Reduce flux block
+                # Compute the inverse of the flux block
+                self.matrix_flux_inv = self.compute_invA(matrix)
+                # form the schur complement
+                self.reduced_matrix = self.compute_Schur_complement(
+                    self.matrix_flux_inv
+                )
 
-            # 2. Reduce to pure pressure system
-            (
-                self.fully_reduced_matrix,
-                self.fully_reduced_rhs,
-            ) = self.eliminate_lagrange_multiplier(
-                self.reduced_matrix,
-                self.reduced_rhs,
+            self.reduced_rhs = self.compute_reduced_rhs(self.matrix_flux_inv, rhs)
+            self.fully_reduced_rhs = self.eliminate_lagrange_multiplier_rhs(
+                self.reduced_rhs
             )
 
-            # 3. Build linear solver for pure pressure system
             if setup_linear_solver:
+                # 2. Reduce to pure pressure system
+                self.fully_reduced_matrix = self.eliminate_lagrange_multiplier_matrix(
+                    self.reduced_matrix
+                )
+
+                # 3. Build linear solver for pure pressure system
                 if self.linear_solver_type == "direct":
                     self.setup_direct_solver(self.fully_reduced_matrix)
 
@@ -1202,20 +1301,15 @@ def linear_solve(
                     self.setup_cg_solver(self.fully_reduced_matrix)
 
                 elif self.linear_solver_type == "ksp":
-                    if not hasattr(self, "linear_solver"):
-                        self.setup_ksp_solver(self.fully_reduced_matrix)
-
-                    if reuse_solver:
-                        self.linear_solver.setup(self.solver_options)
-
-                    # Free memory if solver needs to be re-setup
                     if hasattr(self, "linear_solver"):
                         if not reuse_solver:
                             self.linear_solver.kill()
                             self.setup_ksp_solver(self.fully_reduced_matrix)
                     else:
                         self.setup_ksp_solver(self.fully_reduced_matrix)
 
+                    self.linear_solver.setup(self.solver_options)
+
             # Stop timer to measure setup time
             time_setup = time.time() - tic
 
@@ -1270,6 +1364,31 @@ def eliminate_flux(self, jacobian: sps.csc_matrix, residual: np.ndarray) -> tupl
 
         return reduced_jacobian, reduced_residual, J_inv
 
+    def compute_reduced_rhs(
+        self, J_inv: sps.csc_matrix, residual: np.ndarray
+    ) -> np.ndarray:
+        """
+        Compute the reduced right hand side.
+        """
+        # Gauss elimination on vectors
+        reduced_residual = residual[self.reduced_system_slice].copy()
+        reduced_residual -= self.D.dot(J_inv.dot(residual[self.flux_slice]))
+
+        return reduced_residual
+
+    def compute_invA(self, jacobian: sps.csc_matrix) -> sps.csc_matrix:
+        """
+        Set the inverse of the flux block
+        """
+        return sps.diags(1.0 / jacobian.diagonal()[self.flux_slice])
+
+    def compute_Schur_complement(self, invA: sps.csc_matrix) -> sps.csc_matrix:
+        """
+        Compute the Schur complement, including the lagrange multiplier block
+        """
+        # Build Schur complement wrt flux-block
+        return self.jacobian_subblock + self.D.dot(invA.dot(self.DT))
+
     def eliminate_lagrange_multiplier(
         self, reduced_jacobian, reduced_residual
     ) -> tuple:
@@ -1309,6 +1428,57 @@ def eliminate_lagrange_multiplier(
 
         return self.fully_reduced_jacobian, fully_reduced_residual
 
+    def eliminate_lagrange_multiplier_rhs(self, reduced_residual) -> tuple:
+        """Eliminate the lagrange multiplier from the reduced system.
+
+        Employ a Schur complement/block Gauss elimination approach.
+
+        Args:
+            reduced_residual (np.ndarray): reduced residual
+
+        Returns:
+            np.ndarray: fully reduced residual
+
+        """
+        # Rhs is not affected by Gauss elimination as it is assumed that the residual
+        # is zero in the constrained cell, and the pressure is zero there as well.
+        # If not, we need to do a proper Gauss elimination on the right hand side!
+        if abs(reduced_residual[-1]) > 1e-6:
+            raise NotImplementedError("Implementation requires residual to be zero.")
+        fully_reduced_residual = reduced_residual[
+            self.fully_reduced_system_indices
+        ].copy()
+
+        return fully_reduced_residual
+
+    def eliminate_lagrange_multiplier_matrix(self, reduced_jacobian) -> tuple:
+        """Eliminate the lagrange multiplier from the reduced system.
+
+        Employ a Schur complement/block Gauss elimination approach.
+
+        Args:
+            reduced_jacobian (sps.csc_matrix): reduced jacobian
+            reduced_residual (np.ndarray): reduced residual
+
+        Returns:
+            tuple: fully reduced jacobian, fully reduced residual
+
+        """
+        # Make sure the jacobian is a CSC matrix
+        assert isinstance(
+            reduced_jacobian, sps.csc_matrix
+        ), "Jacobian should be a CSC matrix."
+
+        # Effective Gauss-elimination for the particular case of the lagrange multiplier
+        self.fully_reduced_jacobian.data[:] = np.delete(
+            reduced_jacobian.data.copy(), self.rm_indices
+        )
+        # NOTE: The indices have to be restored if the LU factorization is to be used
+        # FIXME omit if not required
+        self.fully_reduced_jacobian.indices = self.fully_reduced_jacobian_indices.copy()
+
+        return self.fully_reduced_jacobian
+
     def compute_flux_update(self, solution: np.ndarray, rhs: np.ndarray) -> np.ndarray:
         """Compute the flux update from the solution.