Add support for objective function in SNES

[stefanozampini]

Description

Add support for objective functions in NonlinearVariationalProblem and NonlinearVariationalSolver. Needs https://gitlab.com/petsc/petsc/-/merge_requests/9331 to be merged in PETSc

I believe there's a flaw in the pre_apply_bcs=True model. Setting to zero the bc contribution to the update solution using https://gitlab.com/petsc/petsc/-/merge_requests/9340 make gamg tests pass

[pbrubeck]

In order to get the energy at the coarse levels in MG/FAS we just need to add a tiny amount of logic in mg/ufl_utils.py.

I can add it for you, if you don't mind.

[stefanozampini]

In order to get the energy at the coarse levels in MG/FAS we just need to add a tiny amount of logic in mg/ufl_utils.py.

I can add it for you, if you don't mind.

Go ahead

[pbrubeck]

Since the PETSc MR is targeting release, and this PR is incremental and non API-breaking, we can merge this one into our release

[pbrubeck]

This needs to be mentioned in the manual.

It'd be also good to add this to solve( ..., objective=E)

[stefanozampini]

This needs to be mentioned in the manual.

It'd be also good to add this to solve( ..., objective=E)

If you tell me where should I add this information, I'll do it

[pbrubeck]

M had been used before as the objective, however we were completely ignoring it

[pbrubeck]

Suggested change

	# Apply DirichletBC on the solution
	for bc in ctx._problem.dirichlet_bcs():
	bc.apply(ctx._x)
	# Apply DirichletBC on the solution
	if ctx.pre_apply_bcs:
	for bc in ctx._problem.dirichlet_bcs():
	bc.apply(ctx._x)

[stefanozampini]

I don't think what you are proposing is correct, but I also don't think what I have in the current PR is 100% correct. In all my experience with nonlinear solvers, you always inject the correct BCs whenever you evaluate objective functions or residuals. Then, in case of residuals, you modify afterward the part of the vector that is associated with bc, because there you solve a different set of equations, i.e.

F(u) = 0 on \Omega
u - g = 0 on \partial\Omega

As a byproduct, if you do Newton, the boundary conditions will be exact after the first full step (that may not be the first Newton step if you use linesearch). See e.g. how I do it in MFEM https://github.com/mfem/mfem/blob/master/linalg/petsc.cpp#L4996

If ctx._x is supposed to be immutable during each phases of a given SNES step, then you need to rethink the firedrake logic a bit

[pbrubeck]

Is this safe? ctx._x is the solution according to Firedrake (but SNES has a different one) can we treat ctx._x as a buffer?

[stefanozampini]

I don't know, I copied that from the residual evaluation routine. At the end of solve, you have work.copy(u), with u = problem.u_restrict.dat.vec and work the Vec SNES is using for the solution. From the point of view of the solve, this is ok. But if your ctx._x is supposed to be immutable for a given Newton step, then overwriting it can be a problem, yes.

[JHopeCollins]

Can form_objective be called within a single Newton step? i.e. by the KSP? If so then I do not think that this is safe because for anything matrix-free we assume that ctx._x holds the current state.

I think you just need to save the state in ctx._x to a temporary buffer at entry to this function and restore it at exit.

[stefanozampini]

In principle it should not happen, but I agree with you better be safe. Make a suggestion in the diff and I will merge it in

[JHopeCollins]

Will this overwrite an existing user post-solve callback? Or will they both be called in that case?

[pbrubeck]

This is a different one. The existing ones are SNES related (pre/post Jacobian/Function)

[JHopeCollins]

Not the ones passed to the nlvs, the ones on the actual ksp. A user might have set one, or a particular ksp type might define one itself

[stefanozampini]

This is a temporary fix to firedrake's model of pre_apply_bcs=True, which I do not fully understand (or agree with :-) ). May be removed if we understand what's going wrong with the bc model and inexact solves

[stefanozampini]

Not the ones passed to the nlvs the ones on the actual ksp
A user might have set one, or a particular ksp type might define one itself

This is set as the default when firedrake creates the SNES. Users can do whatever they want afterward. Add their own, or change the KSP. And no, KSPs don't have their own since each KSP implements its own Solve; these are user-specific callbacks

As I said, this is temporary and can be removed once we understand what is going wrong with pre_apply_bcs=True

[stefanozampini]

Yes, it should.

[stefanozampini]

I agree, the pre_apply_bcs=True is flawed in the SNES sense (does not let you linearize around a state that does not satisfy the BCs), but also in the KSP sense (there is no way to make sure that the PC would not introduce coupling between the BC and the interior, i.e. the PC needs to invert the identity block exactly). The short-circuit pre_apply_bcs=restrict, fixes the KSP issue, but not the SNES one.

What do you mean. With restrict you do not have BC dofs at all in the system, so it should not cause problems either to KSP or SNES, since you will just solve the proper linear system only approximatively

[pbrubeck]

For option 2, we would change the default to pre_apply_bcs=restrict=False and deprecate the argument pre_apply_bcs from the next release onwards, but still accept two separate arguments for restrict and pre_apply_bcs and make sure that they are equal. We then refactor the internal code to have a unified restrict/pre_apply_bcs flag. Finally, remove the pre_apply_bcs argument on the following release.

[pbrubeck]

What do you mean. With restrict you do not have BC dofs at all in the system, so it should not cause problems either to KSP or SNES, since you will just solve the proper linear system only approximatively

The SNES issue is unavoidable and inherent to the pre_apply_bcs=True model. It might result in NaN at the first residual on some nonlinear problems when the initial guess is not consistent with the BCs. Imposing the BCs for such guesses introduces very sharp gradients. See this PR description, and the test that we added for pre_apply_bcs=False.

[pbrubeck]

I opened #5177 to continue with this discussion

[JHopeCollins]

Suggested change

	pc_types = ("gamg",
	"jacobi",
	"hypre",
	"ilu")
	pc_types = ("gamg",
	"jacobi",
	pytest.param("hypre", marks=skiphypre),
	"ilu")

[JHopeCollins]

See suggestion above to use the existing skippackage marker.

Suggested change

	if pc_type == 'hypre' and not PETSc.Sys.hasExternalPackage("hypre"):
	pytest.xfail("Need PETSc with HYPRE support for this test")

[JHopeCollins]

Do we need to be testing this many PCs? Do we expect significantly different behaviour between them that would indicate a bug Firedrake? Just trying to be conscious of the size of our (already bloated) test suite!

[pbrubeck]

Jacobi and gamg should stay, as they were originally broken (they were too inexact and caused NewtonTR failures). Hypre and ILU never failed so they serve as our reference, so we can at most remove one.

This is a 1D test, and changing the PC does not generate any new code, we should not be worried about timing here.

[JHopeCollins]

Seeing as it's 1D and you want it as a reference, how about replacing hypre and ilu with lu?

[stefanozampini]

These tests exercise approximate linear solves, so lu does not belong here