Custom `Solver`

[KulaginVladimir]

Following the discussion, it might be useful to export some metrics, that can characterise the overall error estimation of the solution (or something else), vs. time step (solver iteration). For example, there is some information on "convergence plots" in COMSOL.

[RemDelaporteMathurin]

Pinging @gabriele-ferrero

[RemDelaporteMathurin]

I'm still struggling to understand exactly what would this convergence plot show. Is it:

at each time step, we record the final error (absolute and relative) of the Newton solver and keep track of it?

or

at each time step, we show a plot with the Newton iterations in X and the errors in Y

?

[KulaginVladimir]

@RemDelaporteMathurin I guess, for a time-dependent problem the final absolute/relative tolerance at each time-step is OK. Another option is to plot 1/time_step vs. time-step if adaptive time-step is used. According to the formulation of the adaptive time-step, it also can indicate that the problem converges.

Another question is related to stationary problems, where there are no time-steps. However, I suppose it is not really needed, but absolute/relative tolerance vs. Newton iterations can be an option.

What do you think?

[RemDelaporteMathurin]

This sounds very good to me and I don't think it would be super hard to implement. Maybe we can start with a prototype in fenics to see how it's done and then think of how to integrate with FESTIM

[KulaginVladimir]

Don't know about not super hard) Possibly, the shown below is not what we need.

As I see from dolfin, NonlinearVariationalSolver (used for solving in FESTIM) creates Newton solver on a call, which in turn returns (_newton_iteration, newton_converged). It appears to me that one or both of these classes have to be change to get the info from Newton solver.

I found hint1 and hint2 on how problem and solver can be adopted for the needs. Maybe, they can help.

Also, there is an example with a CustomSolver that has absolute and relative tolerances as attributes that can be acessed externally. To use this custom class, a custom Problem class is also required, which methods, I suppose, are similar to the ones from here:

import fenics as f
import matplotlib.pyplot as plt

mesh = f.IntervalMesh(1000, 0, 1)

class Left(f.SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and f.near(x[0], 0)

class Right(f.SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and f.near(x[0], 1)

left = Left()
right = Right()
subdomains = f.MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
subdomains.set_all(0)
left.mark(subdomains, 1)
right.mark(subdomains, 2)

ds = f.Measure("ds", domain=mesh, subdomain_data=subdomains)
# ---------- function space ----------------

CG = f.FiniteElement("CG", mesh.ufl_cell(), 1)

V = f.FunctionSpace(mesh, CG)

c = f.Function(V)
c_n = f.Function(V)
v_c = f.TestFunction(V)

# ---------- initial conditions ----------------
c_initial = 1
f.assign(c, f.interpolate(f.Constant(c_initial), V))
f.assign(c_n, f.interpolate(f.Constant(c_initial), V))
# ---------- variational form ----------------
dt = f.Constant(0.001)

F = 0
F += (c - c_n) / dt * v_c * f.dx
F += f.inner(f.grad(c), f.grad(v_c)) * f.dx
J = f.derivative(F, c, f.TrialFunction(c.function_space()))

bcs = [
    f.DirichletBC(V, 0.1, subdomains, 1),
    f.DirichletBC(V, c_initial, subdomains, 2),
]

# ---------- Problem ----------------
class Problem(f.NonlinearProblem):
    def __init__(self, J, F, bcs):
        f.NonlinearProblem.__init__(self)
        self.bilinear_form = J
        self.linear_form = F
        self.bcs = bcs

    def F(self, b, x):
        f.assemble(self.linear_form, tensor=b)
        for bc in self.bcs:
            bc.apply(b, x)

    def J(self, A, x):
        f.assemble(self.bilinear_form, tensor=A)
        for bc in self.bcs:
            bc.apply(A)

problem = Problem(J, F, bcs)
# ------------- Solver ----------------
class CustomSolver(f.NewtonSolver):
    def converged(self, r, problem, iteration):
        if iteration == 0:
            self.r0 = r.norm("l2")
        self.current_absolute_residual = r.norm('l2')
        self.current_relative_residual = r.norm('l2') / self.r0
        return super().converged(r, problem, iteration)

solver = CustomSolver()
# -------------- solve ---------------
fig, ax = plt.subplots()
ax_sub = ax.twinx()

steps = [i for i in range(500)]
abs_res = []
rel_res = []

for i in steps:
    nb_it, converged = solver.solve(problem, c.vector())
    abs_res.append(solver.current_absolute_residual)
    rel_res.append(solver.current_relative_residual)
    c_n.assign(c)

ax.plot(steps, abs_res)
ax_sub.plot(steps, rel_res, color = 'tab:red')
ax.set_xlabel('step')
ax.set_ylabel('Absolute tolerance')
ax_sub.set_ylabel('Relative tolerance', color = 'tab:red')
plt.show()

Resulting plots:

[RemDelaporteMathurin]

@KulaginVladimir this is very interesting! (I think on your second plot you meant to label the y axes with absolute error and relative error right?)

Do we really need to write our own NonLinearProblem class or would writing our own NewtonSolver class suffice? What's the difference between NewtonSolver and NonLinearVariationalSolver?

[KulaginVladimir]

I think on your second plot you meant to label the y axes with absolute error and relative error right?

Yes.

Do we really need to write our own NonLinearProblem class or would writing our own NewtonSolver class suffice?

I think, a custom NewtonSolver is enough if only the Newton method is considered. To use a custom NonLinearVariationalSolver, one would also need to re-write NewtonSolver, since a default NewtonSolver is created:

if (!newton_solver)
{
  dolfin_assert(u->function_space()->mesh());
  MPI_Comm comm = u->function_space()->mesh()->mpi_comm();
  newton_solver = std::make_shared<NewtonSolver>(comm);
}

Though NonLinearVariationalSolver also performs a lot of checks during computation. Noteworthy, the attention should be paid to the properties of a custom NewtonSolver: it must work properly in parallel and be able to accept user-defined properties, allowed to be set in the current version of FESTIM.

What's the difference between NewtonSolver and NonLinearVariationalSolver?

It appears to me that NonLinearVariationalSolver is a method for the life simplification as it sets NewtonSolver and prepares the non-linear problem for solving by NewtonSolver. Probably, someone else has a better knowledge of FEniCS to help with this issue.

Btw, I don't really understand how to directly use fenics.NonlinearProblem. In FEniCSx, it seems to be easier.

[RemDelaporteMathurin]

I understand now that the NewtonSolver object created on call in NonLinearVariationalSolver cannot be retrieved outside the class as it's not stored as an attribute.

Your solution seems to work on this small case.

Maybe the next step is to write a FESTIM script where the solve_once method of HTransportProblem is overwritten to see if it behaves as expected

[KulaginVladimir]

Is Task 3 of Workshop a suitable test case?

[RemDelaporteMathurin]

Sure!

[KulaginVladimir]

So, I guess this can be discussed. I used the script from Task 3 of FESTIM Workshop with minor modifications regarding the simulation model. Then I introduced the previously shown CustomProblem (NonlinearProblem) and CustomSolver that is the same as an ordinary NewtonSolver but has two attributes: current_absolute_error and current_relative_error, that can be accessed externally. After that, CustomHTransportProblem and CustomSimulation classes were introduced to use the CustomSolver class (they minorly differ from parent classes). That is the script that I came up with:

import festim as F
import numpy as np
import fenics as f
import matplotlib.pyplot as plt

# Just to collect errors
global a_er, r_er
a_er = []
r_er = []
kwargs = [
    {
        "ls": "solid", 
        "color": "royalblue", 
        "lw": 3, 
        "label": "festim.Simulation"},
    {
        "ls": "dashed",
        "marker": "x",
        "markevery": 20,
        "color": "red",
        "lw": 1.5,
        "label": "CustomSimulation",
    },
]

fig, ax = plt.subplots()
ax_sub = ax.twinx()

# ---------- Problem ----------------
class CustomProblem(f.NonlinearProblem):
    def __init__(self, J, F, bcs):
        self.bilinear_form = J
        self.linear_form = F
        self.bcs = bcs
        f.NonlinearProblem.__init__(self)

    def F(self, b, x):
        f.assemble(self.linear_form, tensor=b)
        for bc in self.bcs:
            bc.apply(b, x)

    def J(self, A, x):
        f.assemble(self.bilinear_form, tensor=A)
        for bc in self.bcs:
            bc.apply(A)

# ------------- Solver ----------------
class CustomSolver(f.NewtonSolver):
    def converged(self, r, problem, iteration):
        if iteration == 0:
            self.r0 = r.norm("l2")
        self.current_absolute_error = r.norm("l2")
        self.current_relative_error = r.norm("l2") / self.r0
        return super().converged(r, problem, iteration)

# -------------- Custom classes for Simulation  ---------------
class CustomHTransportProblem(F.HTransportProblem):
    def solve_once(self):
        if self.J is None:  # Define the Jacobian
            du = f.TrialFunction(self.u.function_space())
            J = f.derivative(self.F, self.u, du)
        else:
            J = self.J
        problem = CustomProblem(J, self.F, self.bcs)
        solver = CustomSolver()
        solver.parameters["error_on_nonconvergence"] = False
        solver.parameters["absolute_tolerance"] = self.settings.absolute_tolerance
        solver.parameters["relative_tolerance"] = self.settings.relative_tolerance
        solver.parameters["maximum_iterations"] = self.settings.maximum_iterations
        solver.parameters["linear_solver"] = self.settings.linear_solver
        nb_it, converged = solver.solve(problem, self.u.vector())

        # for a plot
        a_er.append(solver.current_absolute_error)
        r_er.append(solver.current_relative_error)

        return nb_it, converged

class CustomSimulation(F.Simulation):
    def initialise(self):
        super().initialise()
        self.h_transport_problem = CustomHTransportProblem(
            self.mobile, self.traps, self.T, self.settings, self.initial_conditions
        )
        self.attribute_boundary_conditions()
        self.attribute_source_terms()
        self.h_transport_problem.initialise(self.mesh, self.materials, self.dt)

# -------------- Simulation  ---------------
for i, model in enumerate([F.Simulation(), CustomSimulation()]):
    my_model = model

    my_model.mesh = F.MeshFromVertices(vertices=np.linspace(0, 3e-4, num=1001))
    my_model.materials = F.Material(id=1, D_0=1.9e-7, E_D=0.2)

    my_model.T = F.Temperature(value=500)
    P_up = 200  # Pa

    my_model.boundary_conditions = [
        F.SievertsBC(surfaces=1, S_0=4.02e21, E_S=1.04, pressure=P_up),
        F.DirichletBC(surfaces=2, value=0, field=0),
    ]
    my_model.settings = F.Settings(
        absolute_tolerance=1e-4,
        relative_tolerance=1e-10,
        maximum_iterations=100,
        final_time=100,  # s
    )
    my_model.dt = F.Stepsize(
        initial_value=0.02,
        stepsize_change_ratio=1.05,
        max_stepsize=lambda t: None if t < 0.1 else 0.1,
    )
# -------------- PP ---------------
    derived_quantities = F.DerivedQuantities([F.HydrogenFlux(surface=2)])
    my_model.exports = [derived_quantities]

    my_model.initialise()
    my_model.run()

    times = derived_quantities.t
    computed_flux = derived_quantities.filter(surfaces=2).data

    D = 1.9e-7 * np.exp(-0.2 / F.k_B / 500)
    S = 4.02e21 * np.exp(-1.04 / F.k_B / 500)

    plt.plot(times, np.abs(computed_flux), **kwargs[i])
plt.ylim(bottom=0)
plt.xlabel("Time (s)")
plt.ylabel("Downstream flux (H/m2/s)")
plt.legend()
plt.show()

Solutions with both Simulation objects match well: What is more important is that the CustomSimulation class provides (a kid of) access to errors:

I also checked that custom classes work with MUMPS.

[RemDelaporteMathurin]

Does it compare well in terms of performance?

[KulaginVladimir]

I think, yes, This example, 10000 mesh elements on my personal PC: F.Simulation - Ellapsed time: 121.2 s, CustomSimulation - Ellapsed time: 118.6 s.

I can perform a more accurate comparison, but I suppose that is not the best problem for tests of the performance.

[RemDelaporteMathurin]

I don't think there's a need to test performance accurately here.

Do you know how would one setup different solver parameters with this method? @allentro was interested

[KulaginVladimir]

Can you, please, specify what parameters? Meanwhile, does this example help?

[RemDelaporteMathurin]

Any parameter really. But I guess it's the same as when you do

        solver.parameters["error_on_nonconvergence"] = False
        solver.parameters["absolute_tolerance"] = self.settings.absolute_tolerance
        solver.parameters["relative_tolerance"] = self.settings.relative_tolerance
        solver.parameters["maximum_iterations"] = self.settings.maximum_iterations
        solver.parameters["linear_solver"] = self.settings.linear_solver

In a nutshell, we are interested in having full control (out-of-the-box) over the solver parameters like linear solver, pre-conditionner, linear-solver tolerances etc.

[KulaginVladimir]

Oh, if this parameters are the NewtonSolver parameters, then I believe thay can be set up similarly, yes. In general, the structure of custom functions can be discussed to cover all the needs.

[RemDelaporteMathurin]

Currently in your code, the solver is redefined at every solve_once call but it could be defined only once (maybe outside of the class) and reused everytime.

[KulaginVladimir]

@RemDelaporteMathurin suggested to consider implementation of a custom Solver that can be set by users for a particular need.

From the private discussion, several points can be highlighted:

• Users should have a possibility to provide their own NewtonSolver or similar object. In particular, this can be used to get the convergence dynamics as was shown earlier in this issue;
• If users do not provide a custom solver, a default one should be created;
• Implementation of a user-defined solver should touch HTransportProblem, HeatTransferProblem, and Traps classes;
• Some refactoring of existing methods might be required to prevent overwriting of the user-defined solver.

@RemDelaporteMathurin, if I missed smth or misunderstood, please correct me.

[Allentro]

This is great!

I was looking into this a little over the weekend following the conversation with had with Tim, @RemDelaporteMathurin!

I had a look at task 3 from the workshop as a sandbox and also didn't see any clear time efficiency for this using 'mumps'. It would be interesting to explore this scaling on more complex problems, though.

Let me know if there is anything I can assist with or test @KulaginVladimir

[RemDelaporteMathurin]

I agree that this problem isn't ideal for testing alternative solvers and preconditioners.

Both mumps and UMFPACK are direct linear solvers, I'd be interested to investigate pre-conditioned iterative linear solvers.

We would need a case that is more expensive with several traps, 3D, a decent amount of cells, several materials, etc

[RemDelaporteMathurin]

@KulaginVladimir I think the script would look like

my_model.initialise()

class CustomSolver:
     .....

custom_solver = CustomSolver()

my_model.h_transport_problem.newton_solver = custom_solver

my_model.run()

[KulaginVladimir]

@RemDelaporteMathurin Right. The solver should be an attribute of HTransportProblem, etc.

I only wonder about NonLinearVariationalSolver that is not really flexible. If we decide to use a custom solver in a similar manner as I showed, than we just have to re-define the problem class, and that's it

[RemDelaporteMathurin]

Yes we just have to move away from NonLinearVariationalSolver and use NewtonSolver instead, right?

[KulaginVladimir]

I think so

We would need a case that is more expensive with several traps, 3D, a decent amount of cells, several materials, etc

Looks like the model from this recent paper is appropriate.

[KulaginVladimir]

@Allentro, @RemDelaporteMathurin, @jhdark

I attempted to implement a custom Newton solver. The current version can be found in my repository. Major changes should allow the users to set the solver preconditioner via F.Settings or provide a manually defined solver based on the fenics.NewtonSolver class for F.HTransportProblem and F.HeatTransferProblem. Though I faced some difficulties.

I performed tests using the model from @RemDelaporteMathurin repository in the transient case with other original settings:

my_model.settings = F.Settings(
    absolute_tolerance=1e4,
    relative_tolerance=1e-5,
    maximum_iterations=15,
    traps_element_type="DG",
    chemical_pot=True,
    transient=True,
    linear_solver="mumps",
)

The results obtained with the custom solver and with FESTIMv1.2.1 (without modifications) correlate in the case when chemical potential conservation is not considered. If chemical_pot = True, the model with a custom solver immediately converges, but concentrations are zero (temperature is not). If one look at errors during NewtonSolver iterations, one can notice that the initial absolute residual is very low: Newton iteration 0: r (abs) = 1.629e+03 (tol = 1.000e+04) r (rel) = 1.000e+00 (tol = 1.000e-05)

In the case of FESTIMv1.2.1, it is: Newton iteration 0: r (abs) = 2.147e+19 (tol = 1.000e+04) r (rel) = 1.000e+00 (tol = 1.000e-05)

I understand that a reduction of absolute_tolerance in settings solves the issue to some extent, but I'd appreciate any help in clarifying this behaviour.

[RemDelaporteMathurin]

@KulaginVladimir do you have a script to reproduce this behaviour?

The reason why concentrations are zero is probably because the newton solver converges in zero iterations. Therefore it "solves nothing" and the functions are unchanged (ie equal to the initial conditions).

The error is a function of the solution fields. When activating chemical_pot=True, the solution field isn't $c_m$ but $\theta = \frac{c_m}{K_S}$. The order of magnitude of the solution fields is therefore different which I believe explains the differences in absolute error