Support Guzman-Neilan and Bernardi-Raugel

alfi/bubble.py

@@ -1,12 +1,24 @@
 from firedrake import *
-from firedrake.assemble import OneFormAssembler
+from functools import partial
 
 """ For the P1 + FacetBubble space in 3D, the standard prolongation does not
 preserve the flux across coarse grid facets.  We fix this with a manual
 scaling. """
 
 
 class BubbleTransfer(object):
+    def __new__(cls, Vc, Vf):
+        if cls is not BubbleTransfer:
+            return super().__new__(cls)
+
+        fe = Vc.finat_element
+        entity_dofs = fe.entity_dofs()
+        dim = fe.cell.get_spatial_dimension()
+        if fe.value_shape == (dim,) and len(entity_dofs[dim-1][0]) == 1:
+            return NormalBubbleTransfer(Vc, Vf)
+
+        return super().__new__(cls)
+
     def __init__(self, Vc, Vf):
         meshc = Vc.mesh()
         meshf = Vf.mesh()
@@ -35,8 +47,7 @@ def __init__(self, Vc, Vf):
         self.ainv = ainv
         L = inner(inner(self.fbc, n)/0.625, inner(test, n)) + inner(self.fbc - inner(self.fbc, n)*n, test-inner(test, n)*n)
         L = L*ds + avg(L) * dS
-        # This is the guts of assemble(L, tensor=self.rhs), but saves a little time
-        self.assemble_rhs = OneFormAssembler(L, tensor=self.rhs).assemble
+        self.assemble_rhs = partial(assemble, L, tensor=self.rhs)
 
         # TODO: Parameterise these by the element definition.
         # These are correct for P1 + FB in 3D.
@@ -263,3 +274,41 @@ def prolong(self, coarse, fine):
                      self.fbf.dat(op2.READ, self.fbf.cell_node_map()),
                      fine.dat(op2.INC, fine.cell_node_map()))
         fine.dat.data[:, :] /= self.countvf.dat.data[:, :]
+
+
+class NormalBubbleTransfer(BubbleTransfer):
+    manager = TransferManager()
+
+    def __init__(self, Vc, Vf):
+        W = Vc
+        depth = W.mesh().topological_dimension() - 1
+
+        mesh_dm = W.mesh().topology_dm
+        W_local_ises_indices = W.dof_dset.local_ises[0].indices
+        section = W.dm.getDefaultSection()
+        indices = []
+        for p in range(*mesh_dm.getDepthStratum(depth)):
+            dof = section.getDof(p)
+            if dof <= 0:
+                continue
+            off = section.getOffset(p)
+            W_indices = slice(off*W.value_size, W.value_size * (off + dof))
+            indices.extend(W_local_ises_indices[W_indices])
+
+        ainv = W.dof_dset.layout_vec.copy()
+        ainv.set(1)
+        ainv[indices] = 1 / 0.625
+        self.ainv = ainv
+
+        self.primal = Function(W)
+        self.dual = Cofunction(W.dual())
+
+    def restrict(self, rf, rc):
+        self.manager.restrict(rf, self.dual)
+        with self.dual.dat.vec_wo as wc, rc.dat.vec_ro as b:
+            wc.pointwiseMult(b, self.ainv)
+
+    def prolong(self, uc, uf):
+        with self.primal.dat.vec_wo as wc, uc.dat.vec_ro as b:
+            wc.pointwiseMult(b, self.ainv)
+        self.manager.prolong(self.primal, uf)

alfi/driver.py

@@ -26,7 +26,7 @@ def get_default_parser():
     parser.add_argument("--stabilisation-type", type=str, default=None,
                         choices=["none", "burman", "gls", "supg"])
     parser.add_argument("--discretisation", type=str, required=True,
-                        choices=["pkp0", "sv"])
+                        choices=["pkp0", "gn", "br", "sv"])
     parser.add_argument("--gamma", type=float, default=1e4)
     parser.add_argument("--clear", dest="clear", default=False,
                         action="store_true")
@@ -50,9 +50,12 @@ def get_default_parser():
 
 def get_solver(args, problem, hierarchy_callback=None):
     solver_t = {"pkp0": ConstantPressureSolver,
+                "gn": ConstantPressureSolver,
+                "br": ConstantPressureSolver,
                 "sv": ScottVogeliusSolver}[args.discretisation]
     solver = solver_t(
         problem,
+        family=args.discretisation,
         solver_type=args.solver_type,
         stabilisation_type=args.stabilisation_type,
         nref=args.nref,
@@ -119,7 +122,7 @@ def run_solver(solver, res, args):
                 with DumbCheckpoint(chkptdir + "nssolution-Re-%s" % (re), mode=FILE_UPDATE) as checkpoint:
                     checkpoint.store(solver.z, name="up_%i" % re)
         if args.paraview:
-            pvdf.write(solver.visprolong(solver.z.split()[0]), time=re)
+            pvdf.write(solver.visprolong(solver.z.subfunctions[0]), time=re)
 
     if comm.rank == 0:
         for re in results:

alfi/relaxation.py

@@ -76,7 +76,7 @@ def visualise(self, mesh, vertex=0, color="blue"):
 
     @staticmethod
     def get_entities(opts, name, dm):
-        sentinel = object()
+        sentinel = -111
         codim = opts.getInt("pc_patch_construction_%s_codim" % name, default=sentinel)
         if codim == sentinel:
             dim = opts.getInt("pc_patch_construction_%s_dim" % name, default=0)
@@ -86,7 +86,7 @@ def get_entities(opts, name, dm):
         return entities
 
     def keyfuncs(self, coords):
-        sentinel = object()
+        sentinel = -111
         sortorders = self.opts.getString("pc_patch_construction_%s_sort_order" % self.name, default=sentinel)
         if sortorders is None or sortorders in [sentinel, "None", ""]:
             return None

alfi/solver.py

@@ -45,14 +45,14 @@ def function_space(self, mesh, k):
 
     def residual(self):
         raise NotImplementedError
-    
+
     def update_wind(self, z):
         raise NotImplementedError
 
     def set_transfers(self):
         raise NotImplementedError
 
-    def __init__(self, problem, nref=1, solver_type="almg",
+    def __init__(self, problem, family=None, nref=1, solver_type="almg",
                  stabilisation_type=None,
                  supg_method="shakib", supg_magic=9.0, gamma=10000, nref_vis=1,
                  k=5, patch="star", hierarchy="bary", use_mkl=False, stabilisation_weight=None,
@@ -62,6 +62,8 @@ def __init__(self, problem, nref=1, solver_type="almg",
                  high_accuracy=False
                  ):
 
+        assert family in {"sv", "pkp0", "gn", "br"}
+        self.family = family
         assert solver_type in {"almg", "allu", "lu", "simple", "lsc"}, "Invalid solver type %s" % solver_type
         if stabilisation_type == "none":
             stabilisation_type = None
@@ -115,7 +117,7 @@ def after(dm, i):
         self.parallel = mh[0].comm.size > 1
         self.tdim = mh[0].topological_dimension()
         self.mh = mh
-        self.area = assemble(Constant(1, domain=mh[0])*dx)
+        self.area = assemble(1*dx(domain=mh[0]))
         nu = Constant(1.0)
         self.nu = nu
         self.char_L = problem.char_length()
@@ -132,7 +134,7 @@ def after(dm, i):
 
         mesh = mh[-1]
         uviss = []
- 
+
         self.mesh = mesh
         self.load_balance(mesh)
         Z = self.function_space(mesh, k)
@@ -171,8 +173,8 @@ def visprolong(u):
         if comm.rank == 0:
             print("Number of velocity degrees of freedom: %s (avg %.2f per core)" % (Vdim, Vdim / size))
         z = Function(Z, name="Solution")
-        z.split()[0].rename("Velocity")
-        z.split()[1].rename("Pressure")
+        z.subfunctions[0].rename("Velocity")
+        z.subfunctions[1].rename("Pressure")
         self.z = z
         (u, p) = split(z)
         (v, q) = split(TestFunction(Z))
@@ -268,14 +270,14 @@ def solve(self, re):
         # self.gamma.assign(1+re)
 
         if self.stabilisation is not None:
-            self.stabilisation.update(self.z.split()[0])
+            self.stabilisation.update(self.z.subfunctions[0])
         start = datetime.now()
         self.solver.solve()
         end = datetime.now()
 
         if self.nsp is not None:
             # Hardcode that pressure integral is zero
-            (u, p) = self.z.split()
+            (u, p) = self.z.subfunctions
             pintegral = assemble(p*dx)
             p.assign(p - Constant(pintegral/self.area))
 
@@ -502,7 +504,7 @@ def get_parameters(self):
 
         if self.solver_type == "lu":
             outer = {**outer_base, **outer_lu}
-        elif self.solver_type == "simple": 
+        elif self.solver_type == "simple":
             outer = {**outer_base, **outer_simple}
         elif self.solver_type == "lsc":
             outer = {**outer_base, **outer_lsc}
@@ -559,12 +561,23 @@ class ConstantPressureSolver(NavierStokesSolver):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    def augmented_lagrangian(self, u, v):
+        from finat.quadrature import make_quadrature
+        V = self.Z[0]
+        if V.ufl_element().family() == "Guzman-Neilan":
+            # Override macro quadrature rule
+            cell = V.finat_element.cell
+            al = inner(div(u), div(v))*dx(scheme=make_quadrature(cell, 0))
+        else:
+            al = inner(cell_avg(div(u)), div(v))*dx(metadata={"mode": "vanilla"})
+        return al
+
     def residual(self):
         u, p = split(self.z)
         v, q = TestFunctions(self.Z)
         F = (
             self.nu * inner(2*sym(grad(u)), grad(v))*dx
-            + self.gamma * inner(cell_avg(div(u)), div(v))*dx(metadata={"mode": "vanilla"})
+            + self.gamma * self.augmented_lagrangian(u, v)
             + self.advect * inner(dot(grad(u), u), v)*dx
             - p * div(v) * dx
             - div(u) * q * dx
@@ -574,9 +587,14 @@ def residual(self):
     def function_space(self, mesh, k):
         tdim = mesh.topological_dimension()
         if k < tdim:
-            Pk = FiniteElement("Lagrange", mesh.ufl_cell(), k)
-            FB = FiniteElement("FacetBubble", mesh.ufl_cell(), tdim)
-            eleu = VectorElement(NodalEnrichedElement(Pk, FB))
+            if self.family == "br":
+                eleu = FiniteElement("BR", mesh.ufl_cell(), k)
+            elif self.family == "gn":
+                eleu = FiniteElement("GN", mesh.ufl_cell(), k)
+            else:
+                Pk = FiniteElement("Lagrange", mesh.ufl_cell(), k)
+                FB = FiniteElement("FacetBubble", mesh.ufl_cell(), tdim)
+                eleu = VectorElement(NodalEnrichedElement(Pk, FB))
         else:
             Pk = FiniteElement("Lagrange", mesh.ufl_cell(), k)
             eleu = VectorElement(Pk)
@@ -592,7 +610,7 @@ def get_transfers(self):
         qtransfer = NullTransfer()
         self.vtransfer = vtransfer
         self.qtransfer = qtransfer
-        transfers = {V.ufl_element(): (vtransfer.prolong, vtransfer.restrict if self.restriction else restrict, inject),
+        transfers = {V.ufl_element(): (vtransfer.prolong, vtransfer.restrict if self.restriction else None, None),
                      Q.ufl_element(): (prolong, restrict, qtransfer.inject)}
         return transfers
 

alfi/transfer.py

@@ -89,6 +89,8 @@ def __call__(self, pc):
 
 
 class AutoSchoeberlTransfer(object):
+    manager = TransferManager()
+
     def __init__(self, parameters, tdim, hierarchy):
         self.solver = {}
         self.bcs = {}
@@ -153,7 +155,7 @@ def nodes(self):
                 return self.nodelist
 
         dim = V.mesh().topological_dimension()
-        bc = FixedDirichletBC(V, ufl.zero(V.ufl_element().value_shape()), nodelist)
+        bc = FixedDirichletBC(V, 0, nodelist)
 
         return bc
 
@@ -199,7 +201,7 @@ def restrict_or_prolong(self, source, target, mode):
             fine = source
             coarse = target
         # Rebuild without any indices
-        V = FunctionSpace(fine.ufl_domain(), fine.function_space().ufl_element())
+        V = fine.function_space().collapse()
         key = V.dim()
 
         firsttime = self.bcs.get(key, None) is None
@@ -213,6 +215,7 @@ def restrict_or_prolong(self, source, target, mode):
             tildeu, rhs = Function(V), Function(V)
 
             bform = self.bform(rhs)
+            # TODO use Cofunction and LinearVariationalSolver
             b = Function(V)
             problem = LinearVariationalProblem(a=a, L=0, u=tildeu, bcs=bcs)
             ctx = _SNESContext(problem, mat_type=self.patchparams["mat_type"],
@@ -283,9 +286,9 @@ def restrict_or_prolong(self, source, target, mode):
 
     def standard_transfer(self, source, target, mode):
         if mode == "prolong":
-            prolong(source, target)
+            self.manager.prolong(source, target)
         elif mode == "restrict":
-            restrict(source, target)
+            self.manager.restrict(source, target)
         else:
             raise NotImplementedError
 
@@ -315,20 +318,30 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.transfers = {}
 
+    def augmented_lagrangian(self, u, v):
+        from finat.quadrature import make_quadrature
+        V = v.function_space()
+        if V.ufl_element().family() == "Guzman-Neilan":
+            # Override macro quadrature rule
+            cell = V.finat_element.cell
+            al = inner(div(u), div(v))*dx(scheme=make_quadrature(cell, 0))
+        else:
+            al = inner(cell_avg(div(u)), div(v))*dx(metadata={"mode": "vanilla"})
+        return al
 
     def form(self, V):
         (nu, gamma) = self.parameters
         u = TrialFunction(V)
         v = TestFunction(V)
-        a = nu * inner(2*sym(grad(u)), grad(v))*dx + gamma*inner(cell_avg(div(u)), div(v))*dx(metadata={"mode": "vanilla"})
+        a = nu * inner(2*sym(grad(u)), grad(v))*dx + gamma * self.augmented_lagrangian(u, v)
         return a
 
     def bform(self, rhs):
         V = rhs.function_space()
         (nu, gamma) = self.parameters
         u = TrialFunction(V)
         v = TestFunction(V)
-        a = gamma*inner(cell_avg(div(u)), div(v))*dx(metadata={"mode": "vanilla"})
+        a = gamma * self.augmented_lagrangian(u, v)
         return action(a, rhs)
 
     def standard_transfer(self, source, target, mode):

examples/Makefile

@@ -1,4 +1,4 @@
-exec=mpirun -n 12 python3
+exec=mpiexec -n 1 python3
 pre:
 	mkdir -p logs/
 
@@ -20,6 +20,14 @@ mms3dpkp0: pre
 	$(exec) mms.py --dim 3 --k 1 --nref 5 --discretisation pkp0 --solver-type almg --mh uniform --patch star --gamma 1e4 --stabilisation-type supg --stabilisation--weight 0.05 --baseN 4 2>&1 | tee logs/mms3dpkp0highgamma.log
 	$(exec) mms.py --dim 3 --k 1 --nref 5 --discretisation pkp0 --solver-type almg --mh uniform --patch star --gamma 1e0 --stabilisation-type supg --stabilisation--weight 0.05 --baseN 4 2>&1 | tee logs/mms3dpkp0lowgamma.log
 
+mms3dbr: pre
+	$(exec) mms.py --dim 3 --k 1 --nref 5 --discretisation br --solver-type almg --mh uniform --patch star --gamma 1e4 --stabilisation-type supg --stabilisation--weight 0.05 --baseN 4 2>&1 | tee logs/mms3dbrhighgamma.log
+	$(exec) mms.py --dim 3 --k 1 --nref 5 --discretisation br --solver-type almg --mh uniform --patch star --gamma 1e0 --stabilisation-type supg --stabilisation--weight 0.05 --baseN 4 2>&1 | tee logs/mms3dbrlowgamma.log
+
+mms3dgn: pre
+	$(exec) mms.py --dim 3 --k 1 --nref 5 --discretisation gn --solver-type almg --mh uniform --patch star --gamma 1e4 --stabilisation-type supg --stabilisation--weight 0.05 --baseN 4 2>&1 | tee logs/mms3dgnhighgamma.log
+	$(exec) mms.py --dim 3 --k 1 --nref 5 --discretisation gn --solver-type almg --mh uniform --patch star --gamma 1e0 --stabilisation-type supg --stabilisation--weight 0.05 --baseN 4 2>&1 | tee logs/mms3dgnlowgamma.log
+
 mms2dbary: pre
 	$(exec) mms.py --dim 2 --nref 6 --k 2 --discretisation sv   --solver-type almg --mh bary --patch macro --gamma 1e4 --stabilisation-type burman --stabilisation-weight 5e-3 2>&1 | tee logs/mms2dbarysv.log
 	$(exec) mms.py --dim 2 --nref 6 --k 2 --discretisation pkp0 --solver-type almg --mh bary --patch star  --gamma 1e4 --stabilisation-type burman --stabilisation-weight 5e-3 2>&1 | tee logs/mms2dbarypkp0.log

examples/graddiv/graddiv.py

@@ -5,6 +5,7 @@
 import argparse
 import numpy
 
+bubbles = {"pkp0", "br", "gn"}
 
 parser = get_default_parser()
 parser.add_argument("--dim", type=int, required=True, choices=[2, 3])
@@ -60,6 +61,10 @@ def after(dm, i):
     Pk = FiniteElement("Lagrange", base.ufl_cell(), k)
     FB = FiniteElement("FacetBubble", base.ufl_cell(), dim)
     eleu = VectorElement(NodalEnrichedElement(Pk, FB))
+elif args.discretisation == "br" and k < dim:
+    eleu = FiniteElement("BR", base.ufl_cell(), k)
+elif args.discretisation == "gn" and k < dim:
+    eleu = FiniteElement("GN", base.ufl_cell(), k)
 else:
     Pk = FiniteElement("Lagrange", base.ufl_cell(), k)
     eleu = VectorElement(Pk)
@@ -75,7 +80,7 @@ def after(dm, i):
 else:
     f = Constant((1, 1, 1))
     bclabels = [1, 2, 3, 4, 5, 6]
-if args.discretisation == "pkp0":
+if args.discretisation in bubbles:
     F = inner(2*sym(grad(u)), grad(v))*dx + gamma*inner(cell_avg(div(u)), div(v))*dx - inner(f, v)*dx
 else:
     F = inner(2*sym(grad(u)), grad(v))*dx + gamma*inner(div(u), div(v))*dx - inner(f, v)*dx
@@ -150,7 +155,7 @@ def after(dm, i):
 
 if args.discretisation == "sv":
     vtransfer = SVSchoeberlTransfer((1, gamma), args.dim, args.mh)
-elif args.discretisation == "pkp0":
+elif args.discretisation in bubbles:
     vtransfer = PkP0SchoeberlTransfer((1, gamma), args.dim, args.mh)
 nvproblem = NonlinearVariationalProblem(F, u, bcs=bcs)
 solver = NonlinearVariationalSolver(nvproblem, solver_parameters=sp, options_prefix="")