Newton solver for blocked problems

[jorgensd]

Describe new/missing feature

Currently, we have no simple interface for solving non-linear blocked problems, such as mixed-dimensional problems.
I would suggest adding something that helps the user with this task, as it is far from trivial.

Suggested user interface

class NewtonSolver:
    max_iterations: int
    bcs: list[dolfinx.fem.DirichletBC]
    A: PETSc.Mat
    b: PETSc.Vec
    J: dolfinx.fem.Form
    b: dolfinx.fem.Form
    dx: PETSc.Vec

    def __init__(
        self,
        F: list[dolfinx.fem.form],
        J: list[list[dolfinx.fem.form]],
        w: list[dolfinx.fem.Function],
        bcs: list[dolfinx.fem.DirichletBC] | None = None,
        max_iterations: int = 5,
        petsc_options: dict[str, str | float | int | None] = None,
        problem_prefix="newton",
    ):
        self.max_iterations = max_iterations
        self.bcs = [] if bcs is None else bcs
        self.b = dolfinx.fem.petsc.create_vector_block(F)
        self.F = F
        self.J = J
        self.A = dolfinx.fem.petsc.create_matrix_block(J)
        self.dx = self.A.createVecLeft()
        self.w = w
        self.x = dolfinx.fem.petsc.create_vector_block(F)

        # Set PETSc options
        opts = PETSc.Options()
        if petsc_options is not None:
            for k, v in petsc_options.items():
                opts[k] = v

        # Define KSP solver
        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
        self._solver.setOperators(self.A)
        self._solver.setFromOptions()

        # Set matrix and vector PETSc options
        self.A.setFromOptions()
        self.b.setFromOptions()

    def solve(self, tol=1e-6, beta=1.0):
        i = 0

        while i < self.max_iterations:
            dolfinx.cpp.la.petsc.scatter_local_vectors(
                self.x,
                [si.x.petsc_vec.array_r for si in self.w],
                [
                    (
                        si.function_space.dofmap.index_map,
                        si.function_space.dofmap.index_map_bs,
                    )
                    for si in self.w
                ],
            )
            self.x.ghostUpdate(
                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
            )

            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
            with self.b.localForm() as b_local:
                b_local.set(0.0)
            constants_L = [form and dolfinx.cpp.fem.pack_constants(form._cpp_object) for form in self.F]
            coeffs_L = [dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self.F]

            constants_a = [
                    [
                        dolfinx.cpp.fem.pack_constants(form._cpp_object)
                        if form is not None
                        else np.array([], dtype=PETSc.ScalarType)
                        for form in forms
                    ]
                    for forms in self.J
                ]

            coeffs_a = [
                    [{} if form is None else dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in forms]
                    for forms in self.J
                ]

            dolfinx.fem.petsc.assemble_vector_block(
                self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0, coeffs_a=coeffs_a, constants_a=constants_a, coeffs_L=coeffs_L, constants_L=constants_L
            )
            self.b.ghostUpdate(
                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
            )

            # Assemble Jacobian
            self.A.zeroEntries()
            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a)

            self._solver.solve(self.b, self.dx)
            # self._solver.view()
            assert (
                self._solver.getConvergedReason() > 0
            ), "Linear solver did not converge"
            offset_start = 0
            for s in self.w:
                num_sub_dofs = (
                    s.function_space.dofmap.index_map.size_local
                    * s.function_space.dofmap.index_map_bs
                )
                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
                )
                s.x.petsc_vec.ghostUpdate(
                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
                )
                offset_start += num_sub_dofs
            # Compute norm of update

            correction_norm = self.dx.norm(0)
            print(f"Iteration {i}: Correction norm {correction_norm}")
            if correction_norm < tol:
                break
            i += 1

    def __del__(self):
        self.A.destroy()
        self.b.destroy()
        self.dx.destroy()
        self._solver.destroy()
        self.x.destroy()

[RemDelaporteMathurin]

There is only one tolerance here as opposed to the current behaviour of NewtonSolver with absolute and relative tolerances. Is this something that can be added?

[jorgensd]

There is only one tolerance here as opposed to the current behaviour of NewtonSolver with absolute and relative tolerances. Is this something that can be added?

Once can customize this method for what kind of tolerances and convergence criteria that should be met.

[jorgensd]

An implementation using the underlying dolfinx.cpp.nls.NewtonSolver is available at:
https://gist.github.com/jorgensd/b016630195fa9f34d67a0fb193ef9453
It highlights that the current Python newton solver is not sufficiently modular to get different create_matrix functions.
Maybe we should add handling of this?