FunctionalGenerator classes to reuse tabulations on facets

FIAT/brezzi_douglas_marini.py

@@ -25,9 +25,10 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             # degree is q
             Q = create_quadrature(facet, interpolant_deg + degree)
             Pq = polynomial_set.ONPolynomialSet(facet, degree)
+            ells = functional.NormalMoments(ref_el, Q, Pq)
             for f in top[sd - 1]:
                 cur = len(nodes)
-                nodes.extend(functional.NormalMoments(ref_el, Q, Pq, f))
+                nodes.extend(ells.generate(sd-1, f))
                 entity_ids[sd - 1][f].extend(range(cur, len(nodes)))
 
         elif variant == "point":
@@ -48,7 +49,8 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             Q = create_quadrature(ref_el, interpolant_deg + degree - 1)
             Nedel = nedelec.Nedelec(ref_el, degree - 1, variant)
             Nedfs = Nedel.get_nodal_basis()
-            nodes.extend(functional.FrobeniusIntegralMoments(ref_el, Q, Nedfs))
+            ells = functional.FrobeniusIntegralMoments(ref_el, Q, Nedfs)
+            nodes.extend(ells.generate(sd, 0))
             entity_ids[sd][0].extend(range(cur, len(nodes)))
 
         super().__init__(nodes, ref_el, entity_ids)

FIAT/functional.py

@@ -756,38 +756,62 @@ def __init__(self, ref_el, Q, P_at_qpts, facet):
         super().__init__(ref_el, shp, pt_dict, {}, "IntegralMomentOfNormalNormalEvaluation")
 
 
-def IntegralMoments(ref_el, Q, P, **kwargs):
-    sd = ref_el.get_spatial_dimension()
-    phis = P.tabulate(Q.get_points())[(0,)*sd]
-    for phi in phis:
-        yield IntegralMoment(ref_el, Q, phi, **kwargs)
+class FunctionalGenerator():
+    def __init__(self, ref_el):
+        self.ref_el = ref_el
+
+    def generate(self, dim, entity):
+        raise NotImplementedError
+
+
+class IntegralMoments(FunctionalGenerator):
+    def __init__(self, ref_el, Q, P, comp=tuple(), shp=None):
+        if shp is None:
+            shp = P.get_shape()
+        sd = P.ref_el.get_spatial_dimension()
+        self.phis = P.tabulate(Q.get_points())[(0,)*sd]
+        self.Q = Q
+        self.comp = comp
+        self.shape = shp
+        super().__init__(ref_el)
+
+    def generate(self, dim, entity):
+        if dim == self.ref_el.get_spatial_dimension():
+            assert entity == 0
+            for phi in self.phis:
+                yield IntegralMoment(self.ref_el, self.Q, phi, comp=self.comp, shp=self.shape)
+
+
+class FrobeniusIntegralMoments(IntegralMoments):
+    def __init__(self, ref_el, Q, P):
+        super().__init__(ref_el, Q, P, comp=slice(None))
 
+    def generate(self, dim, entity):
+        if dim == self.ref_el.get_spatial_dimension():
+            assert entity == 0
+            for phi in self.phis:
+                yield FrobeniusIntegralMoment(self.ref_el, self.Q, phi)
 
-def FrobeniusIntegralMoments(ref_el, Q, P, **kwargs):
-    sd = ref_el.get_spatial_dimension()
-    phis = P.tabulate(Q.get_points())[(0,)*sd]
-    for phi in phis:
-        yield FrobeniusIntegralMoment(ref_el, Q, phi, **kwargs)
 
+class FacetIntegralMoments(FrobeniusIntegralMoments):
+    def get_entity_transform(self, dim, entity):
+        return lambda f: f
 
-def FacetIntegralMoments(ref_el, Q, P, dim, entity_id, transform=None):
-    if transform is None:
-        transform = lambda f: f
-    phis = P.tabulate(Q.get_points())[(0,)*dim]
-    Q_mapped = quadrature.FacetQuadratureRule(ref_el, dim, entity_id, Q)
-    phis /= Q_mapped.jacobian_determinant()
-    for phi in phis:
-        yield FrobeniusIntegralMoment(ref_el, Q_mapped, transform(phi))
+    def generate(self, dim, entity):
+        transform = self.get_entity_transform(dim, entity)
+        Q_mapped = quadrature.FacetQuadratureRule(self.ref_el, dim, entity, self.Q)
+        Jdet = Q_mapped.jacobian_determinant()
+        for phi in self.phis:
+            yield FrobeniusIntegralMoment(self.ref_el, Q_mapped, transform(phi)/Jdet)
 
 
-def NormalMoments(ref_el, Q, P, face_id):
-    dim = ref_el.get_spatial_dimension() - 1
-    n = ref_el.compute_scaled_normal(face_id)
-    transform = lambda f: numpy.multiply(n[..., None], f)
-    yield from FacetIntegralMoments(ref_el, Q, P, dim, face_id, transform)
+class NormalMoments(FacetIntegralMoments):
+    def get_entity_transform(self, dim, entity):
+        n = self.ref_el.compute_scaled_normal(entity)
+        return lambda f: numpy.multiply(n[..., None], f)
 
 
-def TangentialMoments(ref_el, Q, P, dim, entity_id):
-    ts = ref_el.compute_tangents(dim, entity_id)
-    transform = lambda f: numpy.dot(ts.T, f)
-    yield from FacetIntegralMoments(ref_el, Q, P, dim, entity_id, transform)
+class TangentialMoments(FacetIntegralMoments):
+    def get_entity_transform(self, dim, entity):
+        ts = self.ref_el.compute_tangents(dim, entity)
+        return lambda f: numpy.dot(ts.T, f)

FIAT/nedelec.py

@@ -123,9 +123,10 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
                     facet = ref_el.construct_subelement(dim)
                     Q = create_quadrature(facet, interpolant_deg + phi_deg)
                     Pqmd = polynomial_set.ONPolynomialSet(facet, phi_deg, (dim,))
+                    ells = functional.TangentialMoments(ref_el, Q, Pqmd)
                     for entity in top[dim]:
                         cur = len(nodes)
-                        nodes.extend(functional.TangentialMoments(ref_el, Q, Pqmd, dim, entity))
+                        nodes.extend(ells.generate(dim, entity))
                         entity_ids[dim][entity].extend(range(cur, len(nodes)))
 
         elif variant == "point":
@@ -155,7 +156,8 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             cur = len(nodes)
             Q = create_quadrature(ref_el, interpolant_deg + phi_deg)
             Pqmd = polynomial_set.ONPolynomialSet(ref_el, phi_deg, shape=(sd,))
-            nodes.extend(functional.FrobeniusIntegralMoments(ref_el, Q, Pqmd))
+            ells = functional.FrobeniusIntegralMoments(ref_el, Q, Pqmd)
+            nodes.extend(ells.generate(sd, 0))
             entity_ids[sd][0].extend(range(cur, len(nodes)))
 
         super().__init__(nodes, ref_el, entity_ids)

FIAT/nedelec_second_kind.py

@@ -132,12 +132,13 @@ def _generate_facet_dofs(self, dim, cell, degree, offset, variant, interpolant_d
             RT = RaviartThomas(facet, rt_degree, variant)
             P = RT.get_nodal_basis()
 
+        ells = TangentialMoments(cell, Q, P)
         for entity in sorted(top[dim]):
             cur = len(dofs)
             # Construct degrees of freedom as integral moments on this cell,
             # using the face quadrature weighted against the values
             # of the (physical) Raviart--Thomas'es on the face
-            dofs.extend(TangentialMoments(cell, Q, P, dim, entity))
+            dofs.extend(ells.generate(dim, entity))
 
             # Assign identifiers (num RTs per face + previous edge dofs)
             ids[entity].extend(range(cur, len(dofs)))

FIAT/raviart_thomas.py

@@ -70,17 +70,19 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             q = degree - 1
             Q = create_quadrature(facet, interpolant_deg + q)
             Pq = polynomial_set.ONPolynomialSet(facet, q if sd > 1 else 0)
+            ells = functional.NormalMoments(ref_el, Q, Pq)
             for f in top[sd - 1]:
                 cur = len(nodes)
-                nodes.extend(functional.NormalMoments(ref_el, Q, Pq, f))
+                nodes.extend(ells.generate(sd-1, f))
                 entity_ids[sd - 1][f].extend(range(cur, len(nodes)))
 
             # internal nodes. These are \int_T v \cdot p dx where p \in P_{q-1}^d
             if q > 0:
                 cur = len(nodes)
                 Q = create_quadrature(ref_el, interpolant_deg + q - 1)
                 Pqm1 = polynomial_set.ONPolynomialSet(ref_el, q - 1, shape=(sd,))
-                nodes.extend(functional.FrobeniusIntegralMoments(ref_el, Q, Pqm1))
+                ells = functional.FrobeniusIntegralMoments(ref_el, Q, Pqm1)
+                nodes.extend(ells.generate(sd, 0))
                 entity_ids[sd][0].extend(range(cur, len(nodes)))
 
         elif variant == "point":