diff --git a/src/scifem/assembly.py b/src/scifem/assembly.py index af65ddc..e08b624 100644 --- a/src/scifem/assembly.py +++ b/src/scifem/assembly.py @@ -9,8 +9,8 @@ import typing -def _extract_dtype(expr: ufl.form.Form | ufl.core.expr.Expr | dolfinx.fem.Form) -> npt.DTypeLike: - """Given a ufl form, expression or a compiled DOLFINx form extract the data type that should +def _extract_dtype(form: ufl.form.Form | dolfinx.fem.Form) -> npt.DTypeLike: + """Given a ufl form or a compiled DOLFINx form extract the data type that should be used in form compilation. Args: @@ -18,10 +18,10 @@ def _extract_dtype(expr: ufl.form.Form | ufl.core.expr.Expr | dolfinx.fem.Form) """ try: - # Compiled forms, as well as Functions and Constants, carry a dtype - return expr.dtype # type: ignore[union-attr] + # Compiled forms carry a dtype + return form.dtype # type: ignore[union-attr] except AttributeError: - form = typing.cast(ufl.form.Form, expr) + form = typing.cast(ufl.form.Form, form) domain = typing.cast(ufl.Mesh, ufl.domain.extract_domains(form)[0]) geom_type = domain.ufl_cargo().geometry.x.dtype coefficients = form.coefficients() diff --git a/src/scifem/eval.py b/src/scifem/eval.py index 512cd84..5c5fa5a 100644 --- a/src/scifem/eval.py +++ b/src/scifem/eval.py @@ -281,7 +281,7 @@ def compute_LP_average( def compute_extrema( u: ufl.core.expr.Expr, - kind=MinMaxFunc, + kind: MinMaxFunc, p: int = 1, num_candidates: int = 3, x0: npt.NDArray[np.floating] | None = None, @@ -320,7 +320,7 @@ def compute_extrema( elif kind is builtins.min: sign = 1 else: - raise ValueError("Unknown extrema {kind}, supports min and max only.") + raise ValueError(f"Unknown extrema {kind}, supports min and max only.") # Compute cell candidates local_avg = compute_LP_average(u, p, mesh) diff --git a/src/scifem/spaces.py b/src/scifem/spaces.py index 70c114d..7d285a6 100644 --- a/src/scifem/spaces.py +++ b/src/scifem/spaces.py @@ -94,8 +94,24 @@ def create_space_of_simple_functions( # It will be the process that owns the cell with global index 0 is_owner = cell_map.local_range[0] == 0 and cell_map.size_local > 0 + _fe_constructor: ( + type[dolfinx.cpp.fem.FiniteElement_float32] | type[dolfinx.cpp.fem.FiniteElement_float64] + ) + _fem_constructor: ( + type[dolfinx.cpp.fem.FunctionSpace_float32] | type[dolfinx.cpp.fem.FunctionSpace_float64] + ) + dtype = mesh.geometry.x.dtype + if dtype == np.float64: + _fe_constructor = dolfinx.cpp.fem.FiniteElement_float64 + _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float64 + elif dtype == np.float32: + _fe_constructor = dolfinx.cpp.fem.FiniteElement_float32 + _fem_constructor = dolfinx.cpp.fem.FunctionSpace_float32 + else: + raise ValueError(f"Unsupported dtype: {dtype}") + # Determine number of dofs and ghosts on each process - el = basix.ufl.element("P", mesh.basix_cell(), 0, shape=(), discontinuous=True) + el = basix.ufl.element("P", mesh.basix_cell(), 0, shape=(), discontinuous=True, dtype=dtype) num_dofs = len(tags) * is_owner num_ghosts = len(tags) * (not is_owner) (_, ghost_owner_rank) = mesh.comm.allreduce((is_owner, mesh.comm.rank), op=MPI.MAXLOC) @@ -115,16 +131,14 @@ def create_space_of_simple_functions( ) # Create dofmap mapping local dofs to position in `tags` input. - adj_flattened = np.zeros_like(cell_tag.values) + adj_flattened = np.zeros(len(cell_tag.values), dtype=np.int32) for i, tag in enumerate(tags): # numpy integer tags are accepted at runtime; the stub only admits int | float adj_flattened[cell_tag.find(tag)] = i # type: ignore[arg-type] # Create dofmap - # Cell tag values are int32 at runtime; the stub admits any MeshTags value dtype dofmap_adj = dolfinx.cpp.graph.AdjacencyList_int32( - adj_flattened, # type: ignore[arg-type] - np.arange(len(cell_tag.values) + 1, dtype=np.int32), + adj_flattened, np.arange(len(cell_tag.values) + 1, dtype=np.int32) ) cpp_dofmap = dolfinx.cpp.fem.DofMap( e_layout, fenicsx_compat.unwrap_index_map(dof_imap), index_map_bs, dofmap_adj, bs @@ -132,18 +146,16 @@ def create_space_of_simple_functions( # Create function space cpp_el = fenicsx_compat.finite_element_ctor_kwargs( - dolfinx.cpp.fem.FiniteElement_float64, + _fe_constructor, el.basix_element._e, gdim=mesh.geometry.dim, value_shape=value_shape, ) - # Only float64 meshes are supported here - cpp_space = dolfinx.cpp.fem.FunctionSpace_float64( - mesh._cpp_object, # type: ignore[arg-type] - cpp_el, - cpp_dofmap, - ) + # The constructor matches the mesh dtype at runtime, which the stubs cannot express + cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap) # type: ignore[arg-type] - _el = basix.ufl.element("P", mesh.basix_cell(), 0, shape=value_shape, discontinuous=True) + _el = basix.ufl.element( + "P", mesh.basix_cell(), 0, shape=value_shape, discontinuous=True, dtype=dtype + ) V = dolfinx.fem.FunctionSpace(mesh, _el, cpp_space) return V diff --git a/tests/test_eval.py b/tests/test_eval.py index 5a4a5f9..26d3c20 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -156,3 +156,17 @@ def f(x): assert np.isclose(u_ufl_ex, -sign * 0.8, atol=tol) np.testing.assert_allclose(_X_ufl_ex, x_p, atol=tol) + + +def test_extrema_requires_kind(): + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 2, 2) + x = ufl.SpatialCoordinate(mesh) + with pytest.raises(TypeError, match="kind"): + compute_extrema(x[0]) + + +def test_extrema_unknown_kind(): + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 2, 2) + x = ufl.SpatialCoordinate(mesh) + with pytest.raises(ValueError, match=f"Unknown extrema {sum}"): + compute_extrema(x[0], sum) diff --git a/tests/test_material_space.py b/tests/test_material_space.py index 98f1394..a5a9b0b 100644 --- a/tests/test_material_space.py +++ b/tests/test_material_space.py @@ -86,3 +86,26 @@ def top(x, tol=1e-8): diff = ufl.inner(f - f_ex, f - f_ex) * dx_restricted error = assemble_scalar(diff) np.testing.assert_allclose(error, 0.0) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("tag_dtype", [np.int32, np.int64]) +def test_material_space_dtypes(dtype, tag_dtype): + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8, dtype=dtype) + cell_map = mesh.topology.index_map(mesh.topology.dim) + all_cells = np.arange(cell_map.size_local + cell_map.num_ghosts, dtype=np.int32) + tags = [2, 7] + values = np.full(len(all_cells), tags[0], dtype=tag_dtype) + values[dolfinx.mesh.locate_entities(mesh, mesh.topology.dim, lambda x: x[0] <= 0.5)] = tags[1] + ct = dolfinx.mesh.meshtags(mesh, mesh.topology.dim, all_cells, values) + + V = create_space_of_simple_functions(mesh, ct, tags) + assert V.mesh.geometry.x.dtype == dtype + for i, tag in enumerate(tags): + np.testing.assert_allclose(V.dofmap.list[ct.indices[ct.values == tag]].flatten(), i) + + # The left half carries tags[1], so the integral of f is 3 * 0.5 + 5 * 0.5 + f = dolfinx.fem.Function(V, dtype=dtype) + f.x.array[:] = [3, 5] + integral = assemble_scalar(dolfinx.fem.form(f * ufl.dx, dtype=dtype)) + np.testing.assert_allclose(integral, 4.0, rtol=10 * np.finfo(dtype).eps)