added new newton

src/festim/hydrogen_transport_problem.py

@@ -32,6 +32,188 @@
 
 __all__ = ["HydrogenTransportProblem", "HTransportProblemDiscontinuous"]
 
+from mpi4py import MPI
+
+import numpy as np
+import dolfinx
+from petsc4py import PETSc
+import ufl
+import pytest
+from packaging.version import parse as _v
+
+_alpha_kw: str = "alpha"
+if _v(dolfinx.__version__) < _v("0.9"):
+    _alpha_kw = "scale"
+
+
+class NewNewtonSolver(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: The PDE residual F(u, v)
+            u: The unknown
+            bcs: List of Dirichlet boundary conditions
+            J: UFL representation of the Jacobian (Optional)
+            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.
+        Example::
+            problem = LinearProblem(F, u, [bc0, bc1])
+        """
+
+        super().__init__(u[0].function_space.mesh.comm)
+        if petsc_options is None:
+            # Set PETSc options
+            self.krylov_solver.setOptionsPrefix()
+            opts = PETSc.Options()
+            if petsc_options is not None:
+                for k, v in petsc_options.items():
+                    opts[k] = v
+            self.krylov_solver.setFromOptions()
+
+        self._L = 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:
+            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
+
+        # Create structures for holding arrays and matrix
+        self._b = dolfinx.fem.petsc.create_vector_block(self._L)
+        self._A = dolfinx.fem.petsc.create_matrix_block(self._a)
+        self.dx = dolfinx.fem.petsc.create_vector_block(self._L)
+        self.x = dolfinx.fem.petsc.create_vector_block(self._L)
+
+        self.setJ(self.J, self._A)
+        self.setF(self.F, self._b)
+        self.set_form(self.pre_newton_iteration)
+        self.set_update(self.update_function)
+        self.convergence_criterion = "residual"
+
+    def __del__(self):
+        self._A.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. It
+        Args:
+           x: The vector containing the latest solution
+        """
+
+        # 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)
+
+    @property
+    def L(self) -> list[dolfinx.fem.Form]:
+        """Compiled linear form (the residual form)"""
+        return self._L
+
+    @property
+    def a(self) -> list[list[dolfinx.fem.Form]]:
+        """Compiled bilinear form (the Jacobian form)"""
+        return self._a
+
+    def F(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._L,
+            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 J(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):
+        # 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
+
 
 class HydrogenTransportProblem(problem.ProblemBase):
     """