Add dDAE BCs for Nystrom (#131)

* Add dDAE BCs for Nystrom

* Adding tests for implicit BC_types

* Add more informative error message

* Add explicit wave demo

Author: ScottMacLachlan

demos/nystrom/demo_wave_explicit.py.rst

@@ -0,0 +1,109 @@
+Solving the Wave Equation with Irksome and an explicit Nystrom method
+=====================================================================
+
+Let's start with the simple wave equation on :math:`\Omega = [0,1]
+\times [0,1]`, with boundary :math:`\Gamma`:
+
+.. math::
+
+   u_{tt} - \Delta u &= 0
+
+   u & = 0 \quad \textrm{on}\ \Gamma
+
+At each time :math:`t`, the solution
+to this equation will be some function :math:`u\in V`, for a suitable function
+space :math:`V`.
+
+We transform this into weak form by multiplying by an arbitrary test function
+:math:`v\in V` and integrating over :math:`\Omega`.  We know have the
+variational problem of finding :math:`u:[0,T]\rightarrow V` such
+that
+
+.. math::
+
+   (u_{tt}, v) + (\nabla u, \nabla v) = 0
+
+As usual, we need to import firedrake::
+
+  from firedrake import *
+
+We will also need to import certain items from irksome::
+
+  from irksome import Dt, MeshConstant, StageDerivativeNystromTimeStepper, ClassicNystrom4Tableau
+
+Here, we will use the "classic" Nystrom method, a 4-stage explicit time-stepper::
+
+  nystrom_tableau = ClassicNystrom4Tableau()
+
+Now we define the mesh and piecewise linear approximating space in
+standard Firedrake fashion::
+
+  N = 32
+
+  msh = UnitSquareMesh(N, N)
+  V = FunctionSpace(msh, "CG", 2)
+
+We define variables to store the time step and current time value, noting that an explicit scheme requires a small timestep for stability::
+
+  dt = Constant(0.2 / N)
+  t = Constant(0.0)
+
+We define the initial condition for the fully discrete problem, which
+will get overwritten at each time step.  For a second-order problem,
+we need an initial condition for the solution and its time derivative
+(in this case, taken to be zero)::
+
+  x, y = SpatialCoordinate(msh)
+  uinit = sin(pi * x) * cos(pi * y)
+  u = Function(V)
+  u.interpolate(uinit)
+  ut = Function(V)
+
+Now, we will define the semidiscrete variational problem using
+standard UFL notation, augmented by the ``Dt`` operator from Irksome.
+Here, the optional second argument indicates the number of derivatives,
+(defaulting to 1)::
+
+  v = TestFunction(V)
+  F = inner(Dt(u, 2), v)*dx + inner(grad(u), grad(v))*dx
+  bc = DirichletBC(V, 0, "on_boundary")
+
+We're using an explicit scheme, so we only need to solve mass matrices in the update system.  So, some good solver parameters here are to use a fieldsplit (so we only invert the diagonal blocks of the stage-coupled system) with incomplete Cholesky as a preconditioner for the mass matrices on the diagonal blocks::
+
+  mass_params = {"snes_type": "ksponly",
+                 "mat_type": "aij",
+                 "ksp_type": "preonly",
+                 "pc_type": "fieldsplit",
+		 "pc_fieldsplit_type": "multiplicative"}
+
+  per_field = {"ksp_type": "cg",
+               "pc_type": "icc"}
+
+  for s in range(nystrom_tableau.num_stages):
+      mass_params["fieldsplit_%s" % (s,)] = per_field
+
+Most of Irksome's magic happens in the
+:class:`.StageDerivativeNystromTimeStepper`.  It takes our semidiscrete
+form `F` and the tableau and produces the variational form for
+computing the stage unknowns.  Then, it sets up a variational problem to be
+solved for the stages at each time step.  Here, we use `dDAE` style boundary conditions, which impose boundary conditions on the stages to match those of :math:`u_t` on the boundary, consistent with the underlying partitioned scheme::
+
+  stepper = StageDerivativeNystromTimeStepper(
+      F, nystrom_tableau, t, dt, u, ut, bcs=bc, bc_type="dDAE",
+      solver_parameters=mass_params)
+
+The system energy is an important quantity for the wave equation.  It
+won't be conserved with our choice of time-stepping method, but serves as a good diagnostic::
+
+  E = 0.5 * (inner(ut, ut) * dx + inner(grad(u), grad(u)) * dx)
+  print(f"Initial energy: {assemble(E)}")
+  
+Then, we can loop over time steps, much like with 1st order systems::
+
+  while (float(t) < 1.0):
+      if (float(t) + float(dt) > 1.0):
+          dt.assign(1.0 - float(t))
+      stepper.advance()
+      t.assign(float(t) + float(dt))
+      print(f"Time: {float(t)}, Energy: {assemble(E)}")
+

docs/source/index.rst

@@ -103,6 +103,7 @@ and for Nystrom methods for second-order-in-time equations:
    :maxdepth: 1
 
    demos/demo_wave.py
+   demos/demo_wave_explicit.py
    demos/demo_wave_mg.py
 
 and for bounds constraints:

irksome/nystrom_stepper.py

@@ -117,6 +117,8 @@ def bc2gcur(bc, i):
             return replace(gfoo, {t: t + c[i] * dt})
 
     elif bc_type == "DAE":
+        if tableau.is_explicit:
+            raise NotImplementedError("Can't impose DAE BCs for an explicit Nystrom tableau.  You should probably try bc_type = 'dDAE' instead")
         try:
             bA1inv = numpy.linalg.inv(tableau.Abar)
             A1inv = vecconst(bA1inv)
@@ -130,6 +132,34 @@ def bc2gcur(bc, i):
             gcur = (1/dt**2) * sum((replace(gorig, {t: t + c[j]*dt}) - ucur - utcur * (dt * c[j])) * A1inv[i, j]
                                    for j in range(num_stages))
             return gcur
+
+    elif bc_type == "dDAE":
+        if tableau.is_explicit:
+            try:
+                AAb = numpy.vstack((tableau.A, tableau.b))
+                AAb = AAb[1:]
+                bA1inv = numpy.linalg.inv(AAb)
+                A1inv = vecconst(bA1inv)
+                c_one = numpy.append(tableau.c, 1.0)
+                c_ddae = vecconst(c_one[1:])
+            except numpy.linalg.LinAlgError:
+                raise NotImplementedError("Can't have derivative DAE BC's for this method")
+        else:
+            try:
+                bA1inv = numpy.linalg.inv(tableau.A)
+                A1inv = vecconst(bA1inv)
+                c_ddae = vecconst(tableau.c)
+            except numpy.linalg.LinAlgError:
+                raise NotImplementedError("Can't have derivative DAE BC's for this method")
+
+        def bc2gcur(bc, i):
+            gorig = as_ufl(bc._original_arg)
+            gfoo = expand_time_derivatives(Dt(gorig, 1), t=t, timedep_coeffs=(u0,))
+            utcur = bc2space(bc, ut0)
+            gcur = (1/dt) * sum((replace(gfoo, {t: t + c_ddae[j]*dt}) - utcur) * A1inv[i, j]
+                                for j in range(num_stages))
+            return gcur
+
     else:
         raise ValueError(f"Unrecognized BC type: {bc_type}")
 

tests/test_nystrom.py

@@ -1,10 +1,11 @@
+import pytest
 from firedrake import (Constant, DirichletBC, Function, FunctionSpace, SpatialCoordinate,
                        TestFunction, UnitIntervalMesh, VectorFunctionSpace, assemble, div, dx,
                        norm, grad, inner, pi, project, sin, split)
-from irksome import Dt, GaussLegendre, StageDerivativeNystromTimeStepper
+from irksome import Dt, GaussLegendre, StageDerivativeNystromTimeStepper, ClassicNystrom4Tableau
 
 
-def wave(n, deg, time_stages):
+def wave(n, deg, time_stages, bc_type):
     N = 2**n
     msh = UnitIntervalMesh(N)
 
@@ -36,7 +37,7 @@ def wave(n, deg, time_stages):
     E = 0.5 * inner(ut, ut) * dx + 0.5 * inner(grad(u), grad(u)) * dx
 
     stepper = StageDerivativeNystromTimeStepper(
-        F, butcher_tableau, t, dt, u, ut, bcs=bc, solver_parameters=params)
+        F, butcher_tableau, t, dt, u, ut, bcs=bc, solver_parameters=params, bc_type=bc_type)
 
     E0 = assemble(E)
     tf = 1
@@ -49,12 +50,13 @@ def wave(n, deg, time_stages):
     return assemble(E) / E0, norm(u + u0)
 
 
-def test_wave_eq():
+@pytest.mark.parametrize("bc_type", ["ODE", "DAE", "dDAE"])
+def test_wave_eq(bc_type):
     # number of refinements
     n = 5
     deg = 2
     stage_count = 2
-    Erat, diff = wave(n, deg, stage_count)
+    Erat, diff = wave(n, deg, stage_count, bc_type)
     print(Erat, diff)
     assert abs(Erat - 1) < 1.e-8
     assert diff < 3.e-5
@@ -120,3 +122,58 @@ def test_mixed_wave_eq():
     print(Erat, diff)
     assert abs(Erat - 1) < 1.e-8
     assert diff < 3.e-5
+
+
+def explicit_wave(n, deg):
+    N = 2**n
+    msh = UnitIntervalMesh(N)
+
+    params = {"snes_type": "ksponly",
+              "ksp_type": "preonly",
+              "mat_type": "aij",
+              "pc_type": "lu"}
+
+    V = FunctionSpace(msh, "CG", deg)
+    x, = SpatialCoordinate(msh)
+
+    t = Constant(0.0)
+    dt = Constant(0.1 / N)
+
+    uinit = sin(pi * x)
+
+    nystrom_tableau = ClassicNystrom4Tableau()
+
+    u0 = project(uinit, V)
+    u = Function(u0)  # copy
+    ut = Function(V)
+
+    v = TestFunction(V)
+
+    F = inner(Dt(u, 2), v) * dx + inner(grad(u), grad(v)) * dx
+
+    bc = DirichletBC(V, 0, "on_boundary")
+
+    E = 0.5 * inner(ut, ut) * dx + 0.5 * inner(grad(u), grad(u)) * dx
+
+    stepper = StageDerivativeNystromTimeStepper(
+        F, nystrom_tableau, t, dt, u, ut, bcs=bc, solver_parameters=params, bc_type="dDAE")
+
+    E0 = assemble(E)
+    tf = 1
+    while (float(t) < tf):
+        if (float(t) + float(dt) > tf):
+            dt.assign(tf - float(t))
+        stepper.advance()
+        t.assign(float(t) + float(dt))
+
+    return assemble(E) / E0, norm(u + u0)
+
+
+def test_explicit_wave_eq():
+    # Mesh size and polynomial degree
+    n = 5
+    deg = 2
+    Erat, diff = explicit_wave(n, deg)
+    print(Erat, diff)
+    assert abs(Erat - 1) < 1.e-6
+    assert diff < 3.e-5