Add PETSc solver to the Mohr-Coulomb tutorial

Author: a-latyshev

doc/demo/demo_plasticity_mohr_coulomb.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 # ---
 # jupyter:
 #   jupytext:
@@ -84,7 +85,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from mpltools import annotation  # for slope markers
-from solvers import LinearProblem
+from solvers import SNESProblem
 from utilities import find_cell_by_point
 
 import basix
@@ -655,7 +656,26 @@ def F_ext(v):
 # applied to other plasticity models.
 
 # %%
-external_operator_problem = LinearProblem(J_replaced, -F_replaced, Du, bcs=bcs)
+petsc_options = {
+    "snes_type": "vinewtonrsls",
+    "ksp_type": "preonly",
+    "pc_type": "lu",
+    "pc_factor_mat_solver_type": "mumps",
+    "snes_atol": 1.0e-11,
+    "snes_rtol": 1.0e-11,
+    "snes_max_it": 50,
+    "snes_monitor": "",
+    # "snes_monitor_cancel": "",
+}
+
+def constitutive_update():
+    evaluated_operands = evaluate_operands(F_external_operators)
+    ((_, sigma_new),) = evaluate_external_operators(J_external_operators, evaluated_operands)
+    # Direct access to the external operator values
+    sigma.ref_coefficient.x.array[:] = sigma_new
+
+external_operator_problem = SNESProblem(Du, F_replaced, J_replaced, bcs=bcs, petsc_options=petsc_options, system_update=constitutive_update)
+
 
 # %%
 x_point = np.array([[0, H, 0]])
@@ -665,56 +685,33 @@ def F_ext(v):
 # parameters of the manual Newton method
 max_iterations, relative_tolerance = 200, 1e-8
 
-load_steps_1 = np.linspace(2, 21, 40)
-load_steps_2 = np.linspace(21, 22.75, 20)[1:]
-load_steps = np.concatenate([load_steps_1, load_steps_2])
+# load_steps_1 = np.linspace(2, 21, 40)
+# load_steps_2 = np.linspace(21, 22.75, 20)[1:]
+# load_steps = np.concatenate([load_steps_1, load_steps_2])
+# load_steps = np.concatenate([np.linspace(1.1, 22.3, 100)[:-1]])
+load_steps = np.concatenate([np.linspace(1.2, 22.3, 150)[:-1]])
 num_increments = len(load_steps)
 results = np.zeros((num_increments + 1, 2))
 
 # %% tags=["scroll-output"]
-
 for i, load in enumerate(load_steps):
     q.value = load * np.array([0, -gamma])
-    external_operator_problem.assemble_vector()
 
-    residual_0 = external_operator_problem.b.norm()
-    residual = residual_0
-    Du.x.array[:] = 0
+    # Du.x.array[:] = 0
 
     if MPI.COMM_WORLD.rank == 0:
-        print(f"Load increment #{i}, load: {load}, initial residual: {residual_0}")
-
-    for iteration in range(0, max_iterations):
-        if residual / residual_0 < relative_tolerance:
-            break
-
-        if MPI.COMM_WORLD.rank == 0:
-            print(f"\tOuter Newton iteration #{iteration}")
-        external_operator_problem.assemble_matrix()
-        external_operator_problem.solve(du)
-
-        Du.x.petsc_vec.axpy(1.0, du.x.petsc_vec)
-        Du.x.scatter_forward()
-
-        evaluated_operands = evaluate_operands(F_external_operators)
-        ((_, sigma_new),) = evaluate_external_operators(J_external_operators, evaluated_operands)
-        # Direct access to the external operator values
-        sigma.ref_coefficient.x.array[:] = sigma_new
-
-        external_operator_problem.assemble_vector()
-        residual = external_operator_problem.b.norm()
-
-        if MPI.COMM_WORLD.rank == 0:
-            print(f"\tResidual: {residual}\n")
+        print(f"Load increment #{i}, load: {load}")
 
+    external_operator_problem.solve()
+    
     u.x.petsc_vec.axpy(1.0, Du.x.petsc_vec)
     u.x.scatter_forward()
 
     sigma_n.x.array[:] = sigma.ref_coefficient.x.array
 
     if len(points_on_process) > 0:
         results[i + 1, :] = (-u.eval(points_on_process, cells)[0], load)
-
+        
 print(f"Slope stability factor: {-q.value[-1] * H / c}")
 
 # %% [markdown]

doc/demo/solvers.py

@@ -3,7 +3,7 @@
 import dolfinx.fem.petsc  # noqa: F401
 import ufl
 from dolfinx import fem
-
+from typing import List, Union, Dict, Optional, Callable, Tuple
 
 class LinearProblem:
     def __init__(self, dR: ufl.Form, R: ufl.Form, u: fem.Function, bcs: list[fem.dirichletbc] | None = None):
@@ -63,3 +63,113 @@ def __del__(self):
         self.solver.destroy()
         self.A.destroy()
         self.b.destroy()
+
+class SNESProblem:
+    """Solves a nonlinear problem via PETSc.SNES.
+    F(u) = 0
+    J = dF/du
+    b = assemble(F)
+    A = assemble(J)
+    Ax = b
+    """
+    def __init__(
+        self,
+        u: fem.Function,
+        F: ufl.Form,
+        J: ufl.Form,
+        bcs=[],
+        petsc_options={},
+        prefix=None,
+        system_update: Optional[Callable] = None,
+    ):
+        self.u = u
+        V = self.u.function_space
+        self.comm = V.mesh.comm
+
+        self.F_form = fem.form(F)
+        # J = ufl.derivative(F_form, self.u, ufl.TrialFunction(V))
+        self.J_form = fem.form(J)
+        self.b = fem.petsc.create_vector(self.F_form)
+        self.A = fem.petsc.create_matrix(self.J_form)
+
+        self.bcs = bcs
+
+        # Give PETSc solver options a unique prefix
+        if prefix is None:
+            prefix = "snes_{}".format(str(id(self))[0:4])
+
+        self.prefix = prefix
+        self.petsc_options = petsc_options
+        self.system_update = system_update
+
+        self.solver = self.solver_setup()
+
+    def set_petsc_options(self):
+        # Set PETSc options
+        opts = PETSc.Options()
+        opts.prefixPush(self.prefix)
+
+        for k, v in self.petsc_options.items():
+            opts[k] = v
+
+        opts.prefixPop()
+
+    def solver_setup(self):
+        # Create nonlinear solver
+        snes = PETSc.SNES().create(self.comm)
+
+        snes.setOptionsPrefix(self.prefix)
+        self.set_petsc_options()
+        snes.setFromOptions()
+
+        snes.setFunction(self.F, self.b)
+        snes.setJacobian(self.J, self.A)
+
+        return snes
+
+    def F(self, snes: PETSc.SNES, x: PETSc.Vec, b: PETSc.Vec) -> None:
+        """Assemble the residual F into the vector b.
+
+        Parameters
+        ==========
+        snes: the snes object
+        x: Vector containing the latest solution.
+        b: Vector to assemble the residual into.
+        """
+        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+        x.copy(self.u.x.petsc_vec)
+        self.u.x.scatter_forward()
+
+        #TODO: SNES makes the iteration #0, where it calculates the b norm.
+        #`system_update()` can be omitted in that case
+        self.system_update()
+        
+        with b.localForm() as b_local:
+            b_local.set(0.0)
+        fem.petsc.assemble_vector(b, self.F_form)
+
+        fem.petsc.apply_lifting(b, [self.J_form], [self.bcs], [x], -1.0)
+        b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+        fem.petsc.set_bc(b, self.bcs, x, -1.0)
+        
+    def J(self, snes, x: PETSc.Vec, A: PETSc.Mat, P: PETSc.Mat) -> None:
+        """Assemble the Jacobian matrix.
+
+        Parameters
+        ==========
+        x: Vector containing the latest solution.
+        A: Matrix to assemble the Jacobian into.
+        """
+        A.zeroEntries()
+        fem.petsc.assemble_matrix(A, self.J_form, self.bcs)
+        A.assemble()
+
+    def solve(self,) -> Tuple[int, int]:
+        self.solver.solve(None, self.u.x.petsc_vec)
+        self.u.x.scatter_forward()
+        return (self.solver.getIterationNumber(), self.solver.getConvergedReason())
+
+    def __del__(self):
+        self.solver.destroy()
+        self.b.destroy()
+        self.A.destroy()