PC: fix max degree (#3401)

* PC: fix max degree

* HiptmairPC: fix degree of potential space for N2curl and N2div

firedrake/preconditioners/fdm.py

@@ -19,7 +19,7 @@
 from tsfc.finatinterface import create_element
 from pyop2.compilation import load
 from pyop2.sparsity import get_preallocation
-from pyop2.utils import get_petsc_dir
+from pyop2.utils import get_petsc_dir, as_tuple
 
 import firedrake.dmhooks as dmhooks
 import ufl
@@ -1595,11 +1595,7 @@ def assemble_coefficients(self, J, fcp):
         tdim = mesh.topological_dimension()
         Finv = ufl.JacobianInverse(mesh)
 
-        degree = V.ufl_element().degree()
-        try:
-            degree = max(degree)
-        except TypeError:
-            pass
+        degree = max(as_tuple(V.ufl_element().degree()))
         quad_deg = fcp.get("degree", 2*degree+1)
         dx = ufl.dx(degree=quad_deg, domain=mesh)
         family = "Discontinuous Lagrange" if tdim == 1 else "DQ"

firedrake/preconditioners/hiptmair.py

@@ -1,5 +1,6 @@
 import abc
 
+from pyop2.utils import as_tuple
 from firedrake.petsc import PETSc
 from firedrake.preconditioners.base import PCBase
 from firedrake.functionspace import FunctionSpace
@@ -23,9 +24,6 @@ class TwoLevelPC(PCBase):
     should implement:
     - :meth:`coarsen`
     """
-
-    needs_python_pmat = False
-
     @abc.abstractmethod
     def coarsen(self, pc):
         """Return a tuple with coarse bilinear form, coarse
@@ -162,42 +160,42 @@ def coarsen(self, pc):
         V = a.arguments()[-1].function_space()
         mesh = V.mesh()
         element = V.ufl_element()
-        degree = element.degree()
-        try:
-            degree = max(degree)
-        except TypeError:
-            pass
         formdegree = V.finat_element.formdegree
         if formdegree == 1:
             celement = curl_to_grad(element)
-            dminus = ufl.grad
-            G_callback = appctx.get("get_gradient", None)
         elif formdegree == 2:
             celement = div_to_curl(element)
-            dminus = ufl.curl
-            if V.shape:
-                dminus = lambda u: ufl.as_vector([ufl.curl(u[k, ...])
-                                                  for k in range(u.ufl_shape[0])])
-            G_callback = appctx.get("get_curl", None)
         else:
             raise ValueError("Hiptmair decomposition not available for", element)
 
         coarse_space = FunctionSpace(mesh, celement)
         assert coarse_space.finat_element.formdegree + 1 == formdegree
         coarse_space_bcs = tuple(bc.reconstruct(V=coarse_space, g=0) for bc in bcs)
 
+        if element.sobolev_space == ufl.HDiv:
+            G_callback = appctx.get("get_curl", None)
+            dminus = ufl.curl
+            if V.shape:
+                dminus = lambda u: ufl.as_vector([ufl.curl(u[k, ...])
+                                                  for k in range(u.ufl_shape[0])])
+        else:
+            G_callback = appctx.get("get_gradient", None)
+            dminus = ufl.grad
+
         # Get only the zero-th order term of the form
         replace_dict = {ufl.grad(t): ufl.zero(ufl.grad(t).ufl_shape) for t in a.arguments()}
         beta = ufl.replace(expand_derivatives(a), replace_dict)
 
         test = TestFunction(coarse_space)
         trial = TrialFunction(coarse_space)
-        coarse_operator = beta(dminus(test), dminus(trial), coefficients={})
+        coarse_operator = beta(dminus(test), dminus(trial))
 
-        if formdegree > 1 and degree > 1:
+        cdegree = max(as_tuple(celement.degree()))
+        if formdegree > 1 and cdegree > 1:
             shift = appctx.get("hiptmair_shift", None)
             if shift is not None:
-                coarse_operator += beta(test, shift*trial, coefficients={})
+                b = beta(test, shift * trial)
+                coarse_operator += ufl.Form(b.integrals_by_type("cell"))
 
         if G_callback is None:
             interp_petscmat = chop(Interpolator(dminus(test), V, bcs=bcs + coarse_space_bcs).callable().handle)
@@ -224,6 +222,8 @@ def curl_to_grad(ele):
         if family.startswith("Sminus"):
             family = "S"
         else:
+            if family in ["Nedelec 2nd kind H(curl)", "Brezzi-Douglas-Marini"]:
+                degree = degree + 1
             family = "CG"
             if isinstance(degree, tuple) and isinstance(cell, ufl.TensorProductCell):
                 cells = ele.cell.sub_cells()
@@ -258,11 +258,13 @@ def div_to_curl(ele):
         family = ele.family()
         if family in ["Lagrange", "CG", "Q"]:
             family = "DG" if ele.cell.is_simplex() else "DQ"
-            degree = degree-1
+            degree = degree - 1
         elif family in ["Discontinuous Lagrange", "DG", "DQ"]:
             family = "CG"
-            degree = degree+1
+            degree = degree + 1
         else:
+            if family == "Brezzi-Douglas-Marini":
+                degree = degree + 1
             replace_dict = {
                 "Raviart-Thomas": "N1curl",
                 "Brezzi-Douglas-Marini": "N2curl",

firedrake/preconditioners/hypre_ads.py

@@ -6,6 +6,7 @@
 from firedrake.dmhooks import get_function_space
 from firedrake.preconditioners.hypre_ams import chop
 from ufl import grad, curl, SpatialCoordinate
+from pyop2.utils import as_tuple
 
 __all__ = ("HypreADS",)
 
@@ -20,11 +21,7 @@ def initialize(self, obj):
 
         family = str(V.ufl_element().family())
         formdegree = V.finat_element.formdegree
-        degree = V.ufl_element().degree()
-        try:
-            degree = max(degree)
-        except TypeError:
-            pass
+        degree = max(as_tuple(V.ufl_element().degree()))
         if formdegree != 2 or degree != 1:
             raise ValueError("Hypre ADS requires lowest order RT elements! (not %s of degree %d)" % (family, degree))
 

firedrake/preconditioners/hypre_ams.py

@@ -9,6 +9,7 @@
 from firedrake_citations import Citations
 from firedrake import SpatialCoordinate
 from ufl import grad
+from pyop2.utils import as_tuple
 
 __all__ = ("HypreAMS",)
 
@@ -42,11 +43,7 @@ def initialize(self, obj):
 
         family = str(V.ufl_element().family())
         formdegree = V.finat_element.formdegree
-        degree = V.ufl_element().degree()
-        try:
-            degree = max(degree)
-        except TypeError:
-            pass
+        degree = max(as_tuple(V.ufl_element().degree()))
         if formdegree != 1 or degree != 1:
             raise ValueError("Hypre AMS requires lowest order Nedelec elements! (not %s of degree %d)" % (family, degree))
 

firedrake/preconditioners/pmg.py

@@ -13,6 +13,7 @@
 from tsfc import compile_expression_dual_evaluation
 from pyop2 import op2
 from pyop2.caching import cached
+from pyop2.utils import as_tuple
 
 import firedrake
 import finat
@@ -386,22 +387,7 @@ def create_injection(self, dmc, dmf):
     @staticmethod
     def max_degree(ele):
         """Return the maximum degree of a :class:`finat.ufl.finiteelement.FiniteElement`"""
-        if isinstance(ele, (finat.ufl.VectorElement, finat.ufl.TensorElement)):
-            return PMGBase.max_degree(ele._sub_element)
-        elif isinstance(ele, (finat.ufl.MixedElement, finat.ufl.TensorProductElement)):
-            return max(PMGBase.max_degree(sub) for sub in ele.sub_elements)
-        elif isinstance(ele, finat.ufl.EnrichedElement):
-            return max(PMGBase.max_degree(sub) for sub in ele._elements)
-        elif isinstance(ele, finat.ufl.WithMapping):
-            return PMGBase.max_degree(ele.wrapee)
-        elif isinstance(ele, (finat.ufl.HDivElement, finat.ufl.HCurlElement, finat.ufl.BrokenElement, finat.ufl.RestrictedElement)):
-            return PMGBase.max_degree(ele._element)
-        else:
-            degree = ele.degree()
-            try:
-                return max(degree)
-            except TypeError:
-                return degree
+        return max(as_tuple(ele.degree()))
 
     @staticmethod
     def reconstruct_degree(ele, degree):

tests/regression/test_fdm.py

@@ -1,5 +1,6 @@
 import pytest
 from firedrake import *
+from pyop2.utils import as_tuple
 
 ksp = {
     "mat_type": "matfree",
@@ -240,11 +241,7 @@ def test_ipdg_direct_solver(fs):
             u_exact = as_vector([u_exact + Constant(k) for k in range(ncomp)])
         u_bc = u_exact
 
-    degree = fs.ufl_element().degree()
-    try:
-        degree = max(degree)
-    except TypeError:
-        pass
+    degree = max(as_tuple(fs.ufl_element().degree()))
     quad_degree = 2*(degree+1)-1
 
     # problem coefficients