Add many geometrical quantities to ffcx expression generator.

ffcx/codegeneration/access.py

@@ -228,7 +228,7 @@ def reference_normal(self, mt, tabledata, access):
         """Access a reference normal."""
         cellname = ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
         if cellname in ("interval", "triangle", "tetrahedron", "quadrilateral", "hexahedron"):
-            table = L.Symbol(f"{cellname}_reference_facet_normals", dtype=L.DataType.REAL)
+            table = L.Symbol(f"{cellname}_reference_normals", dtype=L.DataType.REAL)
             facet = self.symbols.entity("facet", mt.restriction)
             return table[facet][mt.component[0]]
         else:
@@ -238,7 +238,7 @@ def cell_facet_jacobian(self, mt, tabledata, num_points):
         """Access a cell facet jacobian."""
         cellname = ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
         if cellname in ("triangle", "tetrahedron", "quadrilateral", "hexahedron"):
-            table = L.Symbol(f"{cellname}_reference_facet_jacobian", dtype=L.DataType.REAL)
+            table = L.Symbol(f"{cellname}_cell_facet_jacobian", dtype=L.DataType.REAL)
             facet = self.symbols.entity("facet", mt.restriction)
             return table[facet][mt.component[0]][mt.component[1]]
         elif cellname == "interval":
@@ -250,7 +250,7 @@ def reference_cell_edge_vectors(self, mt, tabledata, num_points):
         """Access a reference cell edge vector."""
         cellname = ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
         if cellname in ("triangle", "tetrahedron", "quadrilateral", "hexahedron"):
-            table = L.Symbol(f"{cellname}_reference_edge_vectors", dtype=L.DataType.REAL)
+            table = L.Symbol(f"{cellname}_reference_cell_edge_vectors", dtype=L.DataType.REAL)
             return table[mt.component[0]][mt.component[1]]
         elif cellname == "interval":
             raise RuntimeError(
@@ -263,9 +263,8 @@ def reference_facet_edge_vectors(self, mt, tabledata, num_points):
         """Access a reference facet edge vector."""
         cellname = ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
         if cellname in ("tetrahedron", "hexahedron"):
-            table = L.Symbol(f"{cellname}_reference_edge_vectors", dtype=L.DataType.REAL)
-            facet = self.symbols.entity("facet", mt.restriction)
-            return table[facet][mt.component[0]][mt.component[1]]
+            table = L.Symbol(f"{cellname}_reference_facet_edge_vectors", dtype=L.DataType.REAL)
+            return table[mt.component[0]][mt.component[1]]
         elif cellname in ("interval", "triangle", "quadrilateral"):
             raise RuntimeError(
                 "The reference cell facet edge vectors doesn't make sense for interval "
@@ -280,7 +279,7 @@ def facet_orientation(self, mt, tabledata, num_points):
         if cellname not in ("interval", "triangle", "tetrahedron"):
             raise RuntimeError(f"Unhandled cell types {cellname}.")
 
-        table = L.Symbol(f"{cellname}_facet_orientations", dtype=L.DataType.INT)
+        table = L.Symbol(f"{cellname}_facet_orientation", dtype=L.DataType.INT)
         facet = self.symbols.entity("facet", mt.restriction)
         return table[facet]
 

ffcx/codegeneration/expression_generator.py

@@ -61,10 +61,14 @@ def generate(self):
 
     def generate_geometry_tables(self):
         """Generate static tables of geometry data."""
-        # Currently we only support circumradius
         ufl_geometry = {
+            ufl.geometry.FacetEdgeVectors: "facet_edge_vectors",
+            ufl.geometry.CellFacetJacobian: "cell_facet_jacobian",
             ufl.geometry.ReferenceCellVolume: "reference_cell_volume",
-            ufl.geometry.ReferenceNormal: "reference_facet_normals",
+            ufl.geometry.ReferenceFacetVolume: "reference_facet_volume",
+            ufl.geometry.ReferenceCellEdgeVectors: "reference_cell_edge_vectors",
+            ufl.geometry.ReferenceFacetEdgeVectors: "reference_facet_edge_vectors",
+            ufl.geometry.ReferenceNormal: "reference_normals",
         }
 
         cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()}  # type: ignore

ffcx/codegeneration/geometry.py

@@ -15,18 +15,18 @@ def write_table(tablename, cellname):
     """Write a table."""
     if tablename == "facet_edge_vertices":
         return facet_edge_vertices(tablename, cellname)
-    if tablename == "reference_facet_jacobian":
-        return reference_facet_jacobian(tablename, cellname)
+    if tablename == "cell_facet_jacobian":
+        return cell_facet_jacobian(tablename, cellname)
     if tablename == "reference_cell_volume":
         return reference_cell_volume(tablename, cellname)
     if tablename == "reference_facet_volume":
         return reference_facet_volume(tablename, cellname)
-    if tablename == "reference_edge_vectors":
-        return reference_edge_vectors(tablename, cellname)
-    if tablename == "facet_reference_edge_vectors":
-        return facet_reference_edge_vectors(tablename, cellname)
-    if tablename == "reference_facet_normals":
-        return reference_facet_normals(tablename, cellname)
+    if tablename == "reference_cell_edge_vectors":
+        return reference_cell_edge_vectors(tablename, cellname)
+    if tablename == "reference_facet_edge_vectors":
+        return reference_facet_edge_vectors(tablename, cellname)
+    if tablename == "reference_normals":
+        return reference_normals(tablename, cellname)
     if tablename == "facet_orientation":
         return facet_orientation(tablename, cellname)
     raise ValueError(f"Unknown geometry table name: {tablename}")
@@ -56,7 +56,7 @@ def facet_edge_vertices(tablename, cellname):
     return L.ArrayDecl(symbol, values=out, const=True)
 
 
-def reference_facet_jacobian(tablename, cellname):
+def cell_facet_jacobian(tablename, cellname):
     """Write a reference facet jacobian."""
     celltype = getattr(basix.CellType, cellname)
     out = basix.cell.facet_jacobians(celltype)
@@ -83,7 +83,7 @@ def reference_facet_volume(tablename, cellname):
     return L.VariableDecl(symbol, volumes[0])
 
 
-def reference_edge_vectors(tablename, cellname):
+def reference_cell_edge_vectors(tablename, cellname):
     """Write reference edge vectors."""
     celltype = getattr(basix.CellType, cellname)
     topology = basix.topology(celltype)
@@ -94,7 +94,7 @@ def reference_edge_vectors(tablename, cellname):
     return L.ArrayDecl(symbol, values=out, const=True)
 
 
-def facet_reference_edge_vectors(tablename, cellname):
+def reference_facet_edge_vectors(tablename, cellname):
     """Write facet reference edge vectors."""
     celltype = getattr(basix.CellType, cellname)
     topology = basix.topology(celltype)
@@ -121,7 +121,7 @@ def facet_reference_edge_vectors(tablename, cellname):
     return L.ArrayDecl(symbol, values=out, const=True)
 
 
-def reference_facet_normals(tablename, cellname):
+def reference_normals(tablename, cellname):
     """Write reference facet normals."""
     celltype = getattr(basix.CellType, cellname)
     out = basix.cell.facet_outward_normals(celltype)
@@ -134,4 +134,4 @@ def facet_orientation(tablename, cellname):
     celltype = getattr(basix.CellType, cellname)
     out = basix.cell.facet_orientations(celltype)
     symbol = L.Symbol(f"{cellname}_{tablename}", dtype=L.DataType.REAL)
-    return L.ArrayDecl(symbol, values=out, const=True)
+    return L.ArrayDecl(symbol, values=np.asarray(out), const=True)

ffcx/codegeneration/integral_generator.py

@@ -197,12 +197,12 @@ def generate_geometry_tables(self):
         """Generate static tables of geometry data."""
         ufl_geometry = {
             ufl.geometry.FacetEdgeVectors: "facet_edge_vertices",
-            ufl.geometry.CellFacetJacobian: "reference_facet_jacobian",
+            ufl.geometry.CellFacetJacobian: "cell_facet_jacobian",
             ufl.geometry.ReferenceCellVolume: "reference_cell_volume",
             ufl.geometry.ReferenceFacetVolume: "reference_facet_volume",
-            ufl.geometry.ReferenceCellEdgeVectors: "reference_edge_vectors",
-            ufl.geometry.ReferenceFacetEdgeVectors: "facet_reference_edge_vectors",
-            ufl.geometry.ReferenceNormal: "reference_facet_normals",
+            ufl.geometry.ReferenceCellEdgeVectors: "reference_cell_edge_vectors",
+            ufl.geometry.ReferenceFacetEdgeVectors: "reference_facet_edge_vectors",
+            ufl.geometry.ReferenceNormal: "reference_normals",
             ufl.geometry.FacetOrientation: "facet_orientation",
         }
         cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()}  # type: ignore

test/test_jit_expression.py

@@ -323,3 +323,110 @@ def test_facet_expression(compile_args):
 
         # Check that facet normal is pointing out of the cell
         assert np.dot(midpoint - coords[i], output) > 0
+
+
+def test_facet_geometry_expressions(compile_args):
+    """Test various geometrical quantities for facet expressions."""
+    cell = basix.CellType.triangle
+    mesh = ufl.Mesh(basix.ufl.element("Lagrange", cell, 1, shape=(2,)))
+    dtype = np.float64
+    points = np.array([[0.5]], dtype=dtype)
+    c_type = "double"
+    c_xtype = "double"
+    ffi = cffi.FFI()
+
+    # Prepare reference geometry and working arrays
+    coords = np.array([[1, 0, 0], [3, 0, 0], [0, 2, 0]], dtype=dtype)
+    u_coeffs = np.array([], dtype=dtype)
+    consts = np.array([], dtype=dtype)
+    entity_index = np.empty(1, dtype=np.intc)
+    quad_perm = np.array([0], dtype=np.dtype("uint8"))
+    ffi_data = {
+        "const": ffi.cast(f"{c_type} *", consts.ctypes.data),
+        "coeff": ffi.cast(f"{c_type} *", u_coeffs.ctypes.data),
+        "coords": ffi.cast(f"{c_xtype} *", coords.ctypes.data),
+        "entity_index": ffi.cast("int *", entity_index.ctypes.data),
+        "quad_perm": ffi.cast("uint8_t *", quad_perm.ctypes.data),
+    }
+
+    def check_expression(expression_class, output_shape, entity_values, reference_values):
+        obj = ffcx.codegeneration.jit.compile_expressions(
+            [(expression_class(mesh), points)], cffi_extra_compile_args=compile_args
+        )[0][0]
+        output = np.zeros(output_shape, dtype=dtype)
+        for i, ref_val in enumerate(reference_values):
+            output[:] = 0
+            entity_index[0] = i
+            obj.tabulate_tensor_float64(
+                ffi.cast(f"{c_type} *", output.ctypes.data),
+                ffi_data["coeff"],
+                ffi_data["const"],
+                ffi_data["coords"],
+                ffi_data["entity_index"],
+                ffi_data["quad_perm"],
+            )
+            np.testing.assert_allclose(output, ref_val)
+
+    check_expression(
+        ufl.geometry.CellFacetJacobian, (2, 1), entity_index, basix.cell.facet_jacobians(cell)
+    )
+    check_expression(
+        ufl.geometry.ReferenceFacetVolume,
+        (1,),
+        entity_index,
+        basix.cell.facet_reference_volumes(cell),
+    )
+    check_expression(
+        ufl.geometry.ReferenceCellEdgeVectors,
+        (3, 2),
+        entity_index,
+        np.array(
+            [
+                [
+                    basix.geometry(cell)[j] - basix.geometry(cell)[i]
+                    for i, j in basix.topology(cell)[1]
+                ]
+            ]
+        ),
+    )
+
+
+def test_facet_geometry_expressions_3D(compile_args):
+    cell = basix.CellType.tetrahedron
+    c_el = basix.ufl.element("Lagrange", cell, 1, shape=(3,))
+    mesh = ufl.Mesh(c_el)
+    dtype = np.float64
+    points = np.array([[0.33, 0.33]], dtype=dtype)
+    c_type = "double"
+    c_xtype = "double"
+    ffi = cffi.FFI()
+
+    # Prepare reference geometry and working arrays
+    coords = np.array([[1, 0, 0], [3, 0, 0], [0, 2, 0], [0, 0, 1]], dtype=dtype)
+    u_coeffs = np.array([], dtype=dtype)
+    consts = np.array([], dtype=dtype)
+    entity_index = np.empty(1, dtype=np.intc)
+    quad_perm = np.array([0], dtype=np.dtype("uint8"))
+
+    # Check ReferenceFacetEdgeVectors
+    output = np.zeros((3, 3))
+    triangle_edges = basix.topology(basix.CellType.triangle)[1]
+    ref_fev = []
+    topology = basix.topology(cell)
+    geometry = basix.geometry(cell)
+    for facet in topology[-2]:
+        ref_fev += [geometry[facet[j]] - geometry[facet[i]] for i, j in triangle_edges]
+
+    ref_fev_code = ffcx.codegeneration.jit.compile_expressions(
+        [(ufl.geometry.ReferenceFacetEdgeVectors(mesh), points)],
+        cffi_extra_compile_args=compile_args,
+    )[0][0]
+    ref_fev_code.tabulate_tensor_float64(
+        ffi.cast(f"{c_type} *", output.ctypes.data),
+        ffi.cast(f"{c_type} *", u_coeffs.ctypes.data),
+        ffi.cast(f"{c_type} *", consts.ctypes.data),
+        ffi.cast(f"{c_xtype} *", coords.ctypes.data),
+        ffi.cast("int *", entity_index.ctypes.data),
+        ffi.cast("uint8_t *", quad_perm.ctypes.data),
+    )
+    np.testing.assert_allclose(output, np.asarray(ref_fev)[:3, :])