diff --git a/demos/demo_references.bib b/demos/demo_references.bib index 9c84a461a5..bf9e39d1d9 100644 --- a/demos/demo_references.bib +++ b/demos/demo_references.bib @@ -404,3 +404,15 @@ @Article{Farrell2015 pages = {A2026--A2045}, doi = {10.1137/140984798}, } + +@article{brune2015composing, + title={Composing scalable nonlinear algebraic solvers}, + author={Brune, Peter R and Knepley, Matthew G and Smith, Barry F and Tu, Xuemin}, + journal={SIAM Review}, + volume={57}, + number={4}, + pages={535--565}, + year={2015}, + publisher={SIAM}, + doi={10.1137/130936725} +} diff --git a/demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst b/demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst new file mode 100644 index 0000000000..ca6361bb84 --- /dev/null +++ b/demos/nonlinear_pc/nonlinear_pc_allen_cahn.py.rst @@ -0,0 +1,292 @@ +Nonlinear preconditioning applied to the Allen-Cahn equation +============================================================ + +Contributed by `Daniel Shapero `_ +and `Josh Hope-Collins `_. + +In the other demonstrations on nonlinear PDE, we've used Newton's method to +solve finite- dimensional systems of nonlinear equations. Suppose we want to +solve a system + +.. math:: + + F(u) = 0. + +At each step of Newton's method, we first compute a *search direction* +:math:`v` by solving the linear system + +.. math:: + + dF(u)v = -F(u). + +We can then use a line search along :math:`v` to obtain the next candidate +solution. When we talked about preconditioning, we have always meant applying +a preconditioner to the linear system for the search direction. A linear +preconditioner transforms a linear system into an equivalent one that is +easier to solve. + +For some highly nonlinear problems, however, Newton's method can stagnate or +fail to converge, even with globalization strategies such as a line search. +For example, if :math:`dF(u)` behaves badly, we might not be able to compute a +search direction at all. + +An alternative strategy is to use *nonlinear preconditioning*, abbreviated as +NPC in the following. Nonlinear preconditioning does the same thing as linear +preconditioning: it transforms the problem into an equivalent one that is +(hopefully) easier to solve. Rather than work on the linear system for the +search direction, however, NPC works directly on the nonlinear system itself. +Typically NPC is applied within a higher-level nonlinear solver. For example, +there is a nonlinear analogue of GMRES. For a review of nonlinear solver and +preconditioning strategies beyond Newton, see :cite:`brune2015composing`. + +Here we will demonstrate how to use nonlinear preconditioning for the steady- +state Allen-Cahn equation + +.. math:: + + F(u) = -\epsilon\Delta u + u^3 - u = 0 + +on the unit interval with Dirichlet boundary conditions :math:`u(0) = -1` and +:math:`u(1) = +1`. The Jacobian of this residual is indefinite wherever +:math:`|u| < 1/\sqrt{3}`. Newton's method can fail to compute a search +direciton when starting from an initial guess that crosses this region. + +This demo is adapted from this `Chebfun example`_. + +.. _Chebfun example: https://www.chebfun.org/examples/ode-nonlin/AllenCahn.html + +:: + + import numpy as np + from firedrake import * + +Here we use a domain of length 10, a small diffusion coefficient of 0.003, +and an initial guess that ramps from +1 at the left-hand boundary to -1 at +the right. + +:: + + nx = 128 + lx = 10.0 + eps = Constant(3e-3) + + mesh = IntervalMesh(nx, lx) + Q = FunctionSpace(mesh, "CG", 1) + + x, = SpatialCoordinate(mesh) + u_1 = Constant(1) + u_2 = Constant(-1) + Lx = Constant(lx) + initial_guess = (1 - x / Lx) * u_1 + x / Lx * u_2 + + bcs = [DirichletBC(Q, u_1, [1]), DirichletBC(Q, u_2, [2])] + + u = Function(Q) + u.interpolate(initial_guess) + +In order to have a good baseline to compare against, we want to use as good +a nonlinear solution strategy as possible. Here we use Newton's method with +the *critical point* line search, which is specially adapted for problems +like Allen-Cahn which can be derived from minimizing some free energy. Even +with 10 iterations of a line search, Newton's method will fail. You can try +other line search methods (secant, backtracking, etc.) and find similar +outcomes. If you pump the number of line search iterations way up, you can +make Newton's method converge... to the wrong solution! + +:: + + v = TestFunction(Q) + F = (eps * inner(grad(u), grad(v)) + inner(u**3 - u, v)) * dx + + problem = NonlinearVariationalProblem(F, u, bcs) + + newton_parameters = { + "snes_type": "newtonls", + "snes_monitor": None, + "snes_converged_reason": None, + "snes_linesearch_type": "cp", + "snes_linesearch_max_it": 10, + "snes_linesearch_monitor": None, + } + solver = NonlinearVariationalSolver(problem, solver_parameters=newton_parameters) + try: + solver.solve() + except ConvergenceError as err: + print(err) + print("--------------------------") + print("Solver failed to converge!") + print("Resetting `u`") + u.interpolate(initial_guess) + +The residual decreases at first but eventually diverges. Here's some of the PETSc +log output: + +.. code-block:: console + + $ python nonlinear_pc_allen_cahn.py + 0 SNES Function norm 2.439229081145e-01 + 1 SNES Function norm 5.663027251974e+01 + 2 SNES Function norm 1.667065795962e+01 + 3 SNES Function norm 4.864679262673e+00 + 4 SNES Function norm 1.388964354434e+00 + 5 SNES Function norm 3.744562754242e-01 + 6 SNES Function norm 8.820230192544e-02 + 7 SNES Function norm 3.682817028776e-02 + 8 SNES Function norm 5.405858541080e-02 + 9 SNES Function norm 5.368629859638e-02 + ... + 26 SNES Function norm 5.280237868253e-02 + 27 SNES Function norm 1.098413882561e+00 + 28 SNES Function norm 1.082944223664e+00 + 29 SNES Function norm 6.520280025999e+03 + Nonlinear firedrake_0_ solve did not converge due to DIVERGED_DTOL iterations 29 + +In other scenarios, the solver doesn't explode as dramatically but rather +stagnates and exceeds the number of allowable Newton iterations. + +To salvage the wreck, we'll use nonlinear preconditioning. Here we use +the simplest solution strategy possible: preconditioned nonlinear +Richardson iterations. +The idea is that if :math:`F(u)=0` is too difficult to solve, we can +instead solve a auxiliary operator :math:`G(u)=0` which is in some sense +*nearby* to :math:`F`, and use that solution to iterate towards a solution +for :math:`F(u)=0`. +At each iteration :math:`k` we solve the following system for the next +iterate :math:`u_{k+1}` using the value of the current iterate :math:`u_k`. + +.. math:: + + G_{k}(u_{k+1}) = G_{k}(u_{k}) - F(u_{k}), + \quad + G_{k}(u) = G(u; u_{k}). + +We can note several properties of this iteration. Firstly, if +:math:`F(u_{*})=0` then :math:`u_{k}=u_{*}` is a fixed point of the iteration. +Secondly, we never have to solve :math:`F(u)=0`, we only have to evaluate +its residual at a given state. Thirdly, we can intuitively hope that if +:math:`G` is *close enough* to :math:`F` then the iteration will converge, +although proving this is much more difficult than in the linear case. +Lastly, we can define a different :math:`G_k` at each iteration by +parameterising with the current iterate, as we will see below. + +Our approach for defining :math:`G` here is to hold back the linear part of +the reaction term in the Allen-Cahn equation to the value at the previous +iteration. In other words, at each step, we define the PDE + +.. math:: + + G(u; u_k) = -\epsilon\Delta u + u^3 - u_k = 0. + +The Jacobian for this problem w.r.t. :math:`u` is symmetric and positive- +definite. We have good guarantees about the convergence of Newton's method +in that case, so we can reuse the solver parameters that we tried the first +time around for this inner problem. + +Here we define a custom nonlinear preconditioner which inherits from +:class:`~firedrake.preconditioners.auxiliary_snes.AuxiliaryOperatorSNES`. +Similar to ``AuxiliaryOperatorPC``, we have to implement the ``form`` method. +This method returns (1) a residual form :math:`G(u; u^{k}, v)` where :math:`v` +is the test function and (2) the boundary conditions for this sub-problem. +The arguments that it takes are first the PETSc SNES object, +then the value :math:`u_k` of the solution at the previous iteration; +the current value :math:`u` to be solved for; and a test function. + +:: + + 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)) + inner(u**3 - u_k, v)) * dx, bcs + +The contract for the ``form`` method requires it to supply the boundary +conditions. We could have obtained the boundary conditions by pulling them out +of the global context. That will work for this particular set of solvers but +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 ``F``, which we aren't using here. For other +nonlinear preconditioners, we might instead be building the preconditioning +form by modifying ``F``. + +We now set ``-snes_type nrichardson`` for nonlinear Richardson iterations, +and specify the additional parameters for the nonlinear preconditioner under +the key ``"npc"``. Firstly we need to specify our python SNES type, and then +in the ``"aux"`` key we specify the parameters to actually solve :math:`G` - +here we use Newton iterations. + +:: + + solver_parameters = { + "snes_type": "nrichardson", + "snes_monitor": None, + "snes_converged_reason": None, + "npc": { + "snes_type": "python", + "snes_python_type": f"{__name__}.AllenCahnAuxSNES", + "aux": newton_parameters, + }, + } + + solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters) + solver.solve() + +Now we actually converge! The convergence on nonlinear Richardson iterations is +usually linear, as opposed to the quadratic convergence of Newton, but with +suitable preconditioning they will usually have a wider basin of convergence than +Newton. From the log output, we can observe that the residual is decreasing by a +factor of 2 or more at each iteration. + +.. code-block:: console + + 0 SNES Function norm 2.439229081145e-01 + 1 SNES Function norm 2.405339859939e-01 + 2 SNES Function norm 1.540442351803e-01 + 3 SNES Function norm 7.166137071498e-02 + 4 SNES Function norm 2.854773301463e-02 + 5 SNES Function norm 1.070593887487e-02 + 6 SNES Function norm 3.943106335233e-03 + 7 SNES Function norm 1.451230558440e-03 + ... + 18 SNES Function norm 3.491990058865e-08 + 19 SNES Function norm 1.349411703695e-08 + 20 SNES Function norm 5.217958176918e-09 + 21 SNES Function norm 2.018693509573e-09 + Nonlinear firedrake_1_ solve converged due to CONVERGED_FNORM_RELATIVE iterations 21 + +We've chosen to use nonlinear Richardson here because it's the simplest scheme +with the fewest knobs to turn. There are more sophisticated strategies which can +offer a faster convergence rate. For example, you can run this demo again with +`"snes_type": "ngmres"` to use a nonlinear variant of the generalised minimum +residual algorithm. NGMRES with the default options cuts the number of +iterations in half for this problem. But it has more algorithmic knobs to turn +and those can require tweaking depending on the problem. + +The Allen-Cahn equation is derivable through minimization of the free energy +functional + +.. math:: + + E(u) = \int_\Omega\left(\frac{\epsilon}{2}|\nabla u|^2 + \frac{1}{4}(1 - u^2)^2\right)dx. + +To close, let's evaluate the free energy at the starting guess and at the +computed solution. + +:: + + E = (0.5 * eps * inner(grad(u), grad(u)) + 0.25 * (1 - u**2) ** 2) * dx + E_initial = firedrake.assemble(firedrake.replace(E, {u: initial_guess})) + E_final = firedrake.assemble(E) + print(f"Initial free energy: {E_initial.real:0.04f}") + print(f"Final: {E_final.real:0.04f}") + +.. code-block:: console + + Initial free energy: 1.3339 + Final: 0.0534 + +This demo can be found as a script in :demo:`nonlinear_pc_allen_cahn.py `. + +.. rubric:: References + +.. bibliography:: demo_references.bib + :filter: docname in docnames diff --git a/docs/source/advanced_tut.rst b/docs/source/advanced_tut.rst index 38d9a3d849..244913e737 100644 --- a/docs/source/advanced_tut.rst +++ b/docs/source/advanced_tut.rst @@ -36,3 +36,4 @@ element systems. Steady multicomponent flow -- microfluidic mixing of hydrocarbons. Deflation techniques for computing multiple solutions of nonlinear problems. Coupled volume-surface reaction-diffusion on a torus using Submesh and geometric multigrid. + Nonlinear preconditioning using an auxiliary SNES for the Allen-Cahn equation. diff --git a/docs/source/conf.py b/docs/source/conf.py index 27fc79e5f2..62138fe058 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -432,6 +432,7 @@ 'h5py': ('http://docs.h5py.org/en/latest/', None), 'h5py.h5p': ('https://api.h5py.org/', None), 'matplotlib': ('https://matplotlib.org/', None), + 'petsc4py': ('https://petsc.org/release/petsc4py/', None), 'python': ('https://docs.python.org/3/', None), 'pyadjoint': ('https://pyadjoint.org/', None), 'numpy': ('https://numpy.org/doc/stable/', None), diff --git a/firedrake/__init__.py b/firedrake/__init__.py index 30e1a57ef0..cdf43c92f5 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -88,7 +88,7 @@ def init_petsc(): MassInvPC, PCDPC, PatchPC, PlaneSmoother, PatchSNES, P1PC, P1SNES, LORPC, GTMGPC, PMGPC, PMGSNES, HypreAMS, HypreADS, FDMPC, PoissonFDMPC, TwoLevelPC, HiptmairPC, FacetSplitPC, BDDCPC, - CovariancePC + CovariancePC, AuxiliaryOperatorSNES ) from firedrake.mesh import ( # noqa: F401 Mesh, ExtrudedMesh, VertexOnlyMesh, RelabeledMesh, diff --git a/firedrake/preconditioners/__init__.py b/firedrake/preconditioners/__init__.py index d3436114df..e66d9defca 100644 --- a/firedrake/preconditioners/__init__.py +++ b/firedrake/preconditioners/__init__.py @@ -25,3 +25,4 @@ from firedrake.preconditioners.facet_split import FacetSplitPC # noqa: F401 from firedrake.preconditioners.bddc import BDDCPC # noqa: F401 from firedrake.preconditioners.covariance import CovariancePC # noqa: F401 +from firedrake.preconditioners.auxiliary_snes import AuxiliaryOperatorSNES # noqa: F401 diff --git a/firedrake/preconditioners/auxiliary_snes.py b/firedrake/preconditioners/auxiliary_snes.py new file mode 100644 index 0000000000..95e4c98e41 --- /dev/null +++ b/firedrake/preconditioners/auxiliary_snes.py @@ -0,0 +1,238 @@ +from firedrake.preconditioners.base import SNESBase +from firedrake.function import Function +from firedrake.cofunction import Cofunction +from firedrake.ufl_expr import Argument, TestFunction +from firedrake.petsc import PETSc +from ufl import replace +from firedrake.dmhooks import get_function_space, get_appctx as get_dm_appctx + +__all__ = ("AuxiliaryOperatorSNES",) + + +class AuxiliaryOperatorSNES(SNESBase): + """ + Solve a residual form :math:`F(u) = 0` using a nonlinear Richardson + iteration preconditioned with an auxiliary form :math:`G(u)`. + This may be used to create nonlinear preconditioners for + iterative methods such as Anderson acceleration or NGMRES. + + The `k`-th nonlinearly preconditioned Richardson iteration is: + + .. math :: + + G(u^{k+1}) = G(u^{k}) - F(u^{k}) + + The solution :math:`u^{*}` of :math:`F(u^{*}) = 0` is a fixed point + of the Richardson iteration. + + .. math :: + + G(u^{k+1}) = G(u^{*}) - F(u^{*}) + + G(u^{k+1}) = G(u^{*}) + + \\implies u^{k+1} = u^{*} + + Options for the inner solve for :math:`G(u^{k+1})` are specified + using the ``"aux_"`` prefix. + + The following solver parameters will specify the above Richardson + iteration, assuming that the user has defined a class + ``UserAuxiliarySNES`` which inherits from + :class:`~.AuxiliaryOperatorSNES` and implements the + :meth:`~.AuxiliaryOperatorSNES.form` method. In this example, the + inner solve uses a Newton method with a relative tolerance of 1e-4. + + .. code-block:: python3 + + solver_parameters = { + "snes_rtol": 1e-8, + "snes_type": "python", + "snes_python_type": f"{__name__}.UserAuxiliarySNES", + "aux": { + "snes_rtol": 1e-4, + "snes_type": "newtonls", + ... + } + } + + The following parameters describe the same Richardson iteration + as the parameters above, but explicitly specifying the auxiliary + form as a nonlinear preconditioner using the ``npc_`` prefix. + + .. code-block:: python3 + + solver_parameters = { + "snes_rtol": 1e-8, + "snes_type": "nrichardson", + "npc_snes_max_it": 1, + "npc_snes_type": "python", + "npc_snes_python_type": f"{__name__}.UserAuxiliarySNES", + "npc_aux": { + "snes_rtol": 1e-4, + "snes_type": "newtonls", + ... + } + } + + Although using ``"npc_"`` to specify the parameters is more verbose + than the original, it allows for a wider variety of methods. For + example, by changing the outer ``"snes_type"`` to ``"anderson"``, + we can use preconditioned Anderson acceleration + (``_) + + Notes + ----- + If the auxiliary form is linear, i.e. :math:`G(u)=Au`, with + :math:`A\\approx\\nabla F` approximating the Jacobian of the + outer residual, then the iteration is an inexact Newton method: + + .. math :: + + Au^{k+1} = Au^{k} - F(u^{k}) + + A\\delta u^{k} = - F(u^{k}) + + u^{k+1} = u^{k} + \\delta u^{k} + """ + + _prefix = "aux_" + + @PETSc.Log.EventDecorator() + def initialize(self, snes): + from firedrake.variational_solver import ( # circular import at file level + NonlinearVariationalSolver, NonlinearVariationalProblem) + + ctx = get_dm_appctx(snes.dm) + + parent_prefix = snes.getOptionsPrefix() or "" + prefix = parent_prefix + self._prefix + + V = get_function_space(snes.getDM()).collapse() + + # buffers for current and next iterates + uk1 = Function(V) + uk = Function(V) + self.uk1 = uk1 + self.uk = uk + + # auxiliary form G(k+1) + test = TestFunction(V) + Gk1, bcs = self.form(snes, uk, uk1, test) + + b = Cofunction(V.dual()) + # This is the form we will solve: + # G(u^{k+1}) - b = 0 + # and we will assemble G(u^{k}) - F(u^{k}) into b. + Gk1b = Gk1 - b + + self.b = b + # a buffer for intermediate values when assembling b = Gk - Fk + self._b_wrk = Cofunction(V.dual()) + + self.solver = NonlinearVariationalSolver( + NonlinearVariationalProblem( + Gk1b, uk1, bcs=bcs, + form_compiler_parameters=ctx.fcp), + nullspace=ctx._nullspace, + transpose_nullspace=ctx._nullspace_T, + near_nullspace=ctx._near_nullspace, + appctx=ctx.appctx, options_prefix=prefix) + + # indent monitor outputs + outer_snes = snes + inner_snes = self.solver.snes + inner_snes.incrementTabLevel(1, parent=outer_snes) + inner_snes.ksp.incrementTabLevel(1, parent=outer_snes) + inner_snes.ksp.pc.incrementTabLevel(1, parent=outer_snes) + + def update(self, snes): + pass + + @PETSc.Log.EventDecorator() + def step(self, snes, x, f, y): + # x = u^{k} is state at current iteration + with self.uk.dat.vec_wo as vec: + x.copy(vec) + # initial guess u^{k+1} = u^{k} + with self.uk1.dat.vec_wo as vec: + x.copy(vec) + + # b = F(u^{k}) + with self.b.dat.vec as vec: + f.copy(vec) + + # At this point we have: + # uk1 = x = u^{k}, and b = F(u^{k}), + # so the solver's assembler computes: + # b_wrk = G(uk1) - b = G(u^{k}) - F(u^{k}) + # which is exactly the forcing we need. + self.solver._ctx._assemble_residual(tensor=self._b_wrk) + + # we assign b = G(u^{k}) - F(u^{k}) so now the + # form in the solver has the correct forcing to + # calculate u^{k+1} using G(uk1) - b = 0 + self.b.assign(self._b_wrk) + + self.solver.solve() + + # y = d = u^{k+1} - u^{k} + with self.uk1.dat.vec_ro as vec: + vec.copy(y) + y.aypx(-1, x) + + @PETSc.Log.EventDecorator() + def form(self, snes, uk: Function, uk1: Function, test: Argument): + """Return the auxiliary residual form and boundary conditions. + Subclasses should override this method. + + The returned form is :math:`G(u^{k+1})` in the Richardson iteration. + + Defaults to returning a copy of :math:`F(u)`, i.e. + :math:`G(u^{k+1})=F(u^{k+1})`. + This means that ``AuxiliaryOperatorSNES`` can be used similarly + to :class:`.AssembledPC`, in that it can be used to specify + an alternative ``snes_type`` for solving the same residual form + as the outer ``SNES``. + + Parameters + ---------- + snes : petsc4py.PETSc.SNES + The PETSc nonlinear solver object. + uk : + The current iterate :math:`u^{k}`. + uk1 : + The next iterate :math:`u^{k+1}` that will be solved for. + test : + The test function. + + Returns + ------- + G : :class:`ufl.BaseForm` + The preconditioning residual form. + bcs : Iterable[:class:`~.firedrake.bcs.DirichletBC`] | None + The boundary conditions. + + Notes + ----- + :math:`G(u^{k+1})` can optionally be parameterised by the current + iterate :math:`u^{k}`, for example in Picard iterations for an + advection term: + + .. math:: + + F(u) = u + u\\cdot\\nabla u + + G(u^{k+1}) = u^{k+1} + u^{k}\\cdot\\nabla u^{k+1} + """ + ctx = get_dm_appctx(snes.dm) + u = ctx._x + form = replace(ctx._problem.F, {u: uk1}) + bcs = tuple(ctx._problem.bcs) + return form, bcs + + def view(self, snes, viewer=None): + super().view(snes, viewer) + if hasattr(self, "solver"): + viewer.printfASCII("SNES to apply a preconditioned Richardson iteration.\n") + self.solver.snes.view(viewer) diff --git a/firedrake/preconditioners/base.py b/firedrake/preconditioners/base.py index 2d373b94fd..57522747a1 100644 --- a/firedrake/preconditioners/base.py +++ b/firedrake/preconditioners/base.py @@ -10,18 +10,24 @@ class PCSNESBase(object, metaclass=abc.ABCMeta): - def __init__(self): - """Create a PC context suitable for PETSc. + """Create a PC or SNES python context suitable for PETSc. - Matrix free preconditioners should inherit from this class and - implement: + Both Python PC and SNES classes should inherit from this class and implement: - - :meth:`~.PCSNESBase.initialize` - - :meth:`~.PCSNESBase.update` - - :meth:`~.PCBase.apply` - - :meth:`~.PCBase.applyTranspose` + - :meth:`~.PCSNESBase.initialize` + - :meth:`~.PCSNESBase.update` - """ + Python PC classes should additionally implement: + + - :meth:`~.PCBase.apply` + - :meth:`~.PCBase.applyTranspose` + + Python SNES classes should additionally implement one of the following: + + - ``SNESBase.step`` + - ``SNESBase.solve`` + """ + def __init__(self): petsctools.cite("Kirby2017") self.initialized = False super(PCSNESBase, self).__init__() @@ -107,8 +113,8 @@ def form(self, obj, *args): return a, bcs @staticmethod - def get_appctx(pc): - return get_appctx(pc.getDM()).appctx + def get_appctx(obj): + return get_appctx(obj.getDM()).appctx @staticmethod def new_snes_ctx( @@ -154,6 +160,16 @@ def new_snes_ctx( class PCBase(PCSNESBase): + """Create a PC python context suitable for PETSc. + + Matrix free preconditioners should inherit from this class and + implement: + + - :meth:`~.PCSNESBase.initialize` + - :meth:`~.PCSNESBase.update` + - :meth:`~.PCBase.apply` + - :meth:`~.PCBase.applyTranspose` + """ _asciiname = "preconditioner" _objectname = "pc" @@ -199,6 +215,44 @@ def setUp(self, pc): class SNESBase(PCSNESBase): + """Create SNES python context suitable for PETSc. + + Python SNES classes should inherit from this class and implement: + + - :meth:`~.PCSNESBase.initialize` + - :meth:`~.PCSNESBase.update` + + Inheriting classes should additionally implement *either*: + + - ``SNESBase.step`` + - ``SNESBase.solve`` + + The required function signatures for each method are shown below: + + .. code-block:: python3 + + def solve(self, snes, b, x): + '''Solve the nonlinear problem using the Vec x as the initial guess and + putting the solution back into x. The Vec b is constant forcing term which + may be None. + ''' + pass + + def step(self, snes, X, F, Y): + '''Apply one iteration of the SNES to the current iterate X, + using the function residual F, and putting the update in Y. + + X, F and Y are PETSc Vecs, Y is not guaranteed to be zero on entry. + ''' + pass + + Notes + ----- + The function signatures for the ``solve`` and ``step`` methods are shown in + the docstring rather than being implemented as abstract methods because + petsc4py will test whether the SNES python context has either a ``step`` + or ``solve`` method to decide what to do when ``snes.solve()`` is called. + """ _asciiname = "nonlinear solver" _objectname = "snes" diff --git a/tests/firedrake/demos/test_demos_run.py b/tests/firedrake/demos/test_demos_run.py index bffd70ff51..f6ea8cf9cd 100644 --- a/tests/firedrake/demos/test_demos_run.py +++ b/tests/firedrake/demos/test_demos_run.py @@ -54,6 +54,7 @@ Demo(('vlasov_poisson_1d', 'vp1d'), []), Demo(('shape_optimization', 'shape_optimization'), ["adjoint", "vtk"]), Demo(('submesh_reaction_diffusion', 'submesh_reaction_diffusion'), ["netgen", "vtk"]), + Demo(('nonlinear_pc', 'nonlinear_pc_allen_cahn'), []), ] PARALLEL_DEMOS = [ Demo(("full_waveform_inversion", "full_waveform_inversion"), ["adjoint"]), diff --git a/tests/firedrake/regression/test_auxiliary_snes.py b/tests/firedrake/regression/test_auxiliary_snes.py new file mode 100644 index 0000000000..68bbff77b6 --- /dev/null +++ b/tests/firedrake/regression/test_auxiliary_snes.py @@ -0,0 +1,93 @@ +import firedrake as fd + + +def f(u, v, w=0.): + if not isinstance(w, fd.Constant): + w = fd.Constant(w) + return ( + fd.inner(u, v)*fd.dx + + w*fd.inner(u*u*u, v)*fd.dx + - fd.inner(fd.Constant(1), v)*fd.dx + ) + + +class AuxiliaryPolynomialSNES(fd.AuxiliaryOperatorSNES): + 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 + + +def test_auxiliary_snes(): + """Check that AuxiliaryOperatorSNES is equivalent to a + hand-rolled preconditioned nonlinear Richardson iteration. + """ + + mesh = fd.UnitIntervalMesh(1) + V = fd.FunctionSpace(mesh, "R", 0) + + uk1 = fd.Function(V).zero() + uk = fd.Function(V).zero() + v = fd.TestFunction(V) + + w = 1.0 + wg = 0.75 + + Fk = f(uk, v, w=w) + Gk = f(uk, v, w=wg) + Gk1 = f(uk1, v, w=wg) + + Fp = Gk1 - (Gk - Fk) + + inner_params = { + "snes_monitor": None, + "snes_converged_reason": None, + "snes_type": "newtonls", + "snes_rtol": 1e-2, + "ksp_type": "gmres", + "pc_type": "none", + } + + aux_params = { + "snes_monitor": None, + "snes_converged_reason": None, + "snes_type": "nrichardson", + "snes_atol": 1e-4, + "npc_snes_max_its": 1, + "npc_snes_type": "python", + "npc_snes_python_type": f"{__name__}.AuxiliaryPolynomialSNES", + "npc_snes_aux_w": wg, + "npc_aux": inner_params + } + + solver_inner = fd.NonlinearVariationalSolver( + fd.NonlinearVariationalProblem(Fp, uk1), + solver_parameters=inner_params, + options_prefix="Finner") + + solver_aux = fd.NonlinearVariationalSolver( + fd.NonlinearVariationalProblem(Fk, uk), + solver_parameters=aux_params, + options_prefix="Faux") + + # Record the residual at each handrolled richardson iteration + uk1.zero() + uk.assign(uk1) + res = [] + with fd.assemble(Fk).dat.vec as rvec: + res.append(rvec.norm()) + for k in range(5): + solver_inner.solve() + uk.assign(uk1) + with fd.assemble(Fk).dat.vec as rvec: + res.append(rvec.norm()) + + # Custom monitor to compare residuals at each aux operator iteration + def comparison_monitor(snes, its, rnorm): + assert abs(rnorm - res[its]) < 1e-14 + + solver_aux.snes.setMonitor(comparison_monitor) + + uk.zero() + solver_aux.solve()