FunctionalBlock: generate nodes on multiple entities

pbrubeck · pbrubeck · commit c477bd02d5b2 · 2025-07-17T11:48:02.000-04:00
diff --git a/FIAT/brezzi_douglas_marini.py b/FIAT/brezzi_douglas_marini.py
@@ -9,8 +9,6 @@
                   polynomial_set, nedelec)
 from FIAT.check_format_variant import check_format_variant
 from FIAT.quadrature_schemes import create_quadrature
-from FIAT.quadrature import FacetQuadratureRule
-import numpy
 
 
 class BDMDualSet(dual_set.DualSet):
@@ -33,14 +31,10 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             Q_ref = create_quadrature(facet, interpolant_deg + degree)
             Pq = polynomial_set.ONPolynomialSet(facet, degree)
             Pq_at_qpts = Pq.tabulate(Q_ref.get_points())[(0,)*(sd - 1)]
+            ells = functional.FacetNormalIntegralMomentBlock(ref_el, Q_ref, Pq_at_qpts)
             for f in top[sd - 1]:
                 cur = len(nodes)
-                Q = FacetQuadratureRule(ref_el, sd - 1, f, Q_ref)
-                Jdet = Q.jacobian_determinant()
-                n = ref_el.compute_scaled_normal(f) / Jdet
-                phis = n[None, :, None] * Pq_at_qpts[:, None, :]
-                nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi)
-                             for phi in phis)
+                nodes.extend(ells.nodes(sd-1, f))
                 entity_ids[sd - 1][f] = list(range(cur, len(nodes)))
 
         elif variant == "point":
@@ -61,13 +55,12 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             cell = ref_el.construct_subelement(sd)
             Q_ref = create_quadrature(cell, interpolant_deg + degree - 1)
             Nedel = nedelec.Nedelec(cell, degree - 1, variant)
+            mapping, = set(Nedel.mapping())
             Ned_at_qpts = Nedel.tabulate(0, Q_ref.get_points())[(0,) * sd]
+            ells = functional.FacetIntegralMomentBlock(ref_el, Q_ref, Ned_at_qpts, mapping=mapping)
             for entity in top[sd]:
-                Q = FacetQuadratureRule(ref_el, sd, entity, Q_ref)
-                Jinv = numpy.linalg.inv(Q.jacobian())
-                phis = numpy.tensordot(Jinv.T, Ned_at_qpts, (1, 1)).transpose((1, 0, 2))
                 cur = len(nodes)
-                nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi) for phi in phis)
+                nodes.extend(ells.nodes(sd, entity))
                 entity_ids[sd][entity] = list(range(cur, len(nodes)))
 
         super().__init__(nodes, ref_el, entity_ids)
diff --git a/FIAT/functional.py b/FIAT/functional.py
@@ -15,7 +15,7 @@
 from itertools import chain
 import numpy
 
-from FIAT import polynomial_set, jacobi, quadrature_schemes
+from FIAT import polynomial_set, jacobi, quadrature, quadrature_schemes
 
 
 def index_iterator(shp):
@@ -25,7 +25,162 @@ def index_iterator(shp):
     return numpy.ndindex(shp)
 
 
-class Functional(object):
+def pullback(mapping, J, detJ, phi):
+    try:
+        formdegree = {
+            "identity": (0,),
+            "L2 piola": (3,),
+            "covariant piola": (1,),
+            "contravariant piola": (2,),
+            "double covariant piola": (1, 1),
+            "double contravariant piola": (2, 2),
+            "covariant contravariant piola": (1, 2),
+            "contravariant covariant piola": (2, 1), }[mapping]
+    except KeyError:
+        raise ValueError(f"Unrecognized mapping {mapping}")
+
+    F1 = numpy.linalg.pinv(J).T
+    F2 = J / detJ
+
+    perm = [*range(1, phi.ndim-1), 0, -1]
+    phi = phi.transpose(perm)
+
+    for i, k in enumerate(formdegree):
+        if k == 0:
+            continue
+        elif k == 3:
+            phi = phi / detJ
+        else:
+            F = F1 if k == 1 else F2
+            phi = numpy.tensordot(F, phi, (1, i))
+
+    perm = [-2, *reversed(range(phi.ndim-2)), -1]
+    phi = phi.transpose(perm)
+    return phi
+
+
+class FunctionalBlock:
+    def __init__(self, ref_el):
+        self.ref_el = ref_el
+
+    def completion(self):
+        return self
+
+    def nodes(self, entity_dim, entity_id):
+        raise NotImplementedError
+
+
+class FunctionalBlockView(FunctionalBlock):
+    def __init__(self, block, subset):
+        self.subset = subset
+        self.block = block
+        super().__init__(block.ref_el)
+
+    def completion(self):
+        return self.block
+
+    def nodes(self, entity_dim, entity_id):
+        nodes = list(self.block.nodes(entity_dim, entity_id))
+        for i in self.subset:
+            yield nodes[i]
+
+
+class PointFunctionalBlock(FunctionalBlock):
+    def __init__(self, ref_el, order):
+        self.order = order
+        super().__init__(ref_el)
+
+    def nodes(self, entity_dim, entity_id):
+        from FIAT.polynomial_set import mis
+        sd = self.ref_el.get_spatial_dimension()
+        for pt in self.ref_el.make_points(entity_dim, entity_id, 0):
+            for alpha in mis(sd, self.order):
+                yield PointDerivative(self.ref_el, pt, alpha)
+
+
+class PointGradient(PointFunctionalBlock):
+    def __init__(self, ref_el, x):
+        super().__init__(ref_el, 1)
+
+
+class PointHessian(PointFunctionalBlock):
+    def __init__(self, ref_el, x):
+        super().__init__(ref_el, 2)
+
+
+class PointDirectionalDerivativeBlock(PointFunctionalBlock):
+    def __init__(self, ref_el, basis):
+        self.basis = basis
+        super().__init__(ref_el)
+
+    def get_basis(self, entity_dim, entity_id):
+        return self.basis
+
+    def nodes(self, entity_dim, entity_id):
+        for pt in self.ref_el.make_points(entity_dim, entity_id, 0):
+            for e in self.get_basis(entity_dim, entity_id):
+                yield PointDirectionalDerivative(self.ref_el, pt, e)
+
+
+class PointNormalTangentialDerivativeBlock(PointDirectionalDerivativeBlock):
+    def __init__(self, ref_el, facet_id):
+        super().__init__(ref_el, None)
+
+    def get_basis(self, entity_dim, entity_id):
+        sd = self.ref_el.get_spatial_dimension()
+        basis = numpy.zeros((sd, sd))
+        basis[0] = self.ref_el.compute_scaled_normal(entity_id)
+        basis[1:] = self.ref_el.compute_tangents(sd-1, entity_id)
+        return basis
+
+
+class PointNormalDerivativeView(FunctionalBlockView):
+    def __init__(self, ref_el):
+        block = PointNormalTangentialDerivativeBlock(ref_el)
+        super().__init__(block, [0])
+
+
+class FacetIntegralMomentBlock(FunctionalBlock):
+    def __init__(self, ref_el, Q_ref, Phis, mapping="L2 piola"):
+        self.Q_ref = Q_ref
+        self.Phis = Phis
+        self.mapping = mapping
+        super().__init__(ref_el)
+
+    def get_functions(self, entity_dim, entity_id):
+        return self.Phis
+
+    def nodes(self, entity_dim, entity_id):
+        Q = quadrature.FacetQuadratureRule(self.ref_el, entity_dim, entity_id, self.Q_ref)
+        Phis = self.get_functions(entity_dim, entity_id)
+        phis = pullback(self.mapping, Q.jacobian(), Q.jacobian_determinant(), Phis)
+        for phi in phis:
+            yield FrobeniusIntegralMoment(self.ref_el, Q, phi)
+
+
+class FacetDirectionalIntegralMomentBlock(FacetIntegralMomentBlock):
+    def __init__(self, ref_el, Q_ref, Phis, direction):
+        self.direction = direction
+        super().__init__(ref_el, Q_ref, Phis)
+
+    def get_direction(self, entity_dim, entity_id):
+        return self.direction
+
+    def get_functions(self, entity_dim, entity_id):
+        udir = self.get_direction(entity_dim, entity_id)
+        Phis = self.Phis
+        return Phis[(slice(None), *(None for _ in range(udir.ndim)), slice(None))] * udir[(*(None for _ in range(Phis.ndim-1)), Ellipsis, None)]
+
+
+class FacetNormalIntegralMomentBlock(FacetDirectionalIntegralMomentBlock):
+    def __init__(self, ref_el, Q_ref, Phis):
+        super().__init__(ref_el, Q_ref, Phis, "normal")
+
+    def get_direction(self, entity_dim, entity_id):
+        return self.ref_el.compute_scaled_normal(entity_id)
+
+
+class Functional(FunctionalBlock):
     r"""Abstract class representing a linear functional.
     All FIAT functionals are discrete in the sense that
     they are written as a weighted sum of (derivatives of components of) their
diff --git a/FIAT/nedelec.py b/FIAT/nedelec.py
@@ -125,17 +125,11 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
                     Q_ref = create_quadrature(facet, interpolant_deg + phi_deg)
                     Pqmd = polynomial_set.ONPolynomialSet(facet, phi_deg, (dim,))
                     Phis = Pqmd.tabulate(Q_ref.get_points())[(0,) * dim]
-                    Phis = numpy.transpose(Phis, (0, 2, 1))
-
+                    mapping = "contravariant piola"
+                    ells = functional.FacetIntegralMomentBlock(ref_el, Q_ref, Phis, mapping=mapping)
                     for entity in top[dim]:
                         cur = len(nodes)
-                        Q = FacetQuadratureRule(ref_el, dim, entity, Q_ref)
-                        Jdet = Q.jacobian_determinant()
-                        R = numpy.array(ref_el.compute_tangents(dim, entity))
-                        phis = numpy.dot(Phis, R / Jdet)
-                        phis = numpy.transpose(phis, (0, 2, 1))
-                        nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi)
-                                     for phi in phis)
+                        nodes.extend(ells.nodes(dim, entity))
                         entity_ids[dim][entity] = list(range(cur, len(nodes)))
 
         elif variant == "point":
diff --git a/FIAT/nedelec_second_kind.py b/FIAT/nedelec_second_kind.py
@@ -5,15 +5,12 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
-import numpy
-
 from FIAT.finite_element import CiarletElement
 from FIAT.dual_set import DualSet
 from FIAT.polynomial_set import ONPolynomialSet
 from FIAT.functional import PointEdgeTangentEvaluation as Tangent
-from FIAT.functional import FrobeniusIntegralMoment as IntegralMoment
+from FIAT.functional import FacetIntegralMomentBlock
 from FIAT.raviart_thomas import RaviartThomas
-from FIAT.quadrature import FacetQuadratureRule
 from FIAT.quadrature_schemes import create_quadrature
 from FIAT.check_format_variant import check_format_variant
 
@@ -128,42 +125,30 @@ def _generate_facet_dofs(self, dim, cell, degree, offset, variant, interpolant_d
         # Construct quadrature scheme for the reference facet
         ref_facet = cell.construct_subelement(dim)
         Q_ref = create_quadrature(ref_facet, interpolant_deg + rt_degree)
+
+        # Construct Raviart-Thomas on the reference facet
         if dim == 1:
-            Phi = ONPolynomialSet(ref_facet, rt_degree, (dim,))
+            Phi = ONPolynomialSet(ref_facet, rt_degree, shape=(dim,))
+            mapping = "contravariant piola"
         else:
-            # Construct Raviart-Thomas on the reference facet
             RT = RaviartThomas(ref_facet, rt_degree, variant)
             Phi = RT.get_nodal_basis()
+            mapping, = set(RT.mapping())
 
         # Evaluate basis functions at reference quadrature points
         Phis = Phi.tabulate(Q_ref.get_points())[(0,) * dim]
-        # Note: Phis has dimensions:
-        # num_basis_functions x num_components x num_quad_points
-        Phis = numpy.transpose(Phis, (0, 2, 1))
-        # Note: Phis has dimensions:
-        # num_basis_functions x num_quad_points x num_components
+        ells = FacetIntegralMomentBlock(cell, Q_ref, Phis, mapping=mapping)
 
         # Iterate over the facets
-        cur = offset
-        for facet in range(num_facets):
-            # Get the quadrature and Jacobian on this facet
-            Q_facet = FacetQuadratureRule(cell, dim, facet, Q_ref)
-            J = Q_facet.jacobian()
-            detJ = Q_facet.jacobian_determinant()
-
-            # Map Phis -> phis (reference values to physical values)
-            piola_map = J / detJ
-            phis = numpy.dot(Phis, piola_map.T)
-            phis = numpy.transpose(phis, (0, 2, 1))
-
+        for entity in range(num_facets):
+            cur = offset + 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(IntegralMoment(cell, Q_facet, phi) for phi in phis)
+            dofs.extend(ells.nodes(dim, entity))
 
             # Assign identifiers (num RTs per face + previous edge dofs)
-            ids[facet].extend(range(cur, cur + len(phis)))
-            cur += len(phis)
+            ids[entity].extend(range(cur, offset + len(dofs)))
 
         return (dofs, ids)
 
diff --git a/FIAT/raviart_thomas.py b/FIAT/raviart_thomas.py
@@ -77,14 +77,10 @@ def __init__(self, ref_el, degree, variant, interpolant_deg):
             Q_ref = create_quadrature(facet, interpolant_deg + q)
             Pq = polynomial_set.ONPolynomialSet(facet, q if sd > 1 else 0)
             Pq_at_qpts = Pq.tabulate(Q_ref.get_points())[(0,)*(sd - 1)]
+            ells = functional.FacetNormalIntegralMomentBlock(ref_el, Q_ref, Pq_at_qpts)
             for f in top[sd - 1]:
                 cur = len(nodes)
-                Q = FacetQuadratureRule(ref_el, sd-1, f, Q_ref)
-                Jdet = Q.jacobian_determinant()
-                n = ref_el.compute_scaled_normal(f) / Jdet
-                phis = n[None, :, None] * Pq_at_qpts[:, None, :]
-                nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi)
-                             for phi in phis)
+                nodes.extend(ells.nodes(sd-1, f))
                 entity_ids[sd - 1][f] = list(range(cur, len(nodes)))
 
             # internal nodes. These are \int_T v \cdot p dx where p \in P_{q-1}^d
