Custom solver

festim/__init__.py

@@ -102,6 +102,7 @@
 from .concentration.concentration import Concentration
 from .initial_condition import InitialCondition
 from .concentration.mobile import Mobile
+from .nonlinear_problem import Problem
 from .concentration.theta import Theta
 
 from .concentration.traps.trap import Trap

festim/concentration/traps/extrinsic_trap.py

@@ -1,4 +1,6 @@
 from festim import Trap, as_constant_or_expression
+from fenics import NewtonSolver, MPI
+import warnings
 
 
 class ExtrinsicTrapBase(Trap):
@@ -14,6 +16,7 @@ def __init__(
         relative_tolerance=1e-10,
         maximum_iterations=30,
         linear_solver=None,
+        preconditioner=None,
         **kwargs,
     ):
         """Inits ExtrinsicTrap
@@ -36,18 +39,48 @@ def __init__(
                 If None, the default fenics linear solver will be used ("umfpack").
                 More information can be found at: https://fenicsproject.org/pub/tutorial/html/._ftut1017.html.
                 Defaults to None.
+            preconditioner (str, optional): preconditioning method for the newton solver,
+                options can be veiwed by print(list_krylov_solver_preconditioners()).
+                Defaults to "default".
         """
         super().__init__(k_0, E_k, p_0, E_p, materials, density=None, id=id)
         self.absolute_tolerance = absolute_tolerance
         self.relative_tolerance = relative_tolerance
         self.maximum_iterations = maximum_iterations
         self.linear_solver = linear_solver
+        self.preconditioner = preconditioner
 
+        self.newton_solver = None
         for name, val in kwargs.items():
             setattr(self, name, as_constant_or_expression(val))
         self.density_previous_solution = None
         self.density_test_function = None
 
+    @property
+    def newton_solver(self):
+        return self._newton_solver
+
+    @newton_solver.setter
+    def newton_solver(self, value):
+        if value is None:
+            self._newton_solver = value
+        elif isinstance(value, NewtonSolver):
+            if self._newton_solver:
+                print("Settings for the Newton solver will be overwritten")
+            self._newton_solver = value
+        else:
+            raise TypeError("accepted type for newton_solver is fenics.NewtonSolver")
+
+    def define_newton_solver(self):
+        """Creates the Newton solver and sets its parameters"""
+        self.newton_solver = NewtonSolver(MPI.comm_world)
+        self.newton_solver.parameters["error_on_nonconvergence"] = False
+        self.newton_solver.parameters["absolute_tolerance"] = self.absolute_tolerance
+        self.newton_solver.parameters["relative_tolerance"] = self.relative_tolerance
+        self.newton_solver.parameters["maximum_iterations"] = self.maximum_iterations
+        self.newton_solver.parameters["linear_solver"] = self.linear_solver
+        self.newton_solver.parameters["preconditioner"] = self.preconditioner
+
 
 class ExtrinsicTrap(ExtrinsicTrapBase):
     """

festim/concentration/traps/traps.py

@@ -109,26 +109,23 @@ def define_variational_problem_extrinsic_traps(self, dx, dt, T):
                 self.extrinsic_formulations.append(trap.form_density)
         self.sub_expressions.extend(expressions_extrinsic)
 
+    def define_newton_solver_extrinsic_traps(self):
+        for trap in self:
+            if isinstance(trap, festim.ExtrinsicTrapBase):
+                trap.define_newton_solver()
+
     def solve_extrinsic_traps(self):
         for trap in self:
             if isinstance(trap, festim.ExtrinsicTrapBase):
                 du_t = f.TrialFunction(trap.density[0].function_space())
                 J_t = f.derivative(trap.form_density, trap.density[0], du_t)
-                problem = f.NonlinearVariationalProblem(
-                    trap.form_density, trap.density[0], [], J_t
-                )
-                solver = f.NonlinearVariationalSolver(problem)
-                solver.parameters["newton_solver"][
-                    "absolute_tolerance"
-                ] = trap.absolute_tolerance
-                solver.parameters["newton_solver"][
-                    "relative_tolerance"
-                ] = trap.relative_tolerance
-                solver.parameters["newton_solver"][
-                    "maximum_iterations"
-                ] = trap.maximum_iterations
-                solver.parameters["newton_solver"]["linear_solver"] = trap.linear_solver
-                solver.solve()
+                problem = festim.Problem(J_t, trap.form_density, [])
+
+                f.begin(
+                    "Solving nonlinear variational problem."
+                )  # Add message to fenics logs
+                trap.newton_solver.solve(problem, trap.density[0].vector())
+                f.end()
 
     def update_extrinsic_traps_density(self):
         for trap in self:

festim/h_transport_problem.py

@@ -1,5 +1,6 @@
 from fenics import *
 import festim
+import warnings
 
 
 class HTransportProblem:
@@ -22,6 +23,7 @@ class HTransportProblem:
             ct2, ...)
         v (fenics.TestFunction): the test function
         u_n (fenics.Function): the "previous" function
+        newton_solver (fenics.NewtonSolver): Newton solver for solving the nonlinear problem
         bcs (list): list of fenics.DirichletBC for H transport
     """
 
@@ -36,13 +38,29 @@ def __init__(self, mobile, traps, T, settings, initial_conditions) -> None:
         self.u = None
         self.v = None
         self.u_n = None
+        self.newton_solver = None
 
         self.boundary_conditions = []
         self.bcs = None
         self.V = None
         self.V_CG1 = None
         self.expressions = []
 
+    @property
+    def newton_solver(self):
+        return self._newton_solver
+
+    @newton_solver.setter
+    def newton_solver(self, value):
+        if value is None:
+            self._newton_solver = value
+        elif isinstance(value, NewtonSolver):
+            if self._newton_solver:
+                print("Settings for the Newton solver will be overwritten")
+            self._newton_solver = value
+        else:
+            raise TypeError("accepted type for newton_solver is fenics.NewtonSolver")
+
     def initialise(self, mesh, materials, dt=None):
         """Assigns BCs, create suitable function space, initialise
         concentration fields, define variational problem
@@ -59,6 +77,7 @@ def initialise(self, mesh, materials, dt=None):
             self.mobile.volume_markers = mesh.volume_markers
             self.mobile.T = self.T
         self.attribute_flux_boundary_conditions()
+
         # Define functions
         self.define_function_space(mesh)
         self.initialise_concentrations()
@@ -71,12 +90,14 @@ def initialise(self, mesh, materials, dt=None):
             self.mobile.create_form_post_processing(self.V_DG1, materials, mesh.dx)
 
         self.define_variational_problem(materials, mesh, dt)
+        self.define_newton_solver()
 
         # Boundary conditions
         print("Defining boundary conditions")
         self.create_dirichlet_bcs(materials, mesh)
         if self.settings.transient:
             self.traps.define_variational_problem_extrinsic_traps(mesh.dx, dt, self.T)
+            self.traps.define_newton_solver_extrinsic_traps()
 
     def define_function_space(self, mesh):
         """Creates a suitable function space for H transport problem
@@ -201,6 +222,22 @@ def define_variational_problem(self, materials, mesh, dt=None):
         self.F = F
         self.expressions = expressions
 
+    def define_newton_solver(self):
+        """Creates the Newton solver and sets its parameters"""
+        self.newton_solver = NewtonSolver(MPI.comm_world)
+        self.newton_solver.parameters["error_on_nonconvergence"] = False
+        self.newton_solver.parameters["absolute_tolerance"] = (
+            self.settings.absolute_tolerance
+        )
+        self.newton_solver.parameters["relative_tolerance"] = (
+            self.settings.relative_tolerance
+        )
+        self.newton_solver.parameters["maximum_iterations"] = (
+            self.settings.maximum_iterations
+        )
+        self.newton_solver.parameters["linear_solver"] = self.settings.linear_solver
+        self.newton_solver.parameters["preconditioner"] = self.settings.preconditioner
+
     def attribute_flux_boundary_conditions(self):
         """Iterates through self.boundary_conditions, checks if it's a FluxBC
         and its field is 0, and assign fluxes to self.mobile
@@ -239,7 +276,6 @@ def update(self, t, dt):
             t (float): the current time (s)
             dt (festim.Stepsize): the stepsize
         """
-
         festim.update_expressions(self.expressions, t)
 
         converged = False
@@ -264,28 +300,16 @@ def solve_once(self):
             int, bool: number of iterations for reaching convergence, True if
                 converged else False
         """
-
         if self.J is None:  # Define the Jacobian
             du = TrialFunction(self.u.function_space())
             J = derivative(self.F, self.u, du)
         else:
             J = self.J
-        problem = NonlinearVariationalProblem(self.F, self.u, self.bcs, J)
-        solver = NonlinearVariationalSolver(problem)
-        solver.parameters["newton_solver"]["error_on_nonconvergence"] = False
-        solver.parameters["newton_solver"][
-            "absolute_tolerance"
-        ] = self.settings.absolute_tolerance
-        solver.parameters["newton_solver"][
-            "relative_tolerance"
-        ] = self.settings.relative_tolerance
-        solver.parameters["newton_solver"][
-            "maximum_iterations"
-        ] = self.settings.maximum_iterations
-        solver.parameters["newton_solver"][
-            "linear_solver"
-        ] = self.settings.linear_solver
-        nb_it, converged = solver.solve()
+        problem = festim.Problem(J, self.F, self.bcs)
+
+        begin("Solving nonlinear variational problem.")  # Add message to fenics logs
+        nb_it, converged = self.newton_solver.solve(problem, self.u.vector())
+        end()
 
         return nb_it, converged
 

festim/nonlinear_problem.py

@@ -0,0 +1,31 @@
+import fenics as f
+
+
+class Problem(f.NonlinearProblem):
+    """
+    Class to set up a nonlinear variational problem (F(u, v)=0) to solve
+    by the Newton method based on the form of the variational problem, the Jacobian
+    form of the variational problem, and the boundary conditions
+
+    Args:
+        J (ufl.Form): the Jacobian form of the variational problem
+        F (ufl.Form): the form of the variational problem
+        bcs (list): list of fenics.DirichletBC
+    """
+
+    def __init__(self, J, F, bcs):
+        self.jacobian_form = J
+        self.residual_form = F
+        self.bcs = bcs
+        self.assembler = f.SystemAssembler(
+            self.jacobian_form, self.residual_form, self.bcs
+        )
+        f.NonlinearProblem.__init__(self)
+
+    def F(self, b, x):
+        """Assembles the RHS in Ax=b and applies the boundary conditions"""
+        self.assembler.assemble(b, x)
+
+    def J(self, A, x):
+        """Assembles the LHS in Ax=b and applies the boundary conditions"""
+        self.assembler.assemble(A)

festim/settings.py

@@ -22,9 +22,12 @@ class Settings:
             the formulation will be computed only once at the beggining.
             Else it will be computed at each time step. Defaults to True.
         linear_solver (str, optional): linear solver method for the newton solver,
-            options can be veiwed by print(list_linear_solver_methods()).
+            options can be viewed by print(list_linear_solver_methods()).
             More information can be found at: https://fenicsproject.org/pub/tutorial/html/._ftut1017.html.
             Defaults to None, for the newton solver this is: "umfpack".
+        preconditioner (str, optional): preconditioning method for the newton solver,
+            options can be viewed by print(list_krylov_solver_preconditioners()).
+            Defaults to "default".
 
     Attributes:
         transient (bool): transient or steady state sim
@@ -40,6 +43,7 @@ class Settings:
         traps_element_type (str): Finite element used for traps.
         update_jacobian (bool):
         linear_solver (str): linear solver method for the newton solver
+        precondtitioner (str): preconditioning method for the newton solver
     """
 
     def __init__(
@@ -54,6 +58,7 @@ def __init__(
         traps_element_type="CG",
         update_jacobian=True,
         linear_solver=None,
+        preconditioner="default",
     ):
         # TODO maybe transient and final_time are redundant
         self.transient = transient
@@ -68,3 +73,4 @@ def __init__(
         self.traps_element_type = traps_element_type
         self.update_jacobian = update_jacobian
         self.linear_solver = linear_solver
+        self.preconditioner = preconditioner