Support Bernstein on tensor product cells (#165)

* Support Bernstein on tensor product cells

* Add a test

* Rename lmbda -> make_finat_element

Author: pbrubeck

finat/element_factory.py

@@ -157,11 +157,18 @@ def convert_finiteelement(element, **kwargs):
 
         codim = 1 if element.family() == "Boundary Quadrature" else 0
         return finat.make_quadrature_element(cell, degree, scheme, codim), set()
-    lmbda = supported_elements[element.family()]
-    if element.family() == "Real" and element.cell.cellname() in {"quadrilateral", "hexahedron"}:
-        lmbda = None
-        element = finat.ufl.FiniteElement("DQ", element.cell, 0)
-    if lmbda is None:
+
+    make_finat_element = supported_elements[element.family()]
+
+    if element.cell.cellname() in {"quadrilateral", "hexahedron"}:
+        # Reconstruct Real and Bernstein on tensor product cells
+        if element.family() == "Real":
+            make_finat_element = None
+            element = finat.ufl.FiniteElement("DQ", element.cell, 0)
+        elif element.family() == "Bernstein":
+            make_finat_element = None
+
+    if make_finat_element is None:
         if element.cell.cellname() == "quadrilateral":
             # Handle quadrilateral short names like RTCF and RTCE.
             element = element.reconstruct(cell=quadrilateral_tpc)
@@ -174,49 +181,51 @@ def convert_finiteelement(element, **kwargs):
         finat_elem, deps = _create_element(element, **kwargs)
         return finat.FlattenedDimensions(finat_elem), deps
 
+    deps = set()
     finat_kwargs = {}
     kind = element.variant()
     if kind is None:
         kind = 'spectral'  # default variant
 
     if element.family() == "Lagrange":
         if kind in ['spectral', 'mimetic']:
-            lmbda = finat.GaussLobattoLegendre
+            make_finat_element = finat.GaussLobattoLegendre
         elif element.cell.cellname() == "interval" and kind in cg_interval_variants:
-            lmbda = cg_interval_variants[kind]
+            make_finat_element = cg_interval_variants[kind]
         elif any(map(kind.startswith, ['integral', 'demkowicz', 'fdm'])):
-            lmbda = finat.IntegratedLegendre
+            make_finat_element = finat.IntegratedLegendre
             finat_kwargs["variant"] = kind
         elif kind in ['mgd', 'feec', 'qb', 'mse']:
-            degree = element.degree()
-            shift_axes = kwargs["shift_axes"]
-            restriction = kwargs["restriction"]
+            make_finat_element = finat.RuntimeTabulated
+            finat_kwargs["variant"] = kind
+            finat_kwargs["shift_axes"] = kwargs["shift_axes"]
+            finat_kwargs["restriction"] = kwargs["restriction"]
             deps = {"shift_axes", "restriction"}
-            return finat.RuntimeTabulated(cell, degree, variant=kind, shift_axes=shift_axes, restriction=restriction), deps
         else:
             # Let FIAT handle the general case
-            lmbda = finat.Lagrange
+            make_finat_element = finat.Lagrange
             finat_kwargs["variant"] = kind
 
     elif element.family() in ["Discontinuous Lagrange", "Discontinuous Lagrange L2"]:
         if kind == 'spectral':
-            lmbda = finat.GaussLegendre
+            make_finat_element = finat.GaussLegendre
         elif kind == 'mimetic':
-            lmbda = finat.Histopolation
+            make_finat_element = finat.Histopolation
         elif element.cell.cellname() == "interval" and kind in dg_interval_variants:
-            lmbda = dg_interval_variants[kind]
+            make_finat_element = dg_interval_variants[kind]
         elif any(map(kind.startswith, ['integral', 'demkowicz', 'fdm'])):
-            lmbda = finat.Legendre
+            make_finat_element = finat.Legendre
             finat_kwargs["variant"] = kind
         elif kind in ['mgd', 'feec', 'qb', 'mse']:
-            degree = element.degree()
-            shift_axes = kwargs["shift_axes"]
-            restriction = kwargs["restriction"]
+            make_finat_element = finat.RuntimeTabulated
+            finat_kwargs["variant"] = kind
+            finat_kwargs["shift_axes"] = kwargs["shift_axes"]
+            finat_kwargs["restriction"] = kwargs["restriction"]
+            finat_kwargs["continuous"] = False
             deps = {"shift_axes", "restriction"}
-            return finat.RuntimeTabulated(cell, degree, variant=kind, shift_axes=shift_axes, restriction=restriction, continuous=False), deps
         else:
             # Let FIAT handle the general case
-            lmbda = finat.DiscontinuousLagrange
+            make_finat_element = finat.DiscontinuousLagrange
             finat_kwargs["variant"] = kind
 
     elif element.family() in {"HDiv Trace", "Bubble", "FacetBubble"}:
@@ -225,7 +234,7 @@ def convert_finiteelement(element, **kwargs):
     elif element.variant() is not None:
         finat_kwargs["variant"] = element.variant()
 
-    return lmbda(cell, element.degree(), **finat_kwargs), set()
+    return make_finat_element(cell, element.degree(), **finat_kwargs), deps
 
 
 # Element modifiers and compound element types

finat/ufl/elementlist.py

@@ -152,7 +152,7 @@ def show_elements():
 register_alias("Lob",
                lambda family, dim, order, degree: ("Gauss-Lobatto-Legendre", order))
 
-register_element("Bernstein", None, 0, H1, "identity", (1, None), simplices)
+register_element("Bernstein", None, 0, H1, "identity", (1, None), any_cell)
 
 
 # Let Nedelec H(div) elements be aliases to BDMs/RTs

finat/ufl/finiteelement.py

@@ -94,8 +94,10 @@ def __new__(cls,
                 return EnrichedElement(HCurl(TensorProductElement(Qc_elt, Id_elt, cell=cell)),
                                        HCurl(TensorProductElement(Qd_elt, Ic_elt, cell=cell)))
 
-            elif family == "Q":
-                return TensorProductElement(*[FiniteElement("CG", c, degree, variant=variant)
+            elif family in {"Q", "Bernstein"}:
+                if family == "Q":
+                    family = "CG"
+                return TensorProductElement(*[FiniteElement(family, c, degree, variant=variant)
                                               for c in cell.sub_cells()],
                                             cell=cell)
 

test/finat/test_create_finat_element.py

@@ -106,6 +106,12 @@ def test_quadrilateral_variant_spectral_q():
     assert isinstance(element.product.factors[1], finat.GaussLobattoLegendre)
 
 
+def test_quadrilateral_bernstein():
+    element = create_element(finat.ufl.FiniteElement('Bernstein', ufl.quadrilateral, 3))
+    assert isinstance(element.product.factors[0], finat.Bernstein)
+    assert isinstance(element.product.factors[1], finat.Bernstein)
+
+
 def test_quadrilateral_variant_spectral_dq():
     element = create_element(finat.ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
     assert isinstance(element.product.factors[0], finat.GaussLegendre)