Skip to content

Commit

Permalink
Add many geometrical quantities to ffcx expression generator. (#725)
Browse files Browse the repository at this point in the history
* Add many geometrical quantities to ffcx expression generator.

Fix various errors in the definition of geometry tables

* Add some tests and remove expressions that I don't think will be used or is sensible

* Rename to match ufl names

* Apply suggestions from code review

Remove test of facet det J

* Ruff formatting

* Simplify (#729)

---------

Co-authored-by: Michal Habera <michal.habera@gmail.com>
  • Loading branch information
jorgensd and michalhabera authored Nov 22, 2024
1 parent 3a6eb07 commit 712a57d
Show file tree
Hide file tree
Showing 5 changed files with 136 additions and 26 deletions.
13 changes: 6 additions & 7 deletions ffcx/codegeneration/access.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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":
Expand All @@ -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(
Expand All @@ -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 "
Expand All @@ -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]

Expand Down
8 changes: 6 additions & 2 deletions ffcx/codegeneration/expression_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
26 changes: 13 additions & 13 deletions ffcx/codegeneration/geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
8 changes: 4 additions & 4 deletions ffcx/codegeneration/integral_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
107 changes: 107 additions & 0 deletions test/test_jit_expression.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, :])

0 comments on commit 712a57d

Please sign in to comment.