Finsberg/fix extrema dtype float32 spaces

src/scifem/assembly.py

@@ -9,19 +9,19 @@
 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:
         form: The form to extract the data type from.
 
     """
     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()

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)

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,35 +131,31 @@ 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
     )
 
     # 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

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)

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)