Merge pull request #50 from firedrakeproject/pbrubeck/fix/enriched-mi…

…xed-degree

Ensure EnrichedElement uses the embedding ExpansionSet

FIAT/barycentric_interpolation.py

@@ -29,19 +29,23 @@ class LagrangeLineExpansionSet(expansions.LineExpansionSet):
     https://doi.org/10.1137/S0036144502417715 Eq. (4.2) & (9.4)
     """
     def __init__(self, ref_el, pts):
-        self.nodes = numpy.array(pts).flatten()
-        self.dmat, self.weights = make_dmat(self.nodes)
+        self.points = pts
+        self.x = numpy.array(pts).flatten()
+        self.dmat, self.weights = make_dmat(self.x)
         super(LagrangeLineExpansionSet, self).__init__(ref_el)
 
     def get_num_members(self, n):
-        return len(self.nodes)
+        return len(self.points)
+
+    def get_points(self):
+        return self.points
 
     def get_dmats(self, degree):
         return [self.dmat.T]
 
     def tabulate(self, n, pts):
-        assert n == len(self.nodes)-1
-        results = numpy.add.outer(-self.nodes, numpy.array(pts).flatten())
+        assert n == len(self.points)-1
+        results = numpy.add.outer(-self.x, numpy.array(pts).flatten())
         with numpy.errstate(divide='ignore', invalid='ignore'):
             numpy.reciprocal(results, out=results)
             numpy.multiply(results, self.weights[:, None], out=results)
@@ -85,14 +89,14 @@ def __init__(self, ref_el, pts, shape=tuple()):
         if shape == tuple():
             coeffs = numpy.eye(num_members)
         else:
-            coeffs_shape = tuple([num_members] + list(shape) + [num_exp_functions])
+            coeffs_shape = (num_members, *shape, num_exp_functions)
             coeffs = numpy.zeros(coeffs_shape, "d")
             # use functional's index_iterator function
             cur_bf = 0
             for idx in index_iterator(shape):
                 n = expansions.polynomial_dimension(ref_el, embedded_degree)
                 for exp_bf in range(n):
-                    cur_idx = tuple([cur_bf] + list(idx) + [exp_bf])
+                    cur_idx = (cur_bf, *idx, exp_bf)
                     coeffs[cur_idx] = 1.0
                     cur_bf += 1
 

FIAT/expansions.py

@@ -9,7 +9,6 @@
 
 import numpy
 import math
-from itertools import product
 from FIAT import reference_element
 from FIAT import jacobi
 
@@ -209,19 +208,16 @@ def get_dmats(self, degree):
 
     def _tabulate_jet(self, degree, pts, order=0):
         from FIAT.polynomial_set import mis
-        result = {}
         D = self.ref_el.get_spatial_dimension()
         lorder = min(2, order)
         vals = self._tabulate(degree, numpy.transpose(pts), order=lorder)
-        base_vals = numpy.array(vals[0])
-        base_alpha = (0,) * D
-        result[base_alpha] = base_vals
+        result = {(0,) * D: numpy.array(vals[0])}
         for r in range(1, 1+lorder):
             vr = numpy.transpose(vals[r], tuple(range(1, r+1)) + (0, r+1))
-            for indices in product(range(D), repeat=r):
+            for indices in numpy.ndindex(vr.shape[:r]):
                 alpha = tuple(map(indices.count, range(D)))
                 if alpha not in result:
-                    result[alpha] = vr[indices].reshape(base_vals.shape)
+                    result[alpha] = vr[indices]
 
         def distance(alpha, beta):
             return sum(ai != bi for ai, bi in zip(alpha, beta))
@@ -261,10 +257,10 @@ def tabulate_jet(self, n, pts, order=1):
         v0 = vals[(0,)*D]
         data = [v0]
         for r in range(1, order+1):
-            v = numpy.zeros((D,)*r + v0.shape, dtype=v0.dtype)
-            for index in product(range(D), repeat=r):
-                v[index] = vals[tuple(map(index.count, range(D)))]
-            data.append(v.transpose((r, r+1) + tuple(range(r))))
+            vr = numpy.zeros((D,)*r + v0.shape, dtype=v0.dtype)
+            for index in numpy.ndindex(vr.shape[:r]):
+                vr[index] = vals[tuple(map(index.count, range(D)))]
+            data.append(vr.transpose((r, r+1) + tuple(range(r))))
         return data
 
 

FIAT/lagrange.py

@@ -48,7 +48,7 @@ class Lagrange(finite_element.CiarletElement):
 
     def __init__(self, ref_el, degree, variant="equispaced"):
         dual = LagrangeDualSet(ref_el, degree, variant=variant)
-        if ref_el.shape == LINE and variant != "equispaced":
+        if ref_el.shape == LINE:
             # In 1D we can use the primal basis as the expansion set,
             # avoiding any round-off coming from a basis transformation
             points = []

FIAT/nodal_enriched.py

@@ -9,6 +9,7 @@
 from FIAT.polynomial_set import PolynomialSet
 from FIAT.dual_set import DualSet
 from FIAT.finite_element import CiarletElement
+from FIAT.barycentric_interpolation import LagrangeLineExpansionSet
 
 __all__ = ['NodalEnrichedElement']
 
@@ -35,28 +36,38 @@ def __init__(self, *elements):
                              "of NodalEnrichedElement are nodal")
 
         # Extract common data
-        ref_el = elements[0].get_reference_element()
-        expansion_set = elements[0].get_nodal_basis().get_expansion_set()
-        degree = min(e.get_nodal_basis().get_degree() for e in elements)
+        degree = max(e.get_nodal_basis().get_degree() for e in elements)
         embedded_degree = max(e.get_nodal_basis().get_embedded_degree()
                               for e in elements)
         order = max(e.get_order() for e in elements)
-        mapping = elements[0].mapping()[0]
         formdegree = None if any(e.get_formdegree() is None for e in elements) \
             else max(e.get_formdegree() for e in elements)
-        value_shape = elements[0].value_shape()
+        # LagrangeExpansionSet has fixed degree, ensure we grab the embedding one
+        elem = next(e for e in elements
+                    if e.get_nodal_basis().get_embedded_degree() == embedded_degree)
+        ref_el = elem.get_reference_element()
+        expansion_set = elem.get_nodal_basis().get_expansion_set()
+        mapping = elem.mapping()[0]
+        value_shape = elem.value_shape()
 
         # Sanity check
         assert all(e.get_nodal_basis().get_reference_element() ==
                    ref_el for e in elements)
-        assert all(type(e.get_nodal_basis().get_expansion_set()) ==
-                   type(expansion_set) for e in elements)
         assert all(e_mapping == mapping for e in elements
                    for e_mapping in e.mapping())
         assert all(e.value_shape() == value_shape for e in elements)
 
         # Merge polynomial sets
-        coeffs = _merge_coeffs([e.get_coeffs() for e in elements])
+        if isinstance(expansion_set, LagrangeLineExpansionSet):
+            # Obtain coefficients via interpolation
+            points = expansion_set.get_points()
+            coeffs = [e.tabulate(0, points)[(0,)] for e in elements]
+        else:
+            assert all(type(e.get_nodal_basis().get_expansion_set()) ==
+                       type(expansion_set) for e in elements)
+            coeffs = [e.get_coeffs() for e in elements]
+
+        coeffs = _merge_coeffs(coeffs)
         poly_set = PolynomialSet(ref_el,
                                  degree,
                                  embedded_degree,

FIAT/polynomial_set.py

@@ -95,7 +95,7 @@ def get_embedded_degree(self):
 
     def get_dmats(self):
         if len(self.dmats) == 0:
-            self.dmats = self.expansion_set.get_dmats(self.degree)
+            self.dmats = self.expansion_set.get_dmats(self.embedded_degree)
         return self.dmats
 
     def get_reference_element(self):
@@ -132,13 +132,12 @@ def __init__(self, ref_el, degree, shape=tuple()):
         embedded_degree = degree
         expansion_set = expansions.ExpansionSet(ref_el)
 
+        # set up coefficients
         if shape == tuple():
             coeffs = numpy.eye(num_members)
         else:
-            # set up coefficients
             coeffs_shape = (num_members, *shape, num_exp_functions)
             coeffs = numpy.zeros(coeffs_shape, "d")
-
             # use functional's index_iterator function
             cur_bf = 0
             for idx in index_iterator(shape):
@@ -148,8 +147,8 @@ def __init__(self, ref_el, degree, shape=tuple()):
                     coeffs[cur_idx] = 1.0
                     cur_bf += 1
 
-        PolynomialSet.__init__(self, ref_el, degree, embedded_degree,
-                               expansion_set, coeffs)
+        super(ONPolynomialSet, self).__init__(ref_el, degree, embedded_degree,
+                                              expansion_set, coeffs)
 
 
 def project(f, U, Q):
@@ -240,5 +239,5 @@ def __init__(self, ref_el, degree, size=None):
                     coeffs[cur_bf, j, i, exp_bf] = 1.0
                     cur_bf += 1
 
-        PolynomialSet.__init__(self, ref_el, degree, embedded_degree,
-                               expansion_set, coeffs)
+        super(ONSymTensorPolynomialSet, self).__init__(ref_el, degree, embedded_degree,
+                                                       expansion_set, coeffs)

test/unit/test_fiat.py

@@ -379,17 +379,21 @@ def test_empty_bubble():
         Bubble(S, 3)
 
 
-def test_nodal_enriched_implementation():
+@pytest.mark.parametrize('elements', [
+    (Lagrange(I, 2), Lagrange(I, 1), Bubble(I, 2)),
+    (GaussLobattoLegendre(I, 3), Lagrange(I, 1),
+     RestrictedElement(GaussLobattoLegendre(I, 3), restriction_domain="interior")),
+    (RaviartThomas(T, 2),
+     RestrictedElement(RaviartThomas(T, 2), restriction_domain='facet'),
+     RestrictedElement(RaviartThomas(T, 2), restriction_domain='interior')),
+])
+def test_nodal_enriched_implementation(elements):
     """Following element pair should be the same.
     This might be fragile to dof reordering but works now.
     """
 
-    e0 = RaviartThomas(T, 2)
-
-    e1 = NodalEnrichedElement(
-        RestrictedElement(RaviartThomas(T, 2), restriction_domain='facet'),
-        RestrictedElement(RaviartThomas(T, 2), restriction_domain='interior')
-    )
+    e0 = elements[0]
+    e1 = NodalEnrichedElement(*elements[1:])
 
     for attr in ["degree",
                  "get_reference_element",