From 9d227a48b48e01982fd87c2527a9a202ab208cf3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= Date: Mon, 14 Oct 2024 12:19:20 +0000 Subject: [PATCH 1/6] Add many geometrical quantities to ffcx expression generator. Fix various errors in the definition of geometry tables --- ffcx/codegeneration/access.py | 7 +++---- ffcx/codegeneration/expression_generator.py | 8 +++++++- ffcx/codegeneration/geometry.py | 2 +- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/ffcx/codegeneration/access.py b/ffcx/codegeneration/access.py index 18bcf9c89..5cd361bea 100644 --- a/ffcx/codegeneration/access.py +++ b/ffcx/codegeneration/access.py @@ -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}_facet_reference_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] diff --git a/ffcx/codegeneration/expression_generator.py b/ffcx/codegeneration/expression_generator.py index a3c0c56f0..62f9ee45a 100644 --- a/ffcx/codegeneration/expression_generator.py +++ b/ffcx/codegeneration/expression_generator.py @@ -61,10 +61,16 @@ 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_vertices", + ufl.geometry.CellFacetJacobian: "reference_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.FacetJacobianDeterminant: "reference_facet_jacobian", ufl.geometry.ReferenceNormal: "reference_facet_normals", + ufl.geometry.FacetOrientation: "facet_orientation", } cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()} # type: ignore diff --git a/ffcx/codegeneration/geometry.py b/ffcx/codegeneration/geometry.py index 39a5f5901..dc0b280b6 100644 --- a/ffcx/codegeneration/geometry.py +++ b/ffcx/codegeneration/geometry.py @@ -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) From 24bdca97e8bb0ab97404002c728b5afb8fbb75e5 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Sun, 17 Nov 2024 20:59:40 +0000 Subject: [PATCH 2/6] Add some tests and remove expressions that I don't think will be used or is sensible --- ffcx/codegeneration/expression_generator.py | 2 - test/test_jit_expression.py | 136 ++++++++++++++++++++ 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/ffcx/codegeneration/expression_generator.py b/ffcx/codegeneration/expression_generator.py index 62f9ee45a..6712f01a5 100644 --- a/ffcx/codegeneration/expression_generator.py +++ b/ffcx/codegeneration/expression_generator.py @@ -68,9 +68,7 @@ def generate_geometry_tables(self): ufl.geometry.ReferenceFacetVolume: "reference_facet_volume", ufl.geometry.ReferenceCellEdgeVectors: "reference_edge_vectors", ufl.geometry.ReferenceFacetEdgeVectors: "facet_reference_edge_vectors", - ufl.geometry.FacetJacobianDeterminant: "reference_facet_jacobian", ufl.geometry.ReferenceNormal: "reference_facet_normals", - ufl.geometry.FacetOrientation: "facet_orientation", } cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()} # type: ignore diff --git a/test/test_jit_expression.py b/test/test_jit_expression.py index 9dd3724ef..523b27abf 100644 --- a/test/test_jit_expression.py +++ b/test/test_jit_expression.py @@ -323,3 +323,139 @@ 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 expressiors.""" + cell = basix.CellType.triangle + c_el = basix.ufl.element("Lagrange", cell, 1, shape=(2,)) + mesh = ufl.Mesh(c_el) + 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_const = ffi.cast(f"{c_type} *", consts.ctypes.data) + ffi_coeff = ffi.cast(f"{c_type} *", u_coeffs.ctypes.data) + ffi_coords = ffi.cast(f"{c_xtype} *", coords.ctypes.data) + ffi_ei = ffi.cast("int *", entity_index.ctypes.data) + ffi_qp = ffi.cast("uint8_t *", quad_perm.ctypes.data) + + # Check CellFacetJacobian + obj_fj = ffcx.codegeneration.jit.compile_expressions( + [(ufl.geometry.CellFacetJacobian(mesh), points)], cffi_extra_compile_args=compile_args + )[0][0] + output = np.zeros((2, 1), dtype=dtype) + obj_det = ffcx.codegeneration.jit.compile_expressions( + [(ufl.geometry.FacetJacobianDeterminant(mesh), points)], + cffi_extra_compile_args=compile_args, + )[0][0] + output_det = np.zeros(1, dtype=dtype) + reference_facet_jacobians = basix.cell.facet_jacobians(cell) + for i, ref_fj in enumerate(reference_facet_jacobians): + output[:] = 0 + entity_index[0] = i + obj_fj.tabulate_tensor_float64( + ffi.cast(f"{c_type} *", output.ctypes.data), + ffi_coeff, + ffi_const, + ffi_coords, + ffi_ei, + ffi_qp, + ) + np.testing.assert_allclose(output, ref_fj) + + obj_det.tabulate_tensor_float64( + ffi.cast(f"{c_type} *", output_det.ctypes.data), + ffi_coeff, + ffi_const, + ffi_coords, + ffi_ei, + ffi_qp, + ) + + # Check CellVolume + ref_fj_code = ffcx.codegeneration.jit.compile_expressions( + [(ufl.geometry.ReferenceFacetVolume(mesh), points)], cffi_extra_compile_args=compile_args + )[0][0] + output_fv = np.zeros(1, dtype=dtype) + reference_facet_volumes = basix.cell.facet_reference_volumes(cell) + for i, ref_fv in enumerate(reference_facet_volumes): + output_fv[:] = 0 + entity_index[0] = i + ref_fj_code.tabulate_tensor_float64( + ffi.cast(f"{c_type} *", output_fv.ctypes.data), + ffi_coeff, + ffi_const, + ffi_coords, + ffi_ei, + ffi_qp, + ) + np.testing.assert_allclose(output_fv, ref_fv) + + # Check ReferenceCellEdgeVectors + ref_ev_code = ffcx.codegeneration.jit.compile_expressions( + [(ufl.geometry.ReferenceCellEdgeVectors(mesh), points)], + cffi_extra_compile_args=compile_args, + )[0][0] + output_ev = np.zeros((3, 2), dtype=dtype) + topology = basix.topology(cell) + geometry = basix.geometry(cell) + edge_vectors = np.array([geometry[j] - geometry[i] for i, j in topology[1]]) + ref_ev_code.tabulate_tensor_float64( + ffi.cast(f"{c_type} *", output_ev.ctypes.data), + ffi_coeff, + ffi_const, + ffi_coords, + ffi_ei, + ffi_qp, + ) + np.testing.assert_allclose(output_ev, edge_vectors) + + +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, :]) From f45178bf97ed87ebf183725f69695a672a3f8b6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= Date: Mon, 18 Nov 2024 09:15:08 +0000 Subject: [PATCH 3/6] Rename to match ufl names --- ffcx/codegeneration/access.py | 8 +++---- ffcx/codegeneration/expression_generator.py | 10 ++++----- ffcx/codegeneration/geometry.py | 24 ++++++++++----------- ffcx/codegeneration/integral_generator.py | 8 +++---- test/test_jit_expression.py | 4 ++-- 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/ffcx/codegeneration/access.py b/ffcx/codegeneration/access.py index 5cd361bea..c74560e3a 100644 --- a/ffcx/codegeneration/access.py +++ b/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,7 +263,7 @@ 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}_facet_reference_edge_vectors", dtype=L.DataType.REAL) + 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( diff --git a/ffcx/codegeneration/expression_generator.py b/ffcx/codegeneration/expression_generator.py index 6712f01a5..fdfe2aba1 100644 --- a/ffcx/codegeneration/expression_generator.py +++ b/ffcx/codegeneration/expression_generator.py @@ -62,13 +62,13 @@ def generate(self): 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.FacetEdgeVectors: "facet_edge_vectors", + 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", } cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()} # type: ignore diff --git a/ffcx/codegeneration/geometry.py b/ffcx/codegeneration/geometry.py index dc0b280b6..952238e25 100644 --- a/ffcx/codegeneration/geometry.py +++ b/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) diff --git a/ffcx/codegeneration/integral_generator.py b/ffcx/codegeneration/integral_generator.py index c94b2ef5a..8545e49c1 100644 --- a/ffcx/codegeneration/integral_generator.py +++ b/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 diff --git a/test/test_jit_expression.py b/test/test_jit_expression.py index 523b27abf..5744e8fee 100644 --- a/test/test_jit_expression.py +++ b/test/test_jit_expression.py @@ -358,8 +358,8 @@ def test_facet_geometry_expressions(compile_args): cffi_extra_compile_args=compile_args, )[0][0] output_det = np.zeros(1, dtype=dtype) - reference_facet_jacobians = basix.cell.facet_jacobians(cell) - for i, ref_fj in enumerate(reference_facet_jacobians): + cell_facet_jacobians = basix.cell.facet_jacobians(cell) + for i, ref_fj in enumerate(cell_facet_jacobians): output[:] = 0 entity_index[0] = i obj_fj.tabulate_tensor_float64( From 6e3ff7aa1fa5aa8632e501288dbe516b237c9479 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= Date: Mon, 18 Nov 2024 13:13:56 +0100 Subject: [PATCH 4/6] Apply suggestions from code review Remove test of facet det J --- test/test_jit_expression.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/test/test_jit_expression.py b/test/test_jit_expression.py index 5744e8fee..284a84841 100644 --- a/test/test_jit_expression.py +++ b/test/test_jit_expression.py @@ -353,11 +353,6 @@ def test_facet_geometry_expressions(compile_args): [(ufl.geometry.CellFacetJacobian(mesh), points)], cffi_extra_compile_args=compile_args )[0][0] output = np.zeros((2, 1), dtype=dtype) - obj_det = ffcx.codegeneration.jit.compile_expressions( - [(ufl.geometry.FacetJacobianDeterminant(mesh), points)], - cffi_extra_compile_args=compile_args, - )[0][0] - output_det = np.zeros(1, dtype=dtype) cell_facet_jacobians = basix.cell.facet_jacobians(cell) for i, ref_fj in enumerate(cell_facet_jacobians): output[:] = 0 @@ -372,14 +367,6 @@ def test_facet_geometry_expressions(compile_args): ) np.testing.assert_allclose(output, ref_fj) - obj_det.tabulate_tensor_float64( - ffi.cast(f"{c_type} *", output_det.ctypes.data), - ffi_coeff, - ffi_const, - ffi_coords, - ffi_ei, - ffi_qp, - ) # Check CellVolume ref_fj_code = ffcx.codegeneration.jit.compile_expressions( From 24f2094fb7e1a3c2283bad8eb31c510327c94aff Mon Sep 17 00:00:00 2001 From: jorgensd Date: Mon, 18 Nov 2024 18:18:42 +0000 Subject: [PATCH 5/6] Ruff formatting --- test/test_jit_expression.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/test_jit_expression.py b/test/test_jit_expression.py index 284a84841..66398dad2 100644 --- a/test/test_jit_expression.py +++ b/test/test_jit_expression.py @@ -367,7 +367,6 @@ def test_facet_geometry_expressions(compile_args): ) np.testing.assert_allclose(output, ref_fj) - # Check CellVolume ref_fj_code = ffcx.codegeneration.jit.compile_expressions( [(ufl.geometry.ReferenceFacetVolume(mesh), points)], cffi_extra_compile_args=compile_args From dbd5f2ca8fabb764e9badaf6581802280badcd20 Mon Sep 17 00:00:00 2001 From: Michal Habera Date: Wed, 20 Nov 2024 16:00:49 +0100 Subject: [PATCH 6/6] Simplify (#729) --- test/test_jit_expression.py | 113 ++++++++++++++++-------------------- 1 file changed, 49 insertions(+), 64 deletions(-) diff --git a/test/test_jit_expression.py b/test/test_jit_expression.py index 66398dad2..1bf84a4b5 100644 --- a/test/test_jit_expression.py +++ b/test/test_jit_expression.py @@ -326,10 +326,9 @@ def test_facet_expression(compile_args): def test_facet_geometry_expressions(compile_args): - """Test various geometrical quantities for facet expressiors.""" + """Test various geometrical quantities for facet expressions.""" cell = basix.CellType.triangle - c_el = basix.ufl.element("Lagrange", cell, 1, shape=(2,)) - mesh = ufl.Mesh(c_el) + mesh = ufl.Mesh(basix.ufl.element("Lagrange", cell, 1, shape=(2,))) dtype = np.float64 points = np.array([[0.5]], dtype=dtype) c_type = "double" @@ -342,68 +341,54 @@ def test_facet_geometry_expressions(compile_args): consts = np.array([], dtype=dtype) entity_index = np.empty(1, dtype=np.intc) quad_perm = np.array([0], dtype=np.dtype("uint8")) - ffi_const = ffi.cast(f"{c_type} *", consts.ctypes.data) - ffi_coeff = ffi.cast(f"{c_type} *", u_coeffs.ctypes.data) - ffi_coords = ffi.cast(f"{c_xtype} *", coords.ctypes.data) - ffi_ei = ffi.cast("int *", entity_index.ctypes.data) - ffi_qp = ffi.cast("uint8_t *", quad_perm.ctypes.data) - - # Check CellFacetJacobian - obj_fj = ffcx.codegeneration.jit.compile_expressions( - [(ufl.geometry.CellFacetJacobian(mesh), points)], cffi_extra_compile_args=compile_args - )[0][0] - output = np.zeros((2, 1), dtype=dtype) - cell_facet_jacobians = basix.cell.facet_jacobians(cell) - for i, ref_fj in enumerate(cell_facet_jacobians): - output[:] = 0 - entity_index[0] = i - obj_fj.tabulate_tensor_float64( - ffi.cast(f"{c_type} *", output.ctypes.data), - ffi_coeff, - ffi_const, - ffi_coords, - ffi_ei, - ffi_qp, - ) - np.testing.assert_allclose(output, ref_fj) - - # Check CellVolume - ref_fj_code = ffcx.codegeneration.jit.compile_expressions( - [(ufl.geometry.ReferenceFacetVolume(mesh), points)], cffi_extra_compile_args=compile_args - )[0][0] - output_fv = np.zeros(1, dtype=dtype) - reference_facet_volumes = basix.cell.facet_reference_volumes(cell) - for i, ref_fv in enumerate(reference_facet_volumes): - output_fv[:] = 0 - entity_index[0] = i - ref_fj_code.tabulate_tensor_float64( - ffi.cast(f"{c_type} *", output_fv.ctypes.data), - ffi_coeff, - ffi_const, - ffi_coords, - ffi_ei, - ffi_qp, - ) - np.testing.assert_allclose(output_fv, ref_fv) - - # Check ReferenceCellEdgeVectors - ref_ev_code = ffcx.codegeneration.jit.compile_expressions( - [(ufl.geometry.ReferenceCellEdgeVectors(mesh), points)], - cffi_extra_compile_args=compile_args, - )[0][0] - output_ev = np.zeros((3, 2), dtype=dtype) - topology = basix.topology(cell) - geometry = basix.geometry(cell) - edge_vectors = np.array([geometry[j] - geometry[i] for i, j in topology[1]]) - ref_ev_code.tabulate_tensor_float64( - ffi.cast(f"{c_type} *", output_ev.ctypes.data), - ffi_coeff, - ffi_const, - ffi_coords, - ffi_ei, - ffi_qp, + 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] + ] + ] + ), ) - np.testing.assert_allclose(output_ev, edge_vectors) def test_facet_geometry_expressions_3D(compile_args):