`AuxiliaryOperatorSNES` by JHopeCollins · Pull Request #5119 · firedrakeproject/firedrake

JHopeCollins · 2026-05-15T13:48:49Z

If I am trying to solve F(u)=0 but for whatever reason Newton isn't working or is slow to converge, this snes type enables specifying an auxiliary nonlinear operator G(u) that can be used as a nonlinear preconditioner.

The simplest form of nonlinear preconditioning is a for a nonlinear Richardson iteration:

$$G(u^{k+1}) = G(u^{k}) - F(u^{k})$$

where the solution uhat of F(uhat)=0 is clearly a fixed point.

There's a demo for the Allen-Cahn equation where Newton fails but using a G that lags a particular term in F succeeds. AuxiliaryOpertorSNES can also be used for continuation methods like Reynolds continuation, and many other methods.

This cherrypicks the AuxiliaryOperatorSNES from #4544 because it is ready to go but the FieldsplitSNES could do with a little more stress testing/tidying up.

Pointing at release because this is strictly new functionality. The edits to firedrake/preconditioners/base.py are just updating the docstrings to be clearer about the distinction between PCBase and SNESBase.

…firedrake into JHopeCollins/auxopsnes

dham

minor changes pablo has requested to the demo.

pbrubeck · 2026-05-19T16:18:32Z

+        # with forcing b = G(k) - F(k)
+        Gk = replace(Gk1, {uk1: uk})
+        b = Cofunction(V.dual())
+        Gk1 -= b


Gk1 is not G(u^{k+1}) anymore. Is this the right name? How about F_aux?

pbrubeck · 2026-05-19T16:19:19Z

+        self.Gk = Gk
+        self.Gk1 = Gk1


Do these need to be stored as class attributes?

pbrubeck · 2026-05-19T16:21:49Z

+
+        # G(u^{k+1}) - b = 0
+        self.uk1.assign(self.uk)
+        self.solver.solve()


We might want to enable solver.solve(rhs=b) so the logic with Gk1 becomes clearer. Obviously a separate PR.

pbrubeck · 2026-05-19T16:26:07Z

+            Gk, bcs=bcs,
+            form_compiler_parameters=ctx.fcp,
+            options_prefix=prefix
+        ).assemble


Potentially one could reuse the residual assembler in the _SNESContext of the solver defined below, similar to what we did here for AssembledPC

https://github.com/firedrakeproject/firedrake/pull/5079/changes

Up to you if you want to defer this job to a separate PR.

pbrubeck · 2026-05-19T20:20:33Z

+can create problems for others, for example when using the multigrid method.
+Instead, we've opted to call the parent class's `form` method, which will
+return the original variational form and boundary conditions. We then discard
+the original variational form, which we aren't using here.


Maybe mention that we can also manipulate F

Is that safe? Then the inner problem is mutating the outer problem, which seems impolite

I mean creating a new G from F, e.g. G = F + extra_terms

ufl.Forms are supposed to be immutable

Also F += extra_terms will not modify F "in-place", it'd just assign a new Form to F.

Ok, I added a note explaining that you can also come up with the form by altering F instead of creating it from whole cloth.

pbrubeck · 2026-05-19T21:39:46Z

+    def form(self, snes, state, func, test):
+        prefix = (snes.getOptionsPrefix() or "") + f"snes_{self._prefix}"
+        w = fd.PETSc.Options().getScalar(prefix + "w")
+        return f(func, test, w=w), None


Suggested change

def form(self, snes, state, func, test):

prefix = (snes.getOptionsPrefix() or "") + f"snes_{self._prefix}"

w = fd.PETSc.Options().getScalar(prefix + "w")

return f(func, test, w=w), None

def form(self, snes, state, func, test):

F, bcs = super().form(snes, state, func, test)

prefix = (snes.getOptionsPrefix() or "") + f"snes_{self._prefix}"

w = fd.PETSc.Options(prefix).getScalar("w")

return f(func, test, w=w), bcs

pbrubeck · 2026-05-19T21:42:46Z

+            y.aypx(-1, x)
+
+    @PETSc.Log.EventDecorator()
+    def form(self, snes, uk: Function, uk1: Function, test: Argument):


Maybe uold, unew would be more clear here

uk1 and uk matches the equations describing the iteration in the class docs and the demo so they are clearer here than old/new

Sure, with the context of the docstrings, you could tell that uk1 refers to u_{k+1}. But a subclass will not include the docstring, so the context will get lost easily in user code. If a user cargo-culted this notation, it be hard to tell whether uk1 refers to u_{k-1} or u_{k+1}.

I had to puzzle over this for a moment when I was writing the demo. I agree that a different naming scheme could be beneficial. Is there a convention established elsewhere in Firedrake?

pbrubeck · 2026-05-20T07:21:09Z

+  class AllenCahnAuxSNES(firedrake.AuxiliaryOperatorSNES):
+      def form(self, snes, u_k, u, v):
+          F, bcs = super().form(snes, u_k, u, v)
+          return (eps * inner(grad(u), grad(v)) + (u**3 - u_k) * v) * dx, bcs


Please correct me if I'm wrong, but "lagging" the negative term here is equivalent to dropping it from G.

Also having G(u^k, u^{k+1}) = F(u^k) when u^{k+1} = u^k gives you no extra source term (b).

Maybe it'd be worth adding a comment.

To your second point, this is the exactly the point about the stationary point of the iteration

pbrubeck · 2026-05-20T07:27:42Z

+
+.. math::
+
+  G(u) = \int_\Omega\left(\frac{\epsilon}{2}|\nabla u|^2 + \frac{1}{4}(1 - u^2)^2\right)dx.


Now G(u) means a different thing. A better name would be E(u)

I've changed this as you suggest. It would be a bigger change but I could instead define E(u) at the very beginning and work almost entirely from this, defining F = firedrake.derivative(E, u).

If you do that, then the demo does not run in complex mode

pbrubeck · 2026-05-20T13:40:48Z

+        with self.uk.dat.vec_wo as vec:
+            x.copy(vec)
+
+        # b = G(u^{k})
+        self.assemble_gk(tensor=self.b)
+
+        # b = G(u^{k}) - F(u^{k})
+        with self.b.dat.vec as vec:
+            vec -= f
+
+        # G(u^{k+1}) - b = 0
+        self.uk1.assign(self.uk)
+        self.solver.solve()


We can get Gk-f out of Gk1-b (no need to generate code twice). For the moment, one needs an extra buffer self.b1, but this extra buffer might go away once #5130 lands

Suggested change

with self.uk.dat.vec_wo as vec:

x.copy(vec)

# b = G(u^{k})

self.assemble_gk(tensor=self.b)

# b = G(u^{k}) - F(u^{k})

with self.b.dat.vec as vec:

vec -= f

# G(u^{k+1}) - b = 0

self.uk1.assign(self.uk)

self.solver.solve()

# u^{k+1} = u^{k} = x

with self.uk.dat.vec_wo as vec:

x.copy(vec)

with self.uk1.dat.vec_wo as vec:

x.copy(vec)

# b = f

with self.b.dat.vec_wo as vec:

f.copy(vec)

# b1 = G(u^{k}; u^{k}) - F(u^{k})

self.solver._ctx._assemble_residual(tensor=self.b1)

# b = b1

self.b.assign(self.b1)

# solve G(u^{k}; u^{k+1}) - b = 0 for u^{k+1}

self.solver.solve()

pbrubeck · 2026-05-20T13:59:25Z

+::
+
+  solver_parameters = {
+      "snes_type": "nrichardson",


We could try anderson as well, it reduces iterations down to 11 for this example

Suggested change

"snes_type": "nrichardson",

"snes_type": "anderson",

"snes_anderson_m": 2,

I considered this but wanted to keep it simple at first. I could add Anderson or NGMRES at the end. But I was planning another demo showing more advanced techniques using the Navier-Stokes equations in a later PR.

How about just a sentence along the lines of "we could improve the convergence by swapping to an accelerated SNES type such as Anderson or ngmres, which you may want to experiment with in practice/in your own application"

pbrubeck · 2026-05-20T14:15:36Z

+      },
+  }
+
+  solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters)


It is unclear that the failed solve did not update the value of u. We should make it obvious that we are starting from the same initial guess

Suggested change

solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters)

u.interpolate(initial_guess)

solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters)

I have this further up, immediately after the initial Newton solve failed.

pbrubeck · 2026-05-20T17:01:34Z

+    def initialize(self, snes):
+        from firedrake import (  # circular import if this is at file level
+            NonlinearVariationalSolver, NonlinearVariationalProblem,
+            Function, TestFunction, Cofunction)


Function is imported at the top of this file. Are we sure all of these give circular imports?

AuxiliaryOperatorSNES

c713dec

JHopeCollins assigned danshapero and JHopeCollins May 15, 2026

JHopeCollins added the enhancement label May 15, 2026

danshapero mentioned this pull request May 15, 2026

Lag label firedrakeproject/Irksome#228

Merged

JHopeCollins added 3 commits May 15, 2026 16:28

docs links

26f7c09

AuxiliaryOperatorSNES

a9f770b

docs links

0d93801

danshapero force-pushed the JHopeCollins/auxopsnes branch from 26f7c09 to 0d93801 Compare May 15, 2026 20:57

JHopeCollins and others added 7 commits May 16, 2026 12:14

Merge branch 'JHopeCollins/auxopsnes' of github.com:firedrakeproject/…

bbf35e9

…firedrake into JHopeCollins/auxopsnes

Add Allen-Cahn demo of NPC

c50f8b9

Merge branch 'JHopeCollins/auxopsnes' of github.com:firedrakeproject/…

3004e63

…firedrake into JHopeCollins/auxopsnes

Merge branch 'release' into JHopeCollins/auxopsnes

3066ff7

add AuxiliaryOperatorSNES demo to website and tests

3de898d

updates to AuxiliaryOperatorSNES demo

f12861a

updates to AuxiliaryOperatorSNES demo

f37c18b

JHopeCollins marked this pull request as ready for review May 19, 2026 13:58

Fix citation in Allen-Cahn demo

06fa563

dham requested changes May 19, 2026

View reviewed changes

pbrubeck reviewed May 19, 2026

View reviewed changes

Comment thread demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst

pbrubeck reviewed May 19, 2026

View reviewed changes

Comment thread demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst

pbrubeck reviewed May 19, 2026

View reviewed changes

Show some PETSc log output

6347def

pbrubeck reviewed May 19, 2026

View reviewed changes

pbrubeck reviewed May 20, 2026

View reviewed changes

Comment thread demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst Outdated

Comment thread demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst Outdated

Comment thread demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst Outdated

pbrubeck reviewed May 20, 2026

View reviewed changes

Wordsmithing demo

b416461

pbrubeck reviewed May 20, 2026

View reviewed changes


		.. math::

		G(u) = \int_\Omega\left(\frac{\epsilon}{2}\|\nabla u\|^2 + \frac{1}{4}(1 - u^2)^2\right)dx.

	"snes_type": "nrichardson",
	"snes_type": "anderson",
	"snes_anderson_m": 2,

	solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters)
	u.interpolate(initial_guess)
	solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters)

Conversation

JHopeCollins commented May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dham left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pbrubeck May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pbrubeck May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pbrubeck May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

JHopeCollins commented May 15, 2026 •

edited

Loading

pbrubeck May 20, 2026 •

edited

Loading

pbrubeck May 19, 2026 •

edited

Loading

pbrubeck May 20, 2026 •

edited

Loading