Galerkin: support macro elements

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 .pc import LowOrderRefinedGalerkinPC   # noqa: F401
+from .pc import LowOrderRefinedDiscontinuousGalerkinPC   # noqa: F401
 from .stage_value import StageValueTimeStepper  # noqa: F401
 from .stepper import TimeStepper            # noqa: F401
 from .tools import MeshConstant             # noqa: F401

irksome/base_time_stepper.py

@@ -97,6 +97,7 @@ def __init__(self, F, t, dt, u0, num_stages,
                          bcs=bcs, appctx=appctx, nullspace=nullspace)
 
         self.num_stages = num_stages
+        self.appctx["num_stages"] = num_stages
         if butcher_tableau:
             assert num_stages == butcher_tableau.num_stages
             self.appctx["butcher_tableau"] = butcher_tableau

irksome/discontinuous_galerkin_stepper.py

@@ -1,7 +1,8 @@
 from FIAT import (Bernstein, DiscontinuousElement,
                   DiscontinuousLagrange,
                   Legendre,
-                  make_quadrature, ufc_simplex)
+                  ufc_simplex)
+from FIAT.quadrature_schemes import create_quadrature
 from ufl.constantvalue import as_ufl
 from .base_time_stepper import StageCoupledTimeStepper
 from .bcs import stage2spaces4bc
@@ -12,8 +13,7 @@
 from firedrake import TestFunction
 
 
-def getFormDiscGalerkin(F, L, Q, t, dt, u0, stages, bcs=None):
-
+def getFormDiscontinuousGalerkin(F, L, Q, t, dt, u0, stages, bcs=None):
     """Given a time-dependent variational form, trial and test spaces, and
     a quadrature rule, produce UFL for the Discontinuous Galerkin-in-Time method.
 
@@ -121,6 +121,20 @@ def getFormDiscGalerkin(F, L, Q, t, dt, u0, stages, bcs=None):
     return Fnew, bcsnew
 
 
+def getElementDiscontinuousGalerkin(basis_type, order):
+    ufc_line = ufc_simplex(1)
+    if order == 0:
+        return DiscontinuousLagrange(ufc_line, order)
+    elif basis_type == "Bernstein":
+        return DiscontinuousElement(Bernstein(ufc_line, order))
+    elif basis_type == "integral":
+        return Legendre(ufc_line, order)
+    else:
+        # Let recursivenodes handle the general case
+        variant = None if basis_type == "Lagrange" else basis_type
+        return DiscontinuousLagrange(ufc_line, order, variant=variant)
+
+
 class DiscontinuousGalerkinTimeStepper(StageCoupledTimeStepper):
     """Front-end class for advancing a time-dependent PDE via a Discontinuous Galerkin
     in time method
@@ -171,25 +185,15 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None,
         V = u0.function_space()
         self.num_fields = len(V)
 
-        ufc_line = ufc_simplex(1)
-
-        if order == 0:
-            self.el = DiscontinuousLagrange(ufc_line, order)
-        elif basis_type == "Bernstein":
-            self.el = DiscontinuousElement(Bernstein(ufc_line, order))
-        elif basis_type == "integral":
-            self.el = Legendre(ufc_line, order)
-        else:
-            # Let recursivenodes handle the general case
-            variant = None if basis_type == "Lagrange" else basis_type
-            self.el = DiscontinuousLagrange(ufc_line, order, variant=variant)
+        self.el = getElementDiscontinuousGalerkin(basis_type, order)
 
         if quadrature is None:
-            quadrature = make_quadrature(ufc_line, order+1)
+            ref_complex = self.el.get_reference_complex()
+            quadrature = create_quadrature(ref_complex, 2*order)
         self.quadrature = quadrature
         assert np.size(quadrature.get_points()) >= order+1
 
-        num_stages = order+1
+        num_stages = self.el.space_dimension()
 
         self.update_b = vecconst(self.el.tabulate(0, (1.0,))[(0,)])
 
@@ -198,9 +202,9 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None,
                          appctx=appctx, nullspace=nullspace)
 
     def get_form_and_bcs(self, stages):
-        return getFormDiscGalerkin(self.F, self.el,
-                                   self.quadrature, self.t, self.dt, self.u0, stages,
-                                   self.orig_bcs)
+        return getFormDiscontinuousGalerkin(self.F, self.el, self.quadrature,
+                                            self.t, self.dt, self.u0, stages,
+                                            self.orig_bcs)
 
     def _update(self):
         for i, u0bit in enumerate(self.u0.subfunctions):

irksome/galerkin_stepper.py

@@ -1,6 +1,7 @@
 from FIAT import (Bernstein, DiscontinuousElement, DiscontinuousLagrange,
                   IntegratedLegendre, Lagrange, Legendre,
-                  make_quadrature, ufc_simplex)
+                  ufc_simplex)
+from FIAT.quadrature_schemes import create_quadrature
 from ufl import zero
 from ufl.constantvalue import as_ufl
 from .base_time_stepper import StageCoupledTimeStepper
@@ -38,8 +39,8 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, bcs=None):
        - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
          on the Galerkin-in-time solution,
     """
-    assert L_test.get_reference_element() == Q.ref_el
-    assert L_trial.get_reference_element() == Q.ref_el
+    assert L_test.get_reference_element() <= Q.ref_el
+    assert L_trial.get_reference_element() <= Q.ref_el
     assert Q.ref_el.get_spatial_dimension() == 1
     assert L_trial.get_order() == L_test.get_order() + 1
 
@@ -107,6 +108,26 @@ def getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, bcs=None):
     return Fnew, bcsnew
 
 
+def getElementGalerkin(basis_type, order):
+    ufc_line = ufc_simplex(1)
+    if basis_type == "Bernstein":
+        trial_el = Bernstein(ufc_line, order)
+        if order == 1:
+            test_el = DiscontinuousLagrange(ufc_line, 0)
+        else:
+            test_el = DiscontinuousElement(Bernstein(ufc_line, order-1))
+    elif basis_type == "integral":
+        trial_el = IntegratedLegendre(ufc_line, order)
+        test_el = Legendre(ufc_line, order-1)
+    else:
+        # Let recursivenodes handle the general case
+        variant = None if basis_type == "Lagrange" else basis_type
+        trial_el = Lagrange(ufc_line, order, variant=variant)
+        test_el = DiscontinuousLagrange(ufc_line, order-1, variant=variant)
+
+    return trial_el, test_el
+
+
 class GalerkinTimeStepper(StageCoupledTimeStepper):
     """Front-end class for advancing a time-dependent PDE via a Galerkin
     in time method
@@ -157,29 +178,17 @@ def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None,
         V = u0.function_space()
         self.num_fields = len(V)
 
-        ufc_line = ufc_simplex(1)
-        if basis_type == "Bernstein":
-            self.trial_el = Bernstein(ufc_line, order)
-            if order == 1:
-                self.test_el = DiscontinuousLagrange(ufc_line, 0)
-            else:
-                self.test_el = DiscontinuousElement(
-                    Bernstein(ufc_line, order-1))
-        elif basis_type == "integral":
-            self.trial_el = IntegratedLegendre(ufc_line, order)
-            self.test_el = Legendre(ufc_line, order-1)
-        else:
-            # Let recursivenodes handle the general case
-            variant = None if basis_type == "Lagrange" else basis_type
-            self.trial_el = Lagrange(ufc_line, order, variant=variant)
-            self.test_el = DiscontinuousLagrange(ufc_line, order-1, variant=variant)
-
+        self.trial_el, self.test_el = getElementGalerkin(basis_type, order)
         if quadrature is None:
-            quadrature = make_quadrature(ufc_line, order)
+            ref_complex = self.test_el.get_reference_complex()
+            quadrature_degree = self.trial_el.degree() + self.test_el.degree()
+            quadrature = create_quadrature(ref_complex, quadrature_degree)
         self.quadrature = quadrature
         assert np.size(quadrature.get_points()) >= order
 
-        super().__init__(F, t, dt, u0, order, bcs=bcs,
+        num_stages = self.test_el.space_dimension()
+
+        super().__init__(F, t, dt, u0, num_stages, bcs=bcs,
                          solver_parameters=solver_parameters,
                          appctx=appctx, nullspace=nullspace)
 

irksome/pc.py

@@ -1,6 +1,7 @@
 import copy
 
 import numpy
+from FIAT.quadrature_schemes import create_quadrature
 from firedrake import AuxiliaryOperatorPC, derivative
 from firedrake.dmhooks import get_appctx
 from ufl import replace
@@ -9,6 +10,9 @@
 from irksome.stage_derivative import getForm
 from irksome.stage_value import getFormStage
 
+from irksome.galerkin_stepper import getFormGalerkin, getElementGalerkin
+from irksome.discontinuous_galerkin_stepper import getFormDiscontinuousGalerkin, getElementDiscontinuousGalerkin
+
 
 # Oddly, we can't turn pivoting off in scipy?
 def ldu(A):
@@ -197,3 +201,60 @@ def getAtildes(self, A, Abar):
                 "ClinesLD preconditioner failed for for this tableau.  Please try again with GaussLegendre or RadauIIA methods")
         Abartilde = Lbar @ Dbar
         return Atilde, Abartilde
+
+
+class BaseGalerkinPC(AuxiliaryOperatorPC):
+    """Base class that inherits from Firedrake's AuxiliaryOperatorPC class and
+    provides the preconditioning bilinear form associated with an auxiliary
+    Form and/or equivalent finite element (which are provided by subclasses).
+    """
+
+    def getNewForm(self, pc, u0, test):
+        """Derived classes can optionally provide an auxiliary Form."""
+        raise NotImplementedError
+
+    def get_stage_residual(self, F, t, dt, u0, stages, bcs):
+        raise NotImplementedError
+
+    def form(self, pc, test, trial):
+        """Implements the interface for AuxiliaryOperatorPC."""
+        appctx = self.get_appctx(pc)
+        F = appctx["F"]
+        t = appctx["t"]
+        dt = appctx["dt"]
+        u0 = appctx["u0"]
+        bcs = appctx["bcs"]
+        v0, = F.arguments()
+
+        try:
+            # use new Form if provided
+            F, bcs = self.getNewForm(pc, u0, v0)
+        except NotImplementedError:
+            pass
+
+        # get stages
+        ctx = get_appctx(pc.getDM())
+        stages = ctx._x
+        Fnew, bcnew = self.get_stage_residual(F, t, dt, u0, stages, bcs, appctx)
+        # Now we get the Jacobian for the modified system,
+        # which becomes the auxiliary operator!
+        Jnew = derivative(Fnew, stages, du=trial)
+        return Jnew, bcnew
+
+
+class LowOrderRefinedGalerkinPC(BaseGalerkinPC):
+    def get_stage_residual(self, F, t, dt, u0, stages, bcs, appctx):
+        order = appctx["num_stages"]
+        basis_type = f"iso({order})"
+        L_trial, L_test = getElementGalerkin(basis_type, 1)
+        Q = create_quadrature(L_test.get_reference_complex(), L_trial.degree() + L_test.degree())
+        return getFormGalerkin(F, L_trial, L_test, Q, t, dt, u0, stages, bcs)
+
+
+class LowOrderRefinedDiscontinuousGalerkinPC(BaseGalerkinPC):
+    def get_stage_residual(self, F, t, dt, u0, stages, bcs, appctx):
+        order = appctx["num_stages"] - 1
+        basis_type = f"iso({order})"
+        L = getElementDiscontinuousGalerkin(basis_type, 0)
+        Q = create_quadrature(L.get_reference_complex(), 2*L.degree())
+        return getFormDiscontinuousGalerkin(F, L, Q, t, dt, u0, stages, bcs)

tests/test_disc_galerkin.py

@@ -7,9 +7,7 @@
 import FIAT
 
 
-@pytest.mark.parametrize("order", [0, 1, 2])
-@pytest.mark.parametrize("basis_type", ["Lagrange", "Bernstein", "spectral", "integral"])
-def test_1d_heat_dirichletbc(order, basis_type):
+def run_1d_heat_dirichletbc(order, basis_type, solver_parameters=None):
     # Boundary values
     u_0 = Constant(2.0)
     u_1 = Constant(3.0)
@@ -49,11 +47,13 @@ def test_1d_heat_dirichletbc(order, basis_type):
         DirichletBC(V, u_0, 1),
     ]
 
-    luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu"}
+    if solver_parameters is None:
+        luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu"}
+        solver_parameters = luparams
 
     stepper = DiscontinuousGalerkinTimeStepper(
         F, order, t, dt, u, bcs=bc, basis_type=basis_type,
-        solver_parameters=luparams
+        solver_parameters=solver_parameters
     )
 
     t_end = 2.0
@@ -68,6 +68,17 @@ def test_1d_heat_dirichletbc(order, basis_type):
         assert isclose(u.at(x1), u_1)
 
 
+@pytest.mark.parametrize("order", [0, 1, 2])
+@pytest.mark.parametrize("basis_type", ["Lagrange", "Bernstein", "spectral", "integral"])
+def test_1d_heat_dirichletbc(order, basis_type):
+    run_1d_heat_dirichletbc(order, basis_type)
+
+
+@pytest.mark.parametrize("order", [1, 2])
+def test_1d_heat_dirichletbc_iso(order):
+    run_1d_heat_dirichletbc(0, f"iso({order})")
+
+
 @pytest.mark.parametrize("order", [0, 1, 2])
 def test_1d_heat_neumannbc(order):
     N = 20

tests/test_galerkin.py

@@ -2,14 +2,12 @@
 
 import pytest
 from firedrake import *
-from irksome import Dt, MeshConstant, GalerkinTimeStepper
-from irksome import TimeStepper, GaussLegendre
+from irksome import (Dt, MeshConstant, GalerkinTimeStepper,
+                     TimeStepper, GaussLegendre)
 from FIAT import make_quadrature, ufc_simplex
 
 
-@pytest.mark.parametrize("order", [1, 2, 3])
-@pytest.mark.parametrize("basis_type", ["Lagrange", "Bernstein", "integral"])
-def test_1d_heat_dirichletbc(order, basis_type):
+def run_1d_heat_dirichletbc(order, basis_type, solver_parameters=None):
     # Boundary values
     u_0 = Constant(2.0)
     u_1 = Constant(3.0)
@@ -49,24 +47,61 @@ def test_1d_heat_dirichletbc(order, basis_type):
         DirichletBC(V, u_0, 1),
     ]
 
-    luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
+    if solver_parameters is None:
+        luparams = {"mat_type": "aij", "ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
+        solver_parameters = luparams
 
     stepper = GalerkinTimeStepper(
         F, order, t, dt, u, bcs=bc, basis_type=basis_type,
-        solver_parameters=luparams
+        solver_parameters=solver_parameters
     )
 
     t_end = 2.0
+    stats = []
     while float(t) < t_end:
         if float(t) + float(dt) > t_end:
             dt.assign(t_end - float(t))
         stepper.advance()
         t.assign(float(t) + float(dt))
         # Check solution and boundary values
+        stats.append(stepper.solver_stats())
         assert errornorm(uexact, u) / norm(uexact) < 10.0 ** -3
         assert isclose(u.at(x0), u_0)
         assert isclose(u.at(x1), u_1)
 
+    linits = [stat[2] / stat[0] / stat[1] for stat in stats]
+    return max(linits)
+
+
+@pytest.mark.parametrize("order", [1, 2, 3])
+@pytest.mark.parametrize("basis_type", ["Lagrange", "Bernstein", "integral"])
+def test_1d_heat_dirichletbc(order, basis_type):
+    run_1d_heat_dirichletbc(order, basis_type)
+
+
+@pytest.mark.parametrize("order", [2, 3])
+def test_1d_heat_dirichletbc_iso(order):
+    run_1d_heat_dirichletbc(1, f"iso({order})")
+
+
+@pytest.mark.parametrize("order", [2, 3])
+def test_low_order_refined_galerkin_pc(order):
+    params = {"mat_type": "matfree",
+              "snes_type": "ksponly",
+              "snes_lag_jacobian": -2,
+              "ksp_type": "gmres",
+              "ksp_converged_reason": None,
+              "pc_type": "python",
+              "pc_python_type": "irksome.LowOrderRefinedGalerkinPC",
+              "aux_mat_type": "aij",
+              "aux_pc_type": "lu",
+              "aux_pc_factor_mat_solver_type": "mumps",
+              }
+
+    basis_type = "spectral"
+    its = run_1d_heat_dirichletbc(order, basis_type, solver_parameters=params)
+    assert its <= order + 2
+
 
 @pytest.mark.parametrize("order", [1, 2, 3])
 @pytest.mark.parametrize("num_quad_points", [3, 4])