Add higher order Lagrange on pyramids (#95)

* Generalise pyramid lagrange

* flake

* Test pyramid lagrange

Author: mscroggs

symfem/core/polynomials.py

@@ -270,20 +270,14 @@ def pyramid_polynomial_set_1d(dim, order):
     assert dim == 3
     if order == 0:
         return [one]
-    if order >= 3:
-        raise NotImplementedError()
-
-    poly = []
-    for r in range(order + 1):
-        for p in range(order - r + 1):
-            for q in range(order - r + 1):
-                if r == 0 and p + q < order:
-                    poly.append(x[0] ** p * x[1] ** q)
-                else:
-                    new_p = (2 * x[0] + x[2]) ** p
-                    new_p *= (2 * x[1] + x[2]) ** q
-                    new_p *= x[2] ** r
-                    poly.append(new_p)
+
+    poly = polynomial_set_1d(3, order)
+
+    poly = [x[0] ** a * x[1] ** b * x[2] ** c / (1 - x[2]) ** (a + b + c - order)
+            for c in range(order)
+            for a in range(order + 1 - c) for b in range(order + 1 - c)]
+    poly.append(x[2] ** order)
+
     return poly
 
 

symfem/elements/lagrange_prism_pyramid.py renamed to symfem/elements/lagrange_prism.py

@@ -1,10 +1,10 @@
-"""Lagrange elements on a prism and pyramid."""
+"""Lagrange elements on a prism."""
 
 import sympy
 from itertools import product
 from ..core.symbolic import one, zero
 from ..core.finite_element import CiarletElement
-from ..core.polynomials import prism_polynomial_set, pyramid_polynomial_set
+from ..core.polynomials import prism_polynomial_set
 from ..core.functionals import PointEvaluation, DotPointEvaluation
 from ..core.quadrature import get_quadrature
 
@@ -51,35 +51,21 @@ def __init__(self, reference, order, variant):
                                       for j, o in enumerate(entity.origin)),
                                 entity=(2, e_n)))
 
-            if reference.name == "prism":
-                # Interior
-                for i in product(range(1, order), repeat=3):
-                    if i[0] + i[1] < order:
-                        dofs.append(
-                            PointEvaluation(
-                                tuple(o + sum(a[j] * points[b]
-                                              for a, b in zip(reference.axes, i))
-                                      for j, o in enumerate(reference.origin)),
-                                entity=(3, 0)))
-            elif reference.name == "pyramid":
-                # Interior
-                for i in product(range(1, order), repeat=3):
-                    if max(i[0], i[1]) + i[2] < order:
-                        dofs.append(
-                            PointEvaluation(
-                                tuple(o + sum(a[j] * points[b]
-                                              for a, b in zip(reference.axes, i))
-                                      for j, o in enumerate(reference.origin)),
-                                entity=(3, 0)))
-
-        if reference.name == "prism":
-            poly = prism_polynomial_set(reference.tdim, 1, order)
-        elif reference.name == "pyramid":
-            poly = pyramid_polynomial_set(reference.tdim, 1, order)
+            # Interior
+            for i in product(range(1, order), repeat=3):
+                if i[0] + i[1] < order:
+                    dofs.append(
+                        PointEvaluation(
+                            tuple(o + sum(a[j] * points[b]
+                                          for a, b in zip(reference.axes, i))
+                                  for j, o in enumerate(reference.origin)),
+                            entity=(3, 0)))
+
+        poly = prism_polynomial_set(reference.tdim, 1, order)
         super().__init__(reference, order, poly, dofs, reference.tdim, 1)
 
     names = ["Lagrange", "P"]
-    references = ["prism", "pyramid"]
+    references = ["prism"]
     min_order = 0
     continuity = "C0"
 

symfem/elements/lagrange_pyramid.py

@@ -0,0 +1,74 @@
+"""Lagrange elements on a pyramid.
+
+This element's definition appears in https://doi.org/10.1007/s10915-009-9334-9
+(Bergot, Cohen, Durufle, 2010)
+"""
+
+import sympy
+from itertools import product
+from ..core.finite_element import CiarletElement
+from ..core.polynomials import pyramid_polynomial_set
+from ..core.functionals import PointEvaluation
+from ..core.quadrature import get_quadrature
+
+
+class Lagrange(CiarletElement):
+    """Lagrange finite element."""
+
+    def __init__(self, reference, order, variant):
+        if order == 0:
+            dofs = [
+                PointEvaluation(
+                    tuple(
+                        sympy.Rational(1, reference.tdim + 1)
+                        for i in range(reference.tdim)
+                    ),
+                    entity=(reference.tdim, 0)
+                )
+            ]
+        else:
+            points, _ = get_quadrature(variant, order + 1)
+
+            dofs = []
+            # Vertices
+            for v_n, v in enumerate(reference.reference_vertices):
+                dofs.append(PointEvaluation(v, entity=(0, v_n)))
+            # Edges
+            for e_n in range(reference.sub_entity_count(1)):
+                entity = reference.sub_entity(1, e_n)
+                for i in range(1, order):
+                    dofs.append(
+                        PointEvaluation(
+                            tuple(o + entity.axes[0][j] * points[i]
+                                  for j, o in enumerate(entity.origin)),
+                            entity=(1, e_n)))
+            # Faces
+            for e_n in range(reference.sub_entity_count(2)):
+                entity = reference.sub_entity(2, e_n)
+                for i in product(range(1, order), repeat=2):
+                    if len(entity.vertices) == 4 or sum(i) < order:
+                        dofs.append(
+                            PointEvaluation(
+                                tuple(o + sum(a[j] * points[b]
+                                              for a, b in zip(entity.axes, i[::-1]))
+                                      for j, o in enumerate(entity.origin)),
+                                entity=(2, e_n)))
+
+            # Interior
+            for i in product(range(1, order), repeat=3):
+                if max(i[0], i[1]) + i[2] < order:
+                    dofs.append(
+                        PointEvaluation(
+                            tuple(o + sum(a[j] * points[b]
+                                          for a, b in zip(reference.axes, i))
+                                  for j, o in enumerate(reference.origin)),
+                            entity=(3, 0)))
+
+        poly = pyramid_polynomial_set(reference.tdim, 1, order)
+
+        super().__init__(reference, order, poly, dofs, reference.tdim, 1)
+
+    names = ["Lagrange", "P"]
+    references = ["pyramid"]
+    min_order = 0
+    continuity = "C0"

test/test_against_basix.py

@@ -37,8 +37,7 @@
                    ("Qcurl", "Nedelec 1st kind H(curl)", range(1, 3)),
                    ("Sdiv", "Brezzi-Douglas-Marini", range(1, 3)),
                    ("Scurl", "Nedelec 2nd kind H(curl)", range(1, 3))],
-    "prism": [("Lagrange", "Lagrange", range(1, 4))],
-    "pyramid": [("Lagrange", "Lagrange", range(1, 3))]
+    "prism": [("Lagrange", "Lagrange", range(1, 4))]
 }
 
 

test/test_against_computed_by_hand.py

@@ -1,43 +1,36 @@
 import sympy
+from utils import all_symequal
 from symfem import create_element
-
-
-def all_equal(a, b):
-    if isinstance(a, (list, tuple)):
-        for i, j in zip(a, b):
-            if not all_equal(i, j):
-                return False
-        return True
-    return a == b
+from symfem.core.symbolic import x
 
 
 def test_lagrange():
     space = create_element("triangle", "Lagrange", 1)
-    assert all_equal(
+    assert all_symequal(
         space.tabulate_basis([[0, 0], [0, 1], [1, 0]]),
         ((1, 0, 0), (0, 0, 1), (0, 1, 0)),
     )
 
 
 def test_nedelec():
     space = create_element("triangle", "Nedelec", 1)
-    assert all_equal(
+    assert all_symequal(
         space.tabulate_basis([[0, 0], [1, 0], [0, 1]], "xxyyzz"),
         ((0, 0, 1, 0, 1, 0), (0, 0, 1, 1, 0, 1), (-1, 1, 0, 0, 1, 0)),
     )
 
 
 def test_rt():
     space = create_element("triangle", "Raviart-Thomas", 1)
-    assert all_equal(
+    assert all_symequal(
         space.tabulate_basis([[0, 0], [1, 0], [0, 1]], "xxyyzz"),
         ((0, -1, 0, 0, 0, 1), (-1, 0, -1, 0, 0, 1), (0, -1, 0, -1, 1, 0)),
     )
 
 
 def test_Q():
     space = create_element("quadrilateral", "Q", 1)
-    assert all_equal(
+    assert all_symequal(
         space.tabulate_basis([[0, 0], [1, 0], [0, 1], [1, 1]]),
         ((1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1)),
     )
@@ -46,7 +39,7 @@ def test_Q():
 def test_dual0():
     space = create_element("dual polygon(4)", "dual", 0)
     q = sympy.Rational(1, 4)
-    assert all_equal(
+    assert all_symequal(
         space.tabulate_basis([[q, q], [-q, q], [-q, -q], [q, -q]]),
         ((1, ), (1, ), (1, ), (1, ))
     )
@@ -57,9 +50,29 @@ def test_dual1():
     h = sympy.Rational(1, 2)
     q = sympy.Rational(1, 4)
     e = sympy.Rational(1, 8)
-    assert all_equal(
+    assert all_symequal(
         space.tabulate_basis([[0, 0], [q, q], [h, 0]]),
         ((q, q, q, q),
          (sympy.Rational(5, 8), e, e, e),
          (sympy.Rational(3, 8), e, e, sympy.Rational(3, 8)))
     )
+
+
+def test_lagrange_pyramid():
+    space = create_element("pyramid", "Lagrange", 1)
+    x_i = x[0] / (1 - x[2])
+    y_i = x[1] / (1 - x[2])
+    z_i = x[2] / (1 - x[2])
+    basis = [(1 - x_i) * (1 - y_i) / (1 + z_i),
+             x_i * (1 - y_i) / (1 + z_i),
+             (1 - x_i) * y_i / (1 + z_i),
+             x_i * y_i / (1 + z_i),
+             z_i / (1 + z_i)]
+    assert all_symequal(basis, space.get_basis_functions())
+
+    basis = [(1 - x[0] - x[2]) * (1 - x[1] - x[2]) / (1 - x[2]),
+             x[0] * (1 - x[1] - x[2]) / (1 - x[2]),
+             (1 - x[0] - x[2]) * x[1] / (1 - x[2]),
+             x[0] * x[1] / (1 - x[2]),
+             x[2]]
+    assert all_symequal(basis, space.get_basis_functions())

test/test_elements.py

@@ -176,7 +176,7 @@ def test_element_functionals_and_continuity(
                     (0, 1, 0), (0, 0, 1))
         entity_pairs = [[0, (0, 1)], [0, (2, 3)], [0, (4, 4)],
                         [1, (1, 3)], [1, (2, 4)], [1, (6, 7)],
-                        [2, (2, 4)]]
+                        [2, (2, 3)]]
 
     if space.continuity == "L2":
         return

test/utils.py

@@ -128,7 +128,7 @@
         "Lagrange": {"equispaced": range(4)}
     },
     "pyramid": {
-        "Lagrange": {"equispaced": range(3)}
+        "Lagrange": {"equispaced": range(4)}
     }
 }