WIP: SDC python type SNES

demos/sdc_burgers.py

@@ -0,0 +1,60 @@
+from firedrake import *
+from irksome import Dt, TimeStepper, GaussLegendre
+
+nx = 32
+re = 100
+dt = 0.05
+stages = 3
+
+mesh = UnitIntervalMesh(nx)
+V = VectorFunctionSpace(mesh, "CG", 2)
+R = FunctionSpace(mesh, "R", 0)
+t = Function(R).assign(0)
+dt = Function(R).assign(dt)
+
+x, = SpatialCoordinate(mesh)
+ic = 1 + sin(2*pi*x)
+
+u = Function(V, name="u").project(as_vector([ic]))
+v = TestFunction(V)
+
+nu = Constant(1/re)
+
+F = (
+    inner(Dt(u), v)*dx
+    + inner(dot(u, nabla_grad(u)), v)*dx
+    + inner(nu*grad(u), grad(v))*dx
+)
+
+zero = Constant(0)
+bcs = [DirichletBC(V, as_vector([zero]), "on_boundary")]
+newton_params = {
+    "snes_type": "newtonls",
+    "ksp_type": "preonly",
+    "pc_type": "lu",
+}
+
+params = {
+    "snes": {
+        "converged_reason": None,
+        "monitor": None,
+        "rtol": 1e-5,
+    },
+    "snes_type": "python",
+    "snes_python_type": "irksome.SDCLDSNES",
+    "sdc": newton_params,
+    "sdc": {
+        "snes_type": "python",
+        "snes_python_type": "firedrake.FieldsplitSNES",
+        "snes_fieldsplit_type": "multiplicative",
+        "fieldsplit_0": newton_params,
+        "fieldsplit_1": newton_params,
+        "fieldsplit_2": newton_params,
+    }
+}
+
+stepper = TimeStepper(
+    F, GaussLegendre(3), t, dt, u,
+    bcs=bcs, solver_parameters=params)
+
+stepper.advance()

irksome/__init__.py

@@ -23,6 +23,8 @@
 from .pc import NystromAuxiliaryOperatorPC  # noqa: F401
 from .pc import RanaBase, RanaDU, RanaLD    # noqa: F401
 from .pc import IRKAuxiliaryOperatorPC      # noqa: F401
+from .snes import IRKAuxiliaryOperatorSNES      # noqa: F401
+from .snes import SDCLDSNES      # noqa: F401
 from .stage_value import StageValueTimeStepper  # noqa: F401
 from .stepper import TimeStepper            # noqa: F401
 from .tools import MeshConstant             # noqa: F401

irksome/pc.py

@@ -75,6 +75,7 @@ def form(self, pc, test, trial):
         w = ctx._x
 
         Fnew, bcnew = stepper.get_form_and_bcs(w, tableau=butcher, F=F)
+
         Jnew = derivative(Fnew, w, du=trial)
 
         return Jnew, bcnew

irksome/snes.py

@@ -0,0 +1,60 @@
+import copy
+from firedrake import AuxiliaryOperatorSNES
+from firedrake.dmhooks import get_appctx
+from irksome.pc import ldu
+
+
+class IRKAuxiliaryOperatorSNES(AuxiliaryOperatorSNES):
+    """Base class that inherits from Firedrake's AuxiliaryOperatorSNES class and
+    provides the preconditioning nonlinear form associated with an auxiliary
+    Form and/or approximate Butcher matrix (which are provided by subclasses).
+    """
+
+    def getNewForm(self, snes, u0, test):
+        """Derived classes can optionally provide an auxiliary Form.
+        Must return Fnew, bcnew, unew."""
+        raise NotImplementedError
+
+    def getAtilde(self, A):
+        """Derived classes produce a typically structured
+        approximation to A."""
+        raise NotImplementedError
+
+    def form(self, snes, u, v):
+        """Implements the interface for AuxiliaryOperatorSNES."""
+
+        appctx = get_appctx(snes.dm).appctx
+        stepper = appctx["stepper"]
+        butcher = stepper.butcher_tableau
+        F = stepper.F
+        u0 = stepper.u0
+        bcs = stepper.orig_bcs
+        w = stepper.stages.copy(deepcopy=True)
+        v0, = F.arguments()
+
+        try:
+            # use new Form if provided
+            F, bcs = self.getNewForm(snes, u0, v0)
+        except NotImplementedError:
+            pass
+
+        try:
+            # use new ButcherTableau if provided
+            Atilde = self.getAtilde(butcher.A)
+            butcher = copy.deepcopy(butcher)
+            butcher.A = Atilde
+        except NotImplementedError:
+            pass
+
+        Fnew, bcnew = stepper.get_form_and_bcs(w, tableau=butcher, F=F)
+
+        return Fnew, bcnew, w
+
+
+class SDCLDSNES(IRKAuxiliaryOperatorSNES):
+    """Implements SDC-type preconditioner using Atilde = DU where A=LDU. (Qdelta = Atilde)"""
+    prefix = "sdc_"
+
+    def getAtilde(self, A):
+        L, D, U = ldu(A)
+        return L @ D