Add Johnson Mercier as PhysicallyMappedElement (#126)

Co-authored-by: Rob Kirby <robert.c.kirby@gmail.com>

finat/__init__.py

@@ -10,16 +10,19 @@
 from .fiat_elements import HellanHerrmannJohnson, Regge  # noqa: F401
 from .fiat_elements import FacetBubble  # noqa: F401
 from .fiat_elements import KongMulderVeldhuizen  # noqa: F401
+
 from .argyris import Argyris            # noqa: F401
-from .hct import HsiehCloughTocher, ReducedHsiehCloughTocher   # noqa: F401
+from .aw import ArnoldWinther           # noqa: F401
+from .aw import ArnoldWintherNC         # noqa: F401
 from .bell import Bell                  # noqa: F401
+from .hct import HsiehCloughTocher, ReducedHsiehCloughTocher   # noqa: F401
 from .hermite import Hermite            # noqa: F401
+from .johnson_mercier import JohnsonMercier  # noqa: F401
 from .mtw import MardalTaiWinther       # noqa: F401
 from .morley import Morley              # noqa: F401
-from .aw import ArnoldWinther           # noqa: F401
-from .aw import ArnoldWintherNC         # noqa: F401
 from .trace import HDivTrace  # noqa: F401
 from .direct_serendipity import DirectSerendipity  # noqa: F401
+
 from .spectral import GaussLobattoLegendre, GaussLegendre, Legendre, IntegratedLegendre, FDMLagrange, FDMQuadrature, FDMDiscontinuousLagrange, FDMBrokenH1, FDMBrokenL2, FDMHermite  # noqa: F401
 from .tensorfiniteelement import TensorFiniteElement  # noqa: F401
 from .tensor_product import TensorProductElement  # noqa: F401

finat/aw.py

@@ -1,41 +1,81 @@
 """Implementation of the Arnold-Winther finite elements."""
-import numpy
 import FIAT
-from gem import Literal, ListTensor
+import numpy
+from gem import ListTensor, Literal, partial_indexed
+
 from finat.fiat_elements import FiatElement
-from finat.physically_mapped import PhysicallyMappedElement, Citations
+from finat.physically_mapped import Citations, PhysicallyMappedElement
 
 
-def _edge_transform(T, coordinate_mapping):
-    Vsub = numpy.zeros((12, 12), dtype=object)
+def _facet_transform(fiat_cell, facet_moment_degree, coordinate_mapping):
+    sd = fiat_cell.get_spatial_dimension()
+    top = fiat_cell.get_topology()
+    num_facets = len(top[sd-1])
+    dimPk_facet = FIAT.expansions.polynomial_dimension(
+        fiat_cell.construct_subelement(sd-1), facet_moment_degree)
+    dofs_per_facet = sd * dimPk_facet
+    ndofs = num_facets * dofs_per_facet
 
+    Vsub = numpy.eye(ndofs, dtype=object)
     for multiindex in numpy.ndindex(Vsub.shape):
         Vsub[multiindex] = Literal(Vsub[multiindex])
 
-    for i in range(0, 12, 2):
-        Vsub[i, i] = Literal(1)
-
-    # This bypasses the GEM wrapper.
-    that = numpy.array([T.compute_normalized_edge_tangent(i) for i in range(3)])
-    nhat = numpy.array([T.compute_normal(i) for i in range(3)])
-
-    detJ = coordinate_mapping.detJ_at([1/3, 1/3])
-    J = coordinate_mapping.jacobian_at([1/3, 1/3])
-    J_np = numpy.array([[J[0, 0], J[0, 1]],
-                        [J[1, 0], J[1, 1]]])
-    JTJ = J_np.T @ J_np
-
-    for e in range(3):
-        # Compute alpha and beta for the edge.
-        Ghat_T = numpy.array([nhat[e, :], that[e, :]])
-
-        (alpha, beta) = Ghat_T @ JTJ @ that[e, :] / detJ
-        # Stuff into the right rows and columns.
-        (idx1, idx2) = (4*e + 1, 4*e + 3)
-        Vsub[idx1, idx1-1] = Literal(-1) * alpha / beta
-        Vsub[idx1, idx1] = Literal(1) / beta
-        Vsub[idx2, idx2-1] = Literal(-1) * alpha / beta
-        Vsub[idx2, idx2] = Literal(1) / beta
+    bary = [1/(sd+1)] * sd
+    detJ = coordinate_mapping.detJ_at(bary)
+    J = coordinate_mapping.jacobian_at(bary)
+    rns = coordinate_mapping.reference_normals()
+    offset = dofs_per_facet
+    if sd == 2:
+        R = Literal(numpy.array([[0, -1], [1, 0]]))
+
+        for e in range(num_facets):
+            nhat = partial_indexed(rns, (e, ))
+            that = R @ nhat
+            Jn = J @ nhat
+            Jt = J @ that
+
+            # Compute alpha and beta for the edge.
+            alpha = (Jn @ Jt) / detJ
+            beta = (Jt @ Jt) / detJ
+            # Stuff into the right rows and columns.
+            for i in range(dimPk_facet):
+                idx = offset*e + i * dimPk_facet + 1
+                Vsub[idx, idx-1] = Literal(-1) * alpha / beta
+                Vsub[idx, idx] = Literal(1) / beta
+    elif sd == 3:
+        for f in range(num_facets):
+            nhat = fiat_cell.compute_normal(f)
+            nhat /= numpy.linalg.norm(nhat)
+            ehats = fiat_cell.compute_tangents(sd-1, f)
+            rels = [numpy.linalg.norm(ehat) for ehat in ehats]
+            thats = [a / b for a, b in zip(ehats, rels)]
+            vf = fiat_cell.volume_of_subcomplex(sd-1, f)
+
+            scale = 1.0 / numpy.dot(thats[1], numpy.cross(thats[0], nhat))
+            orth_vecs = [scale * numpy.cross(nhat, thats[1]),
+                         scale * numpy.cross(thats[0], nhat)]
+
+            Jn = J @ Literal(nhat)
+            Jts = [J @ Literal(that) for that in thats]
+            Jorth = [J @ Literal(ov) for ov in orth_vecs]
+
+            alphas = [(Jn @ Jts[i] / detJ) * (Literal(rels[i]) / Literal(2*vf)) for i in range(sd-1)]
+            betas = [Jorth[0] @ Jts[i] / detJ for i in range(sd-1)]
+            gammas = [Jorth[1] @ Jts[i] / detJ for i in range(sd-1)]
+
+            det = betas[0] * gammas[1] - betas[1] * gammas[0]
+
+            for i in range(dimPk_facet):
+                idx = offset*f + i * sd
+
+                Vsub[idx+1, idx] = (alphas[1] * gammas[0]
+                                    - alphas[0] * gammas[1]) / det
+                Vsub[idx+1, idx+1] = gammas[1] / det
+                Vsub[idx+1, idx+2] = Literal(-1) * gammas[0] / det
+                Vsub[idx+2, idx] = (alphas[0] * betas[1]
+                                    - alphas[1] * betas[0]) / det
+                Vsub[idx+2, idx+1] = Literal(-1) * betas[1] / det
+                Vsub[idx+2, idx+2] = betas[0] / det
 
     return Vsub
 
@@ -66,13 +106,11 @@ def __init__(self, cell, degree):
     def basis_transformation(self, coordinate_mapping):
         """Note, the extra 3 dofs which are removed here
         correspond to the constraints."""
-
-        T = self.cell
         V = numpy.zeros((18, 15), dtype=object)
         for multiindex in numpy.ndindex(V.shape):
             V[multiindex] = Literal(V[multiindex])
 
-        V[:12, :12] = _edge_transform(T, coordinate_mapping)
+        V[:12, :12] = _facet_transform(self.cell, 1, coordinate_mapping)
 
         # internal dofs
         W = _evaluation_transform(coordinate_mapping)
@@ -126,7 +164,7 @@ def basis_transformation(self, coordinate_mapping):
         # Put into the right rows and columns.
         V[0:3, 0:3] = V[3:6, 3:6] = V[6:9, 6:9] = W
 
-        V[9:21, 9:21] = _edge_transform(self.cell, coordinate_mapping)
+        V[9:21, 9:21] = _facet_transform(self.cell, 1, coordinate_mapping)
 
         # internal DOFs
         detJ = coordinate_mapping.detJ_at([1/3, 1/3])

finat/hct.py

@@ -78,6 +78,13 @@ def __init__(self, cell, degree):
             Citations().register("Clough1965")
         super().__init__(FIAT.HsiehCloughTocher(cell, reduced=True))
 
+        reduced_dofs = deepcopy(self._element.entity_dofs())
+        sd = cell.get_spatial_dimension()
+        fdim = sd - 1
+        for entity in reduced_dofs[fdim]:
+            reduced_dofs[fdim][entity] = []
+        self._entity_dofs = reduced_dofs
+
     def basis_transformation(self, coordinate_mapping):
         # Jacobians at cell center
         J = coordinate_mapping.jacobian_at([1/3, 1/3])
@@ -133,11 +140,7 @@ def basis_transformation(self, coordinate_mapping):
         return ListTensor(V.T)
 
     def entity_dofs(self):
-        edofs = deepcopy(super(ReducedHsiehCloughTocher, self).entity_dofs())
-        dim = 1
-        for entity in edofs[dim]:
-            edofs[dim][entity] = []
-        return edofs
+        return self._entity_dofs
 
     @property
     def index_shape(self):

finat/johnson_mercier.py

@@ -0,0 +1,34 @@
+import FIAT
+import numpy
+from gem import ListTensor, Literal
+
+from finat.aw import _facet_transform
+from finat.fiat_elements import FiatElement
+from finat.physically_mapped import Citations, PhysicallyMappedElement
+
+
+class JohnsonMercier(PhysicallyMappedElement, FiatElement):  # symmetric matrix valued
+    def __init__(self, cell, degree, variant=None):
+        if degree != 1:
+            raise ValueError("Degree must be 1 for Johnson-Mercier element")
+        if Citations is not None:
+            Citations().register("Gopalakrishnan2024")
+        self._indices = slice(None, None)
+        super(JohnsonMercier, self).__init__(FIAT.JohnsonMercier(cell, degree, variant=variant))
+
+    def basis_transformation(self, coordinate_mapping):
+        numbf = self._element.space_dimension()
+        ndof = self.space_dimension()
+        V = numpy.eye(numbf, ndof, dtype=object)
+        for multiindex in numpy.ndindex(V.shape):
+            V[multiindex] = Literal(V[multiindex])
+
+        Vsub = _facet_transform(self.cell, 1, coordinate_mapping)
+        Vsub = Vsub[:, self._indices]
+        m, n = Vsub.shape
+        V[:m, :n] = Vsub
+
+        # Note: that the edge DOFs are scaled by edge lengths in FIAT implies
+        # that they are already have the necessary rescaling to improve
+        # conditioning.
+        return ListTensor(V.T)

finat/physically_mapped.py

@@ -136,6 +136,14 @@
   year={2017},
   institution={Tech. Rep. ICES REPORT 17-28, Institute for Computational Engineering and Sciences}
 }
+""")
+    Citations().add("Gopalakrishnan2024", """
+@article{gopalakrishnan2024johnson,
+  title={{The Johnson-Mercier elasticity element in any dimensions}},
+  author={Gopalakrishnan, J and Guzman, J and Lee, J J},
+  journal={arXiv preprint arXiv:2403.13189},
+  year={2024}
+}
 """)
 except ImportError:
     Citations = None

finat/ufl/elementlist.py

@@ -142,6 +142,8 @@ def show_elements():
 register_element("Regge", "Regge", 2, HEin, "double covariant Piola",
                  (0, None), simplices[1:])
 register_element("HDiv Trace", "HDivT", 0, L2, "identity", (0, None), any_cell)
+register_element("Johnson-Mercier", "JM", 2, HDivDiv,
+                 "double contravariant Piola", (1, 1), simplices[1:])
 register_element("Hellan-Herrmann-Johnson", "HHJ", 2, HDivDiv,
                  "double contravariant Piola", (0, None), ("triangle",))
 register_element("Nonconforming Arnold-Winther", "AWnc", 2, HDivDiv,

test/test_hct.py

@@ -1,49 +1,53 @@
 import FIAT
 import finat
 import numpy as np
+import pytest
 from gem.interpreter import evaluate
 
 from fiat_mapping import MyMapping
 
 
-def test_hct():
-    degree = 3
-    ref_cell = FIAT.ufc_simplex(2)
-    ref_pts = finat.point_set.PointSet(FIAT.reference_element.make_lattice(ref_cell.vertices, degree))
-    ref_element = finat.HsiehCloughTocher(ref_cell, degree, avg=True)
+@pytest.fixture
+def ref_cell(request):
+    K = FIAT.ufc_simplex(2)
+    return K
 
-    phys_cell = FIAT.ufc_simplex(2)
-    phys_cell.vertices = ((0.0, 0.1), (1.17, -0.09), (0.15, 1.84))
 
-    mapping = MyMapping(ref_cell, phys_cell)
-    z = (0, 0)
-    finat_vals_gem = ref_element.basis_evaluation(0, ref_pts, coordinate_mapping=mapping)[z]
-    finat_vals = evaluate([finat_vals_gem])[0].arr
+@pytest.fixture
+def phys_cell(request):
+    K = FIAT.ufc_simplex(2)
+    K.vertices = ((0.0, 0.1), (1.17, -0.09), (0.15, 1.84))
+    return K
 
-    phys_cell_FIAT = FIAT.HsiehCloughTocher(phys_cell, degree)
-    phys_points = FIAT.reference_element.make_lattice(phys_cell.vertices, degree)
-    phys_vals = phys_cell_FIAT.tabulate(0, phys_points)[z]
 
-    assert np.allclose(finat_vals.T, phys_vals)
+def make_unisolvent_points(element):
+    degree = element.degree()
+    ref_complex = element.get_reference_complex()
+    sd = ref_complex.get_spatial_dimension()
+    top = ref_complex.get_topology()
+    pts = []
+    for cell in top[sd]:
+        pts.extend(ref_complex.make_points(sd, cell, degree+sd+1))
+    return pts
 
 
-def test_reduced_hct():
+@pytest.mark.parametrize("reduced", (False, True), ids=("standard", "reduced"))
+def test_hct(ref_cell, phys_cell, reduced):
     degree = 3
-    ref_cell = FIAT.ufc_simplex(2)
-    ref_pts = finat.point_set.PointSet(FIAT.reference_element.make_lattice(ref_cell.vertices, degree))
-    ref_element = finat.ReducedHsiehCloughTocher(ref_cell, degree)
-
-    phys_cell = FIAT.ufc_simplex(2)
-    phys_cell.vertices = ((0.0, 0.1), (1.17, -0.09), (0.15, 1.84))
+    if reduced:
+        ref_element = finat.ReducedHsiehCloughTocher(ref_cell, degree)
+    else:
+        ref_element = finat.HsiehCloughTocher(ref_cell, degree, avg=True)
+    ref_pts = finat.point_set.PointSet(make_unisolvent_points(ref_element._element))
 
     mapping = MyMapping(ref_cell, phys_cell)
-    z = (0, 0)
+    z = (0,) * ref_element.cell.get_spatial_dimension()
     finat_vals_gem = ref_element.basis_evaluation(0, ref_pts, coordinate_mapping=mapping)[z]
     finat_vals = evaluate([finat_vals_gem])[0].arr
 
-    phys_cell_FIAT = FIAT.HsiehCloughTocher(phys_cell, degree, reduced=True)
-    phys_points = FIAT.reference_element.make_lattice(phys_cell.vertices, degree)
-    phys_vals = phys_cell_FIAT.tabulate(0, phys_points)[z]
+    phys_element = FIAT.HsiehCloughTocher(phys_cell, degree, reduced=reduced)
+    phys_points = make_unisolvent_points(phys_element)
+    phys_vals = phys_element.tabulate(0, phys_points)[z]
 
     numbf = ref_element.space_dimension()
     assert np.allclose(finat_vals.T, phys_vals[:numbf])