Merge pull request #67 from scientificcomputing/dokken/new_blocked_ne…

…wton_solver

Add new newton solver that uses DOLFINx newton sovler

src/scifem/__init__.py

@@ -8,7 +8,7 @@
 from .point_source import PointSource
 from .assembly import assemble_scalar
 from . import xdmf
-from .solvers import NewtonSolver
+from .solvers import BlockedNewtonSolver, NewtonSolver
 from .mesh import create_entity_markers
 
 __all__ = [
@@ -23,6 +23,7 @@
     "dof_to_vertexmap",
     "create_entity_markers",
     "NewtonSolver",
+    "BlockedNewtonSolver",
 ]
 
 

src/scifem/solvers.py

@@ -10,7 +10,7 @@
 import dolfinx
 
 
-__all__ = ["NewtonSolver"]
+__all__ = ["NewtonSolver", "BlockedNewtonSolver"]
 
 logger = logging.getLogger(__name__)
 
@@ -243,3 +243,194 @@ def __del__(self):
         self.dx.destroy()
         self._solver.destroy()
         self.x.destroy()
+
+
+class BlockedNewtonSolver(dolfinx.cpp.nls.petsc.NewtonSolver):
+    def __init__(
+        self,
+        F: list[ufl.form.Form],
+        u: list[dolfinx.fem.Function],
+        bcs: list[dolfinx.fem.DirichletBC] = [],
+        J: list[list[ufl.form.Form]] | None = None,
+        form_compiler_options: dict | None = None,
+        jit_options: dict | None = None,
+        petsc_options: dict | None = None,
+        entity_maps: dict | None = None,
+    ):
+        """Initialize solver for solving a non-linear problem using Newton's method.
+        Args:
+            F: List of PDE residuals [F_0(u, v_0), F_1(u, v_1), ...]
+            u: List of unknown functions u=[u_0, u_1, ...]
+            bcs: List of Dirichlet boundary conditions
+            J: UFL representation of the Jacobian (Optional)
+                Note:
+                    If not provided, the Jacobian will be computed using the
+                    assumption that the test functions come from a ``ufl.MixedFunctionSpace``
+            form_compiler_options: Options used in FFCx
+                compilation of this form. Run ``ffcx --help`` at the
+                command line to see all available options.
+            jit_options: Options used in CFFI JIT compilation of C
+                code generated by FFCx. See ``python/dolfinx/jit.py``
+                for all available options. Takes priority over all
+                other option values.
+            petsc_options:
+                Options passed to the PETSc Krylov solver.
+            entity_maps: Maps used to map entities between different meshes.
+                Only needed if the forms have not been compiled a priori,
+                and has coefficients, test, or trial functions that are defined on different meshes.
+        """
+        # Initialize base class
+        super().__init__(u[0].function_space.mesh.comm)
+
+        # Set PETSc options for Krylov solver
+        if petsc_options is None:
+            # Set PETSc options
+            prefix = self.krylov_solver.getOptionsPrefix()
+            opts = PETSc.Options()
+            if petsc_options is not None:
+                for k, v in petsc_options.items():
+                    opts[k] = f"{prefix}{v}"
+            self.krylov_solver.setFromOptions()
+
+        self._F = dolfinx.fem.form(
+            F,
+            form_compiler_options=form_compiler_options,
+            jit_options=jit_options,
+            entity_maps=entity_maps,
+        )
+
+        # Create the Jacobian matrix, dF/du
+        if J is None:
+            if _v(dolfinx.__version__) < _v("0.9"):
+                raise RuntimeError(
+                    "Automatic computation of Jacobian for blocked problem is only"
+                    + "supported in DOLFINx 0.9 and later"
+                )
+            du = ufl.TrialFunctions(ufl.MixedFunctionSpace(*[ui.function_space for ui in u]))
+            J = ufl.extract_blocks(sum(ufl.derivative(sum(F), u[i], du[i]) for i in range(len(u))))
+        self._a = dolfinx.fem.form(
+            J,
+            form_compiler_options=form_compiler_options,
+            jit_options=jit_options,
+            entity_maps=entity_maps,
+        )
+
+        self._bcs = bcs
+        self._u = u
+        self._pre_solve_callback: Callable[["BlockedNewtonSolver"], None] | None = None
+        self._post_solve_callback: Callable[["BlockedNewtonSolver"], None] | None = None
+
+        # Create structures for holding arrays and matrix
+        self._b = dolfinx.fem.petsc.create_vector_block(self._F)
+        self._J = dolfinx.fem.petsc.create_matrix_block(self._a)
+        self._dx = dolfinx.fem.petsc.create_vector_block(self._F)
+        self._x = dolfinx.fem.petsc.create_vector_block(self._F)
+
+        self.setJ(self._assemble_jacobian, self._J)
+        self.setF(self._assemble_residual, self._b)
+        self.set_form(self._pre_newton_iteration)
+        self.set_update(self._update_function)
+
+    def set_pre_solve_callback(self, callback: Callable[["BlockedNewtonSolver"], None]):
+        """Set a callback function that is called before each Newton iteration."""
+        self._pre_solve_callback = callback
+
+    def set_post_solve_callback(self, callback: Callable[["BlockedNewtonSolver"], None]):
+        """Set a callback function that is called after each Newton iteration."""
+        self._post_solve_callback = callback
+
+    @property
+    def L(self) -> list[dolfinx.fem.Form]:
+        """Compiled linear form (the residual form)"""
+        return self._F
+
+    @property
+    def a(self) -> list[list[dolfinx.fem.Form]]:
+        """Compiled bilinear form (the Jacobian form)"""
+        return self._a
+
+    @property
+    def u(self):
+        return self._u
+
+    def __del__(self):
+        self._J.destroy()
+        self._b.destroy()
+        self._dx.destroy()
+        self._x.destroy()
+
+    def _pre_newton_iteration(self, x: PETSc.Vec) -> None:
+        """Function called before the residual or Jacobian is
+        computed.
+        Args:
+           x: The vector containing the latest solution
+        """
+        if self._pre_solve_callback is not None:
+            self._pre_solve_callback(self)
+        # Scatter previous solution `u=[u_0, ..., u_N]` to `x`; the
+        # blocked version used for lifting
+        dolfinx.cpp.la.petsc.scatter_local_vectors(
+            x,
+            [ui.x.petsc_vec.array_r for ui in self._u],
+            [
+                (
+                    ui.function_space.dofmap.index_map,
+                    ui.function_space.dofmap.index_map_bs,
+                )
+                for ui in self._u
+            ],
+        )
+        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+
+    def _assemble_residual(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
+        """Assemble the residual F into the vector b.
+        Args:
+            x: The vector containing the latest solution
+            b: Vector to assemble the residual into
+        """
+        # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
+        with b.localForm() as b_local:
+            b_local.set(0.0)
+        dolfinx.fem.petsc.assemble_vector_block(
+            b,
+            self._F,
+            self._a,
+            bcs=self._bcs,
+            x0=x,
+            # dolfinx 0.8 compatibility
+            # this is called 'scale' in 0.8, 'alpha' in 0.9
+            **{_alpha_kw: -1.0},
+        )
+        b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
+
+    def _assemble_jacobian(self, x: PETSc.Vec, A: PETSc.Mat) -> None:
+        """Assemble the Jacobian matrix.
+        Args:
+            x: The vector containing the latest solution
+        """
+        # Assemble Jacobian
+        A.zeroEntries()
+        dolfinx.fem.petsc.assemble_matrix_block(A, self._a, bcs=self._bcs)
+        A.assemble()
+
+    def _update_function(self, solver, dx: PETSc.Vec, x: PETSc.Vec):
+        if self._post_solve_callback is not None:
+            self._post_solve_callback(self)
+        # Update solution
+        offset_start = 0
+        for ui in self._u:
+            Vi = ui.function_space
+            num_sub_dofs = Vi.dofmap.index_map.size_local * Vi.dofmap.index_map_bs
+            ui.x.petsc_vec.array_w[:num_sub_dofs] -= (
+                self.relaxation_parameter * dx.array_r[offset_start : offset_start + num_sub_dofs]
+            )
+            ui.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+            offset_start += num_sub_dofs
+
+    def solve(self):
+        """Solve non-linear problem into function. Returns the number
+        of iterations and if the solver converged."""
+        dolfinx.log.set_log_level(dolfinx.log.LogLevel.INFO)
+        n, converged = super().solve(self._x)
+        self._x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+        return n, converged

tests/test_blocked_newton_solver.py

@@ -3,9 +3,10 @@
 import numpy as np
 import dolfinx
 from petsc4py import PETSc
-from scifem import NewtonSolver, assemble_scalar
+from scifem import NewtonSolver, assemble_scalar, BlockedNewtonSolver
 import ufl
 import pytest
+from packaging.version import parse as _v
 
 
 @pytest.mark.parametrize("factor", [1, -1])
@@ -16,18 +17,16 @@ def test_NewtonSolver(factor):
     V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
     Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
 
-    # if dolfinx.__version__ == "0.8.0":
-    v = ufl.TestFunction(V)
-    q = ufl.TestFunction(Q)
-    du = ufl.TrialFunction(V)
-    dp = ufl.TrialFunction(Q)
-    # elif dolfinx.__version__ == "0.9.0.0":
-    # Relies on: https://github.com/FEniCS/ufl/pull/308
-    # W = ufl.MixedFunctionSpace(V, Q)
-    # v, q = ufl.TestFunctions(W)
-    # du, dp = ufl.TrialFunctions(W)
-    # else:
-    #     raise ValueError("Unsupported version of dolfinx")
+    backward_compat = _v(dolfinx.__version__) < _v("0.9")
+    if backward_compat:
+        v = ufl.TestFunction(V)
+        q = ufl.TestFunction(Q)
+        du = ufl.TrialFunction(V)
+        dp = ufl.TrialFunction(Q)
+    else:
+        W = ufl.MixedFunctionSpace(V, Q)
+        v, q = ufl.TestFunctions(W)
+        du, dp = ufl.TrialFunctions(W)
 
     # Set nonzero initial guess
     # Choose initial guess acording to the input factor
@@ -42,17 +41,16 @@ def test_NewtonSolver(factor):
     c1 = dolfinx.fem.Constant(mesh, ftype(0.82))
     F0 = ufl.inner(c0 * u**2, v) * ufl.dx - ufl.inner(u_expr, v) * ufl.dx
     F1 = ufl.inner(c1 * p**2, q) * ufl.dx - ufl.inner(p_expr, q) * ufl.dx
-    # if dolfinx.__version__ == "0.8.0":
-    F = [F0, F1]
-    J = [
-        [ufl.derivative(F0, u, du), ufl.derivative(F0, p, dp)],
-        [ufl.derivative(F1, u, du), ufl.derivative(F1, p, dp)],
-    ]
-    # elif dolfinx.__version__ == "0.9.0.0":
-    # Relies on: https://github.com/FEniCS/ufl/pull/308
-    # F_ = F0 + F1
-    # F = list(ufl.extract_blocks(F_))
-    # J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
+    if backward_compat:
+        F = [F0, F1]
+        J = [
+            [ufl.derivative(F0, u, du), ufl.derivative(F0, p, dp)],
+            [ufl.derivative(F1, u, du), ufl.derivative(F1, p, dp)],
+        ]
+    else:
+        F_ = F0 + F1
+        F = list(ufl.extract_blocks(F_))
+        J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
     petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
     solver = NewtonSolver(F, J, [u, p], max_iterations=25, petsc_options=petsc_options)
     solver.solve()
@@ -64,3 +62,21 @@ def test_NewtonSolver(factor):
     tol = np.finfo(dtype).eps * 1.0e3
     assert assemble_scalar(err_u) < tol
     assert assemble_scalar(err_p) < tol
+
+    # Check consistency with other Newton solver
+    if backward_compat:
+        blocked_solver = dolfinx.nls.NewtonSolver(F, J, [u, p], petsc_options=petsc_options)
+    else:
+        blocked_solver = BlockedNewtonSolver(F, [u, p], J=None, petsc_options=petsc_options)
+
+    dolfinx.log.set_log_level(dolfinx.log.LogLevel.ERROR)
+    u.x.array[:] = factor * 0.1
+    p.x.array[:] = factor * 0.02
+    blocked_solver.convergence_criterion = "incremental"
+    blocked_solver.solve()
+
+    err_u = ufl.inner(u_ex - u, u_ex - u) * ufl.dx
+    err_p = ufl.inner(p_ex - p, p_ex - p) * ufl.dx
+    tol = np.finfo(dtype).eps * 1.0e3
+    assert assemble_scalar(err_u) < tol
+    assert assemble_scalar(err_p) < tol