diff --git a/cpp/dolfinx/fem/assemble_expression_impl.h b/cpp/dolfinx/fem/assemble_expression_impl.h index 7f3dcd9ddd..80a3729428 100644 --- a/cpp/dolfinx/fem/assemble_expression_impl.h +++ b/cpp/dolfinx/fem/assemble_expression_impl.h @@ -61,11 +61,11 @@ namespace dolfinx::fem::impl /// @param[in] perms Entity permutation information for use in `fn`. template void tabulate_expression( - std::span values, fem::FEkernel auto fn, + std::span values, fem::FEkernel auto fn, std::array Xshape, std::size_t value_size, std::size_t num_argument_dofs, md::mdspan> x_dofmap, - std::span> x, + std::span x, md::mdspan> coeffs, std::span constants, fem::MDSpan2 auto entities, std::span cell_info, @@ -153,7 +153,7 @@ void tabulate_expression( /// expression values at the evaluation points. template void tabulate_expression( - std::span values, fem::FEkernel auto fn, + std::span values, fem::FEkernel auto fn, std::array Xshape, std::size_t value_size, md::mdspan> coeffs, std::span constants, const mesh::Mesh& mesh, diff --git a/cpp/dolfinx/fem/assemble_matrix_impl.h b/cpp/dolfinx/fem/assemble_matrix_impl.h index 01f5940968..9257d096cd 100644 --- a/cpp/dolfinx/fem/assemble_matrix_impl.h +++ b/cpp/dolfinx/fem/assemble_matrix_impl.h @@ -62,18 +62,16 @@ using mdspan2_t = md::mdspan>; /// function mesh. /// @param cell_info1 Cell permutation information for the trial /// function mesh. -template +template void assemble_cells_matrix( la::MatSet auto mat_set, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, std::span cells, std::tuple> dofmap0, fem::DofTransformKernel auto P0, std::tuple> dofmap1, fem::DofTransformKernel auto P1T, std::span bc0, - std::span bc1, FEkernel auto kernel, + std::span bc1, FEkernel auto kernel, md::mdspan> coeffs, std::span constants, std::span cell_info0, std::span cell_info1) @@ -90,7 +88,7 @@ void assemble_cells_matrix( const int ndim0 = bs0 * num_dofs0; const int ndim1 = bs1 * num_dofs1; std::vector Ae(ndim0 * ndim1); - std::vector> cdofs(3 * x_dofmap.extent(1)); + std::vector cdofs(3 * x_dofmap.extent(1)); // Iterate over active cells assert(cells0.size() == cells.size()); @@ -226,12 +224,10 @@ void assemble_cells_matrix( /// function mesh. /// @param[in] perms Entity permutation integer. Empty if entity /// permutations are not required. -template +template void assemble_entities( la::MatSet auto mat_set, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, md::mdspan> entities, @@ -245,7 +241,7 @@ void assemble_entities( std::extents>> dofmap1, fem::DofTransformKernel auto P1T, std::span bc0, - std::span bc1, FEkernel auto kernel, + std::span bc1, FEkernel auto kernel, md::mdspan> coeffs, std::span constants, std::span cell_info0, std::span cell_info1, @@ -258,7 +254,7 @@ void assemble_entities( const auto [dmap1, bs1, entities1] = dofmap1; // Data structures used in assembly - std::vector> cdofs(3 * x_dofmap.extent(1)); + std::vector cdofs(3 * x_dofmap.extent(1)); const int num_dofs0 = dmap0.extent(1); const int num_dofs1 = dmap1.extent(1); const int ndim0 = bs0 * num_dofs0; @@ -393,12 +389,10 @@ void assemble_entities( /// function mesh. /// @param[in] perms Facet permutation integer. Empty if facet /// permutations are not required. -template +template void assemble_interior_facets( la::MatSet auto mat_set, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, md::mdspan> facets, @@ -412,7 +406,7 @@ void assemble_interior_facets( std::extents>> dofmap1, fem::DofTransformKernel auto P1T, std::span bc0, - std::span bc1, FEkernel auto kernel, + std::span bc1, FEkernel auto kernel, md::mdspan> coeffs, @@ -427,7 +421,7 @@ void assemble_interior_facets( const auto [dmap1, bs1, facets1] = dofmap1; // Data structures used in assembly - using X = scalar_value_t; + using X = U; std::vector cdofs(2 * x_dofmap.extent(1) * 3); std::span cdofs0(cdofs.data(), x_dofmap.extent(1) * 3); std::span cdofs1(cdofs.data() + x_dofmap.extent(1) * 3, @@ -612,9 +606,7 @@ void assemble_interior_facets( template void assemble_matrix( la::MatSet auto mat_set, const Form& a, - md::mdspan, - md::extents> - x, + md::mdspan> x, std::span constants, const std::map, std::pair, int>>& coefficients, @@ -687,7 +679,7 @@ void assemble_matrix( std::span cells1 = a.domain_arg(IntegralType::cell, 1, i, cell_type_idx); auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i}); assert(cells.size() * cstride == coeffs.size()); - impl::assemble_cells_matrix( + impl::assemble_cells_matrix( mat_set, x_dofmap, x, cells, {dofs0, bs0, cells0}, P0, {dofs1, bs1, cells1}, P1T, bc0, bc1, fn, md::mdspan(coeffs.data(), cells.size(), cstride), constants, @@ -732,7 +724,7 @@ void assemble_matrix( std::span facets0 = a.domain_arg(IntegralType::interior_facet, 0, i, 0); std::span facets1 = a.domain_arg(IntegralType::interior_facet, 1, i, 0); assert((facets.size() / 4) * 2 * cstride == coeffs.size()); - impl::assemble_interior_facets( + impl::assemble_interior_facets( mat_set, x_dofmap, x, mdspanx22_t(facets.data(), facets.size() / 4, 2, 2), {*dofmap0, bs0, @@ -777,7 +769,7 @@ void assemble_matrix( std::span e1 = a.domain_arg(itg_type, 1, i, 0); mdspanx2_t entities1(e1.data(), e1.size() / 2, 2); assert((entities.size() / 2) * cstride == coeffs.size()); - impl::assemble_entities( + impl::assemble_entities( mat_set, x_dofmap, x, entities, {dofs0, bs0, entities0}, P0, {dofs1, bs1, entities1}, P1T, bc0, bc1, fn, md::mdspan(coeffs.data(), entities.extent(0), cstride), constants, diff --git a/cpp/dolfinx/fem/assemble_scalar_impl.h b/cpp/dolfinx/fem/assemble_scalar_impl.h index 23070de9f9..d2c50dbffb 100644 --- a/cpp/dolfinx/fem/assemble_scalar_impl.h +++ b/cpp/dolfinx/fem/assemble_scalar_impl.h @@ -17,20 +17,20 @@ #include #include #include +#include #include namespace dolfinx::fem::impl { /// Assemble functional over cells -template -T assemble_cells(mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, - std::span cells, FEkernel auto fn, - std::span constants, - md::mdspan> coeffs, - std::span> cdofs_b) +template +T assemble_cells( + mdspan2_t x_dofmap, + md::mdspan> x, + std::span cells, FEkernel auto fn, + std::span constants, + md::mdspan> coeffs, + std::span> cdofs_b) { T value(0); if (cells.empty()) @@ -65,19 +65,17 @@ T assemble_cells(mdspan2_t x_dofmap, /// However, entities may be attached to more than one cell. This function /// therefore computes 'one-sided' integrals, i.e. evaluates integrals as seen /// from cell used to define the entity. -template +template T assemble_entities( mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, md::mdspan> entities, - FEkernel auto fn, std::span constants, + FEkernel auto fn, std::span constants, md::mdspan> coeffs, md::mdspan> perms, - std::span> cdofs_b) + std::span> cdofs_b) { T value(0); if (entities.empty()) @@ -106,21 +104,19 @@ T assemble_entities( } /// Assemble functional over interior facets -template +template T assemble_interior_facets( mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, md::mdspan> facets, - FEkernel auto fn, std::span constants, + FEkernel auto fn, std::span constants, md::mdspan> coeffs, md::mdspan> perms, - std::span> cdofs_b) + std::span> cdofs_b) { T value(0); if (facets.empty()) @@ -160,9 +156,7 @@ T assemble_interior_facets( template T assemble_scalar( const fem::Form& M, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, std::span constants, const std::map, std::pair, int>>& coefficients, @@ -171,7 +165,7 @@ T assemble_scalar( std::shared_ptr> mesh = M.mesh(); assert(mesh); - std::vector> cdofs_b(2 * 3 * x_dofmap.extent(1)); + std::vector cdofs_b(2 * 3 * x_dofmap.extent(1)); T value = 0; for (int i = 0; i < M.num_integrals(IntegralType::cell, cell_type_idx); ++i) diff --git a/cpp/dolfinx/fem/assemble_vector_impl.h b/cpp/dolfinx/fem/assemble_vector_impl.h index ce62cb0fda..caf50b0023 100644 --- a/cpp/dolfinx/fem/assemble_vector_impl.h +++ b/cpp/dolfinx/fem/assemble_vector_impl.h @@ -61,18 +61,16 @@ using mdspan2_t = md::mdspan>; /// coefficient for cell `i`. /// @param[in] cell_info0 Cell permutation information for the test /// function mesh. -template ::value_type> requires std::is_same_v::value_type, T> void assemble_cells( fem::DofTransformKernel auto P0, V&& b, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, std::span cells, std::tuple> dofmap, - FEkernel auto kernel, std::span constants, + FEkernel auto kernel, std::span constants, md::mdspan> coeffs, std::span cell_info0) { @@ -83,7 +81,7 @@ void assemble_cells( assert(_bs < 0 or _bs == bs); // Create data structures used in assembly - std::vector> cdofs(3 * x_dofmap.extent(1)); + std::vector cdofs(3 * x_dofmap.extent(1)); std::vector be(bs * dmap.extent(1)); // Iterate over active cells @@ -155,14 +153,12 @@ void assemble_cells( /// function mesh. /// @param[in] perms Entity permutation integer. Empty if entity /// permutations are not required. -template ::value_type> requires std::is_same_v::value_type, T> void assemble_entities( fem::DofTransformKernel auto P0, V&& b, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, md::mdspan> entities, @@ -170,7 +166,7 @@ void assemble_entities( md::mdspan>> dofmap, - FEkernel auto kernel, std::span constants, + FEkernel auto kernel, std::span constants, md::mdspan> coeffs, std::span cell_info0, md::mdspan> perms) @@ -183,7 +179,7 @@ void assemble_entities( // Create data structures used in assembly const int num_dofs = dmap.extent(1); - std::vector> cdofs(3 * x_dofmap.extent(1)); + std::vector cdofs(3 * x_dofmap.extent(1)); std::vector be(bs * num_dofs); assert(entities0.size() == entities.size()); for (std::size_t f = 0; f < entities.extent(0); ++f) @@ -251,14 +247,12 @@ void assemble_entities( /// function mesh. /// @param[in] perms Facet permutation integer. Empty if facet /// permutations are not required. -template ::value_type> requires std::is_same_v::value_type, T> void assemble_interior_facets( fem::DofTransformKernel auto P0, V&& b, mdspan2_t x_dofmap, - md::mdspan, - md::extents> - x, + md::mdspan> x, md::mdspan> facets, @@ -266,14 +260,14 @@ void assemble_interior_facets( md::mdspan>> dofmap, - FEkernel auto kernel, std::span constants, + FEkernel auto kernel, std::span constants, md::mdspan> coeffs, std::span cell_info0, md::mdspan> perms) { - using X = scalar_value_t; + using X = U; if (facets.empty()) return; @@ -567,9 +561,7 @@ template ::value_type, T> void assemble_vector( V&& b, const Form& L, - md::mdspan, - md::extents> - x, + md::mdspan> x, std::span constants, const std::map, std::pair, int>>& coefficients) @@ -765,22 +757,12 @@ void assemble_vector( std::pair, int>>& coefficients) { using mdspanx3_t - = md::mdspan, - md::extents>; + = md::mdspan>; std::shared_ptr> mesh = L.mesh(); assert(mesh); auto x = mesh->geometry().x(); - if constexpr (std::is_same_v>) - { - impl::assemble_vector(b, L, mdspanx3_t(x.data(), x.size() / 3, 3), - constants, coefficients); - } - else - { - std::vector> _x(x.begin(), x.end()); - impl::assemble_vector(b, L, mdspanx3_t(_x.data(), _x.size() / 3, 3), - constants, coefficients); - } + impl::assemble_vector(b, L, mdspanx3_t(x.data(), x.size() / 3, 3), constants, + coefficients); } } // namespace dolfinx::fem::impl diff --git a/cpp/dolfinx/fem/assembler.h b/cpp/dolfinx/fem/assembler.h index 0654e108e0..b9259b8da4 100644 --- a/cpp/dolfinx/fem/assembler.h +++ b/cpp/dolfinx/fem/assembler.h @@ -176,8 +176,7 @@ T assemble_scalar( std::pair, int>>& coefficients) { using mdspanx3_t - = md::mdspan, - md::extents>; + = md::mdspan>; std::shared_ptr> mesh = M.mesh(); assert(mesh); @@ -191,19 +190,9 @@ T assemble_scalar( // Geometry dofmap and data md::mdspan> x_dofmap = mesh->geometry().dofmaps().at(cell_type_idx); - if constexpr (std::is_same_v>) - { - val += impl::assemble_scalar(M, x_dofmap, - mdspanx3_t(x.data(), x.size() / 3, 3), - constants, coefficients, cell_type_idx); - } - else - { - std::vector> _x(x.begin(), x.end()); - val += impl::assemble_scalar(M, x_dofmap, - mdspanx3_t(_x.data(), _x.size() / 3, 3), - constants, coefficients, cell_type_idx); - } + val += impl::assemble_scalar(M, x_dofmap, + mdspanx3_t(x.data(), x.size() / 3, 3), + constants, coefficients, cell_type_idx); } return val; } @@ -469,25 +458,14 @@ void assemble_matrix( { common::Timer t_assm("[Assemble Matrix]"); using mdspanx3_t - = md::mdspan, - md::extents>; + = md::mdspan>; std::shared_ptr> mesh = a.mesh(); assert(mesh); std::span x = mesh->geometry().x(); - if constexpr (std::is_same_v>) - { - impl::assemble_matrix( - mat_add, a, mdspanx3_t(x.data(), x.size() / 3, 3), constants, - coefficients, dof_marker0, dof_marker1); - } - else - { - std::vector> _x(x.begin(), x.end()); - impl::assemble_matrix( - mat_add, a, mdspanx3_t(_x.data(), _x.size() / 3, 3), constants, - coefficients, dof_marker0, dof_marker1); - } + impl::assemble_matrix( + mat_add, a, mdspanx3_t(x.data(), x.size() / 3, 3), constants, + coefficients, dof_marker0, dof_marker1); } /// @brief Assemble bilinear form into a matrix diff --git a/cpp/dolfinx/fem/interpolate.h b/cpp/dolfinx/fem/interpolate.h index a3f6c2f4a8..87434c4ff3 100644 --- a/cpp/dolfinx/fem/interpolate.h +++ b/cpp/dolfinx/fem/interpolate.h @@ -310,7 +310,9 @@ void scatter_values(MPI_Comm comm, std::span src_ranks, template void interpolation_apply(U&& Pi, V&& data, std::span coeffs, int bs) { - using X = typename dolfinx::scalar_value_t; + // Geometry (real) scalar type, taken from the interpolation operator Pi + // rather than scalar_value_t so it is independent of the value scalar T. + using X = typename std::remove_cvref_t::value_type; // Compute coefficients = Pi * x (matrix-vector multiply) if (bs == 1) @@ -421,7 +423,7 @@ void interpolate_same_map(Function& u1, mesh::CellRange auto&& cells1, auto [i_m, im_shape] = element1->create_interpolation_operator(*element0); // Iterate over mesh and interpolate on each cell - using X = typename dolfinx::scalar_value_t; + using X = U; // geometry (real) type, independent of the value scalar T assert(cells0.size() == cells1.size()); for (auto cell0_it = cells0.begin(), cell1_it = cells1.begin(); cell0_it != cells0.end() and cell1_it != cells1.end(); @@ -663,7 +665,7 @@ void interpolate_nonmatching_maps(Function& u1, coeffs0[dof_bs0 * i + k] = array0[dof_bs0 * dofs0[i] + k]; // Evaluate v at the interpolation points (physical space values) - using X = typename dolfinx::scalar_value_t; + using X = U; // geometry (real) type, independent of the value scalar T for (std::size_t p = 0; p < Xshape[0]; ++p) { for (int k = 0; k < bs0; ++k) diff --git a/cpp/dolfinx/fem/kernel.h b/cpp/dolfinx/fem/kernel.h index 3c949864ca..a889043db0 100644 --- a/cpp/dolfinx/fem/kernel.h +++ b/cpp/dolfinx/fem/kernel.h @@ -34,6 +34,9 @@ using kernel_t = std::function> constexpr kernel_t extract_kernel(const ufcx_integral* integral) { + static_assert(std::is_same_v>, + "UFCx kernels require geometry type U == scalar_value_t."); + if constexpr (std::is_same_v) return integral->tabulate_tensor_float32; else if constexpr (std::is_same_v) diff --git a/cpp/dolfinx/fem/traits.h b/cpp/dolfinx/fem/traits.h index 3a8d521e1d..d6915de30a 100644 --- a/cpp/dolfinx/fem/traits.h +++ b/cpp/dolfinx/fem/traits.h @@ -25,10 +25,9 @@ concept DofTransformKernel /// /// Kernel functions that can be passed to an assembler for execution /// must satisfy this concept. -template -concept FEkernel - = std::is_invocable_v*, - const int*, const std::uint8_t*, void*>; +template > +concept FEkernel = std::is_invocable_v; /// @brief Concept for mdspan of rank 1 or 2. template diff --git a/cpp/dolfinx/fem/utils.h b/cpp/dolfinx/fem/utils.h index 9c944459be..3cd60029c0 100644 --- a/cpp/dolfinx/fem/utils.h +++ b/cpp/dolfinx/fem/utils.h @@ -568,7 +568,7 @@ Form create_form_factory( active_coeffs.push_back(j); } - impl::kernel_t k = impl::extract_kernel(integral); + impl::kernel_t k = impl::extract_kernel(integral); if (!k) { throw std::runtime_error( @@ -648,7 +648,7 @@ Form create_form_factory( active_coeffs.push_back(j); } - impl::kernel_t k = impl::extract_kernel(integral); + impl::kernel_t k = impl::extract_kernel(integral); assert(k); // Build list of entities to assembler over @@ -774,7 +774,7 @@ Form create_form_factory( active_coeffs.push_back(j); } - impl::kernel_t k = impl::extract_kernel(integral); + impl::kernel_t k = impl::extract_kernel(integral); // Build list of entities to assembler over auto e_to_c = topology->connectivity(dim, tdim); @@ -971,28 +971,24 @@ Expression create_expression( static_cast(e.entity_dimension)}; std::vector value_shape(e.value_shape, e.value_shape + e.num_components); - std::function*, - const int*, const std::uint8_t*, void*)> + + static_assert(std::is_same_v>, + "UFCx kernels require geometry type U == scalar_value_t."); + + using kptr_t = void (*)(T*, const T*, const T*, const U*, const int*, + const std::uint8_t*, void*); + std::function tabulate_tensor = nullptr; if constexpr (std::is_same_v) - tabulate_tensor = e.tabulate_tensor_float32; -#ifndef DOLFINX_NO_STDC_COMPLEX_KERNELS - else if constexpr (std::is_same_v>) - { - tabulate_tensor = reinterpret_cast*, const int*, - const unsigned char*, void*)>(e.tabulate_tensor_complex64); - } -#endif // DOLFINX_NO_STDC_COMPLEX_KERNELS + tabulate_tensor = reinterpret_cast(e.tabulate_tensor_float32); else if constexpr (std::is_same_v) - tabulate_tensor = e.tabulate_tensor_float64; + tabulate_tensor = reinterpret_cast(e.tabulate_tensor_float64); #ifndef DOLFINX_NO_STDC_COMPLEX_KERNELS + else if constexpr (std::is_same_v>) + tabulate_tensor = reinterpret_cast(e.tabulate_tensor_complex64); else if constexpr (std::is_same_v>) - { - tabulate_tensor = reinterpret_cast*, const int*, - const unsigned char*, void*)>(e.tabulate_tensor_complex128); - } + tabulate_tensor = reinterpret_cast(e.tabulate_tensor_complex128); #endif // DOLFINX_NO_STDC_COMPLEX_KERNELS else throw std::runtime_error("Type not supported."); diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h index 3af3ec35ff..8059d424ab 100644 --- a/cpp/dolfinx/la/MatrixCSR.h +++ b/cpp/dolfinx/la/MatrixCSR.h @@ -352,7 +352,7 @@ class MatrixCSR { const std::size_t nrows = num_all_rows(); const std::size_t ncols = _index_maps[1]->size_global(); - std::vector A(nrows * ncols * _bs[0] * _bs[1], 0.0); + std::vector A(nrows * ncols * _bs[0] * _bs[1], value_type(0)); for (std::size_t r = 0; r < nrows; ++r) { for (std::int32_t j = _row_ptr[r]; j < _row_ptr[r + 1]; ++j) diff --git a/python/dolfinx/fem/bcs.py b/python/dolfinx/fem/bcs.py index 462a6f0521..81884699bf 100644 --- a/python/dolfinx/fem/bcs.py +++ b/python/dolfinx/fem/bcs.py @@ -13,7 +13,7 @@ from __future__ import annotations from collections.abc import Callable, Iterable -from typing import Generic +from typing import ClassVar, Generic import numpy as np import numpy.typing as npt @@ -98,6 +98,14 @@ class DirichletBC(Generic[Scalar]): The conditions are imposed on a linear system. """ + # Matched-precision built-ins (geometry == real scalar part). Public: + # extend with additional (scalar, geometry) dtype pairs as needed. + cpp_types: ClassVar[dict] = { + (np.dtype(np.float32), np.dtype(np.float32)): _cpp.fem.DirichletBC_float32, + (np.dtype(np.float64), np.dtype(np.float64)): _cpp.fem.DirichletBC_float64, + (np.dtype(np.complex64), np.dtype(np.float32)): _cpp.fem.DirichletBC_complex64, + (np.dtype(np.complex128), np.dtype(np.float64)): _cpp.fem.DirichletBC_complex128, + } _cpp_object: ( _cpp.fem.DirichletBC_complex64 | _cpp.fem.DirichletBC_complex128 @@ -200,19 +208,19 @@ def dirichletbc( try: dtype = value.dtype - if np.issubdtype(dtype, np.float32): - bctype = _cpp.fem.DirichletBC_float32 - elif np.issubdtype(dtype, np.float64): - bctype = _cpp.fem.DirichletBC_float64 - elif np.issubdtype(dtype, np.complex64): - bctype = _cpp.fem.DirichletBC_complex64 - elif np.issubdtype(dtype, np.complex128): - bctype = _cpp.fem.DirichletBC_complex128 - else: - raise NotImplementedError(f"Type {value.dtype} not supported.") except AttributeError as err: raise AttributeError("Boundary condition value must have a dtype attribute.") from err + # Geometry type is the mesh geometry type of the function space (or the + # value's space), defaulting to matched precision when neither has one. + if V is not None: + geometry_dtype = np.dtype(V.mesh.geometry.x.dtype) + elif isinstance(value, Function): + geometry_dtype = np.dtype(value.function_space.mesh.geometry.x.dtype) + else: + geometry_dtype = np.dtype(dtype).type(0).real.dtype + bctype = DirichletBC.cpp_types[dtype, geometry_dtype] + # Unwrap value object, if required if isinstance(value, np.ndarray): _value = value diff --git a/python/dolfinx/fem/element.py b/python/dolfinx/fem/element.py index 1d82d341af..c7313b00b9 100644 --- a/python/dolfinx/fem/element.py +++ b/python/dolfinx/fem/element.py @@ -6,7 +6,7 @@ """Finite elements.""" from functools import singledispatch -from typing import Generic +from typing import ClassVar, Generic import numpy as np import numpy.typing as npt @@ -20,6 +20,12 @@ class CoordinateElement(Generic[Real]): """Coordinate element describing the geometry map for mesh cells.""" + # Built-in geometry types. Public: extend with additional geometry + # dtypes as needed. + cpp_types: ClassVar[dict] = { + np.dtype(np.float32): _cpp.fem.CoordinateElement_float32, + np.dtype(np.float64): _cpp.fem.CoordinateElement_float64, + } _cpp_object: _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64 def __init__( @@ -144,12 +150,8 @@ def coordinate_element( Returns: A coordinate element. """ - if np.issubdtype(dtype, np.float32): - return CoordinateElement(_cpp.fem.CoordinateElement_float32(celltype, degree, variant)) - elif np.issubdtype(dtype, np.float64): - return CoordinateElement(_cpp.fem.CoordinateElement_float64(celltype, degree, variant)) - else: - raise RuntimeError("Unsupported dtype.") + cpp_type = CoordinateElement.cpp_types[np.dtype(dtype)] + return CoordinateElement(cpp_type(celltype, degree, variant)) @coordinate_element.register(basix.finite_element.FiniteElement) @@ -164,15 +166,19 @@ def _(e: basix.finite_element.FiniteElement) -> CoordinateElement: Returns: A coordinate element. """ - try: - return CoordinateElement(_cpp.fem.CoordinateElement_float32(e._e)) - except TypeError: - return CoordinateElement(_cpp.fem.CoordinateElement_float64(e._e)) + cpp_type = CoordinateElement.cpp_types[e.dtype] + return CoordinateElement(cpp_type(e._e)) class FiniteElement(Generic[Real]): """A finite element.""" + # Built-in geometry types. Public: extend with additional geometry + # dtypes as needed. + cpp_types: ClassVar[dict] = { + np.dtype(np.float32): _cpp.fem.FiniteElement_float32, + np.dtype(np.float64): _cpp.fem.FiniteElement_float64, + } _cpp_object: _cpp.fem.FiniteElement_float32 | _cpp.fem.FiniteElement_float64 def __init__( @@ -346,12 +352,7 @@ def finiteelement( the selected element. FiniteElement_dtype: Geometry type of the element. """ - if np.issubdtype(FiniteElement_dtype, np.float32): - CppElement = _cpp.fem.FiniteElement_float32 - elif np.issubdtype(FiniteElement_dtype, np.float64): - CppElement = _cpp.fem.FiniteElement_float64 - else: - raise ValueError(f"Unsupported dtype: {FiniteElement_dtype}") + CppElement = FiniteElement.cpp_types[np.dtype(FiniteElement_dtype)] if ufl_e.is_mixed: elements = [ diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 73a41375a3..a54a698737 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -37,6 +37,14 @@ class Form(typing.Generic[Scalar]): """A finite element form.""" + # Matched-precision built-ins (geometry == real scalar part). Public: + # extend with additional (scalar, geometry) dtype pairs as needed. + cpp_types: typing.ClassVar[dict] = { + (np.dtype(np.float32), np.dtype(np.float32)): _cpp.fem.Form_float32, + (np.dtype(np.float64), np.dtype(np.float64)): _cpp.fem.Form_float64, + (np.dtype(np.complex64), np.dtype(np.float32)): _cpp.fem.Form_complex64, + (np.dtype(np.complex128), np.dtype(np.float64)): _cpp.fem.Form_complex128, + } _cpp_object: ( _cpp.fem.Form_complex64 | _cpp.fem.Form_complex128 @@ -222,16 +230,9 @@ def form_cpp_class( This function is for advanced usage, typically when writing custom kernels using Numba or C. """ - if np.issubdtype(dtype, np.float32): - return _cpp.fem.Form_float32 - elif np.issubdtype(dtype, np.float64): - return _cpp.fem.Form_float64 - elif np.issubdtype(dtype, np.complex64): - return _cpp.fem.Form_complex64 - elif np.issubdtype(dtype, np.complex128): - return _cpp.fem.Form_complex128 - else: - raise NotImplementedError(f"Type {dtype} not supported.") + # Geometry is the real part of the scalar type (matched precision). + scalar_dtype = np.dtype(dtype) + return Form.cpp_types[scalar_dtype, scalar_dtype.type(0).real.dtype] _ufl_to_dolfinx_domain = { @@ -280,7 +281,6 @@ def mixed_topology_form( form_compiler_options = dict() form_compiler_options["scalar_type"] = dtype - ftype = form_cpp_class(dtype) # Extract subdomain data from UFL form sd = next(iter(forms)).subdomain_data() @@ -295,6 +295,9 @@ def mixed_topology_form( raise RuntimeError("Expecting to find a Mesh in the form.") comm = mesh.comm if jit_comm is None else jit_comm + # Geometry type is fixed by the mesh. + ftype = Form.cpp_types[np.dtype(dtype), mesh.geometry.x.dtype] + ufcx_forms = [] modules = [] codes = [] @@ -364,7 +367,6 @@ def form( form_compiler_options = dict() form_compiler_options["scalar_type"] = dtype - ftype = form_cpp_class(dtype) def _form(form): """Compile a single UFL form.""" @@ -381,6 +383,9 @@ def _form(form): raise RuntimeError("Expecting to find a Mesh in the form.") comm = msh.comm if jit_comm is None else jit_comm + # Geometry type is fixed by the mesh. + ftype = Form.cpp_types[np.dtype(dtype), msh.geometry.x.dtype] + ufcx_form, module, code = jit.ffcx_jit( comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options ) @@ -443,6 +448,9 @@ def _zero_form(form): assert len(V) > 0 msh = V[0].mesh + # Geometry type is fixed by the mesh. + ftype = Form.cpp_types[np.dtype(dtype), msh.geometry.x.dtype] + f = ftype( spaces=V, integrals={}, diff --git a/python/dolfinx/fem/function.py b/python/dolfinx/fem/function.py index a2b917cdcf..83ae7fc6b9 100644 --- a/python/dolfinx/fem/function.py +++ b/python/dolfinx/fem/function.py @@ -35,6 +35,14 @@ class Constant(ufl.Constant, Generic[Scalar]): """A constant with respect to a domain.""" + # Built-in scalar types. Public: extend with additional scalar dtypes + # as needed. + cpp_types: typing.ClassVar[dict] = { + np.dtype(np.float32): _cpp.fem.Constant_float32, + np.dtype(np.float64): _cpp.fem.Constant_float64, + np.dtype(np.complex64): _cpp.fem.Constant_complex64, + np.dtype(np.complex128): _cpp.fem.Constant_complex128, + } _cpp_object: ( _cpp.fem.Constant_complex64 | _cpp.fem.Constant_complex128 @@ -56,18 +64,10 @@ def __init__( c = np.asarray(c) super().__init__(domain, c.shape) try: - if np.issubdtype(c.dtype, np.complex64): - self._cpp_object = _cpp.fem.Constant_complex64(c) - elif np.issubdtype(c.dtype, np.complex128): - self._cpp_object = _cpp.fem.Constant_complex128(c) - elif np.issubdtype(c.dtype, np.float32): - self._cpp_object = _cpp.fem.Constant_float32(c) - elif np.issubdtype(c.dtype, np.float64): - self._cpp_object = _cpp.fem.Constant_float64(c) - else: - raise RuntimeError("Unsupported dtype") + cpp_type = Constant.cpp_types[c.dtype] except AttributeError as err: raise AttributeError("Constant value must have a dtype attribute.") from err + self._cpp_object = cpp_type(c) @property def value(self): @@ -112,6 +112,15 @@ class Expression(Generic[Scalar]): """ + # Matched-precision built-ins (geometry == real scalar part). An + # ``Expression`` is built through a factory, not a class constructor. + # Public: extend with additional (scalar, geometry) dtype pairs. + cpp_types: typing.ClassVar[dict] = { + (np.dtype(np.float32), np.dtype(np.float32)): _cpp.fem.create_expression_float32, + (np.dtype(np.float64), np.dtype(np.float64)): _cpp.fem.create_expression_float64, + (np.dtype(np.complex64), np.dtype(np.float32)): _cpp.fem.create_expression_complex64, + (np.dtype(np.complex128), np.dtype(np.float64)): _cpp.fem.create_expression_complex128, + } _ufl_expression: ufl.core.expr.Expr _argument_space: FunctionSpace | None _cpp_object: ( @@ -205,17 +214,15 @@ def __init__( else: raise RuntimeError("Expressions with more that one Argument not allowed.") - def _create_expression(dtype): - if np.issubdtype(dtype, np.float32): - return _cpp.fem.create_expression_float32 - elif np.issubdtype(dtype, np.float64): - return _cpp.fem.create_expression_float64 - elif np.issubdtype(dtype, np.complex64): - return _cpp.fem.create_expression_complex64 - elif np.issubdtype(dtype, np.complex128): - return _cpp.fem.create_expression_complex128 - else: - raise NotImplementedError(f"Type {dtype} not supported.") + # Geometry type is fixed by the expression's mesh. + expr_domains = ufl.domain.extract_domains(e) + if len(expr_domains) > 0: + expr_domain = expr_domains[0] + assert isinstance(expr_domain, ufl.Mesh) + geometry_dtype = expr_domain.ufl_cargo().geometry.x.dtype + else: + geometry_dtype = np.dtype(dtype).type(0).real.dtype + create_expression = Expression.cpp_types[np.dtype(dtype), geometry_dtype] _entity_maps = ( [entity_map._cpp_object for entity_map in entity_maps] @@ -223,7 +230,7 @@ def _create_expression(dtype): else [] ) ffi = module.ffi - self._cpp_object = _create_expression(dtype)( + self._cpp_object = create_expression( ffi.cast("uintptr_t", ffi.addressof(self._ufcx_expression)), coeffs, constants, @@ -341,6 +348,14 @@ class Function(ufl.Coefficient, Generic[Scalar]): """ + # Matched-precision built-ins (geometry == real scalar part). Public: + # extend with additional (scalar, geometry) dtype pairs. + cpp_types: typing.ClassVar[dict] = { + (np.dtype(np.float32), np.dtype(np.float32)): _cpp.fem.Function_float32, + (np.dtype(np.float64), np.dtype(np.float64)): _cpp.fem.Function_float64, + (np.dtype(np.complex64), np.dtype(np.float32)): _cpp.fem.Function_complex64, + (np.dtype(np.complex128), np.dtype(np.float64)): _cpp.fem.Function_complex128, + } _cpp_object: ( _cpp.fem.Function_complex64 | _cpp.fem.Function_complex128 @@ -376,27 +391,14 @@ def __init__( if dtype is None: dtype = default_scalar_type - assert np.issubdtype(V.element.dtype, np.dtype(dtype).type(0).real.dtype), ( - "Incompatible FunctionSpace dtype and requested dtype." - ) - - # Create cpp Function - def functiontype(dtype): - if np.issubdtype(dtype, np.float32): - return _cpp.fem.Function_float32 - elif np.issubdtype(dtype, np.float64): - return _cpp.fem.Function_float64 - elif np.issubdtype(dtype, np.complex64): - return _cpp.fem.Function_complex64 - elif np.issubdtype(dtype, np.complex128): - return _cpp.fem.Function_complex128 - else: - raise NotImplementedError(f"Type {dtype} not supported.") - + # Scalar type (dtype) is independent of the geometry type, which + # is fixed by the mesh. + geometry_dtype = V.mesh.geometry.x.dtype + cpp_type = Function.cpp_types[np.dtype(dtype), geometry_dtype] if x is not None: - self._cpp_object = functiontype(dtype)(V._cpp_object, x._cpp_object) # type: ignore + self._cpp_object = cpp_type(V._cpp_object, x._cpp_object) # type: ignore else: - self._cpp_object = functiontype(dtype)(V._cpp_object) # type: ignore + self._cpp_object = cpp_type(V._cpp_object) # type: ignore # Initialize the ufl.FunctionSpace super().__init__(V.ufl_function_space()) @@ -699,11 +701,9 @@ def functionspace( "Mesh and element dtype are not compatible." ) - # Initialize the cpp.FunctionSpace - try: - cppV = _cpp.fem.FunctionSpace_float64(mesh._cpp_object, element._cpp_object, cpp_dofmap) # type: ignore - except TypeError: - cppV = _cpp.fem.FunctionSpace_float32(mesh._cpp_object, element._cpp_object, cpp_dofmap) # type: ignore + # Initialize the cpp.FunctionSpace. Geometry type is fixed by the mesh. + cpp_type = FunctionSpace.cpp_types[dtype] + cppV = cpp_type(mesh._cpp_object, element._cpp_object, cpp_dofmap) # type: ignore return FunctionSpace(mesh, ufl_e, cppV) @@ -711,6 +711,12 @@ def functionspace( class FunctionSpace(ufl.FunctionSpace, Generic[Real]): """A space on which Functions (fields) can be defined.""" + # Built-in geometry types. Public: extend with additional geometry + # dtypes as needed. + cpp_types: typing.ClassVar[dict] = { + np.dtype(np.float32): _cpp.fem.FunctionSpace_float32, + np.dtype(np.float64): _cpp.fem.FunctionSpace_float64, + } _cpp_object: _cpp.fem.FunctionSpace_float32 | _cpp.fem.FunctionSpace_float64 _mesh: Mesh[Real] @@ -759,14 +765,9 @@ def clone(self) -> FunctionSpace[Real]: Returns: A new function space that shares data """ # noqa: D301 - try: - Vcpp = _cpp.fem.FunctionSpace_float64( - self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap - ) # type: ignore - except TypeError: - Vcpp = _cpp.fem.FunctionSpace_float32( - self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap - ) # type: ignore + # Geometry type is fixed by the mesh. + cpp_type = FunctionSpace.cpp_types[self._mesh.geometry.x.dtype] + Vcpp = cpp_type(self._cpp_object.mesh, self._cpp_object.element, self._cpp_object.dofmap) # type: ignore return FunctionSpace(self._mesh, self.ufl_element(), Vcpp) @property diff --git a/python/dolfinx/wrappers/dolfinx_wrappers/fem.h b/python/dolfinx/wrappers/dolfinx_wrappers/fem.h index 94cfdbb94a..189fabe896 100644 --- a/python/dolfinx/wrappers/dolfinx_wrappers/fem.h +++ b/python/dolfinx/wrappers/dolfinx_wrappers/fem.h @@ -58,17 +58,6 @@ namespace dolfinx_wrappers namespace nb = nanobind; namespace md = MDSPAN_IMPL_STANDARD_NAMESPACE; -template -struct geom_type -{ - typedef T value_type; -}; -template -struct geom_type> -{ - typedef typename T::value_type value_type; -}; - // Copy a container of pointers to a // std::vector> auto ptr_to_ref_wrapper_vec(auto& x) @@ -344,12 +333,41 @@ void declare_function_space(nb::module_& m, std::string type) } } -// Declare DirichletBC objects for type T template -void declare_objects(nb::module_& m, std::string type) +void declare_constant(nb::module_& m, std::string type) { - using U = typename dolfinx::scalar_value_t; + // dolfinx::fem::Constant + std::string pyclass_name_constant = std::string("Constant_") + type; + nb::class_>( + m, pyclass_name_constant.c_str(), + "Value constant with respect to integration domain") + .def( + "__init__", + [](dolfinx::fem::Constant* cp, + nb::ndarray c) + { + std::vector shape(c.shape_ptr(), + c.shape_ptr() + c.ndim()); + new (cp) + dolfinx::fem::Constant(std::span(c.data(), c.size()), shape); + }, + nb::arg("c").noconvert(), "Create a constant from a value array") + .def_prop_ro("dtype", [](const dolfinx::fem::Constant&) + { return dolfinx_wrappers::numpy_dtype_v; }) + .def_prop_ro( + "value", + [](dolfinx::fem::Constant& self) + { + return nb::ndarray( + self.value.data(), self.shape.size(), self.shape.data()); + }, + nb::rv_policy::reference_internal); +} +// Declare DirichletBC objects for type T +template > +void declare_objects(nb::module_& m, std::string type) +{ // dolfinx::fem::DirichletBC std::string pyclass_name = std::string("DirichletBC_") + type; nb::class_> dirichletbc( @@ -653,33 +671,6 @@ void declare_objects(nb::module_& m, std::string type) .def_prop_ro("function_space", &dolfinx::fem::Function::function_space); - // dolfinx::fem::Constant - std::string pyclass_name_constant = std::string("Constant_") + type; - nb::class_>( - m, pyclass_name_constant.c_str(), - "Value constant with respect to integration domain") - .def( - "__init__", - [](dolfinx::fem::Constant* cp, - nb::ndarray c) - { - std::vector shape(c.shape_ptr(), - c.shape_ptr() + c.ndim()); - new (cp) - dolfinx::fem::Constant(std::span(c.data(), c.size()), shape); - }, - nb::arg("c").noconvert(), "Create a constant from a value array") - .def_prop_ro("dtype", [](const dolfinx::fem::Constant&) - { return dolfinx_wrappers::numpy_dtype_v; }) - .def_prop_ro( - "value", - [](dolfinx::fem::Constant& self) - { - return nb::ndarray( - self.value.data(), self.shape.size(), self.shape.data()); - }, - nb::rv_policy::reference_internal); - // dolfinx::fem::Expression std::string pyclass_name_expr = std::string("Expression_") + type; nb::class_>(m, pyclass_name_expr.c_str(), @@ -700,9 +691,8 @@ void declare_objects(nb::module_& m, std::string type) argument_space) { auto tabulate_expression_ptr - = (void (*)(T*, const T*, const T*, - const typename geom_type::value_type*, - const int*, const std::uint8_t*, void*))fn_addr; + = (void (*)(T*, const T*, const T*, const U*, const int*, + const std::uint8_t*, void*))fn_addr; new (ex) dolfinx::fem::Expression( coefficients, constants, std::span(X.data(), X.size()), {X.shape(0), X.shape(1)}, tabulate_expression_ptr, value_shape, @@ -743,22 +733,26 @@ void declare_objects(nb::module_& m, std::string type) const std::vector& entity_maps, std::shared_ptr> argument_space) { - const ufcx_expression* p - = reinterpret_cast(expression); - return dolfinx::fem::create_expression( - *p, coefficients, constants, ptr_to_ref_wrapper_vec(entity_maps), - argument_space); + if constexpr (std::is_same_v>) + { + const ufcx_expression* p + = reinterpret_cast(expression); + return dolfinx::fem::create_expression( + *p, coefficients, constants, ptr_to_ref_wrapper_vec(entity_maps), + argument_space); + } + else + throw std::runtime_error( + "create_expression requires geometry type scalar_value_t."); }, nb::arg("expression"), nb::arg("coefficients"), nb::arg("constants"), nb::arg("entity_maps"), nb::arg("argument_space").none(), "Create Expression from a pointer to ufc_form."); } -template +template > void declare_form(nb::module_& m, std::string type) { - using U = typename dolfinx::scalar_value_t; - // dolfinx::fem::Form std::string pyclass_name_form = std::string("Form_") + type; nb::class_>(m, pyclass_name_form.c_str(), @@ -784,7 +778,7 @@ void declare_form(nb::module_& m, std::string type) std::shared_ptr> mesh) { std::map, - dolfinx::fem::integral_data> + dolfinx::fem::integral_data> _integrals; // Loop over kernel for each entity type @@ -793,9 +787,8 @@ void declare_form(nb::module_& m, std::string type) for (auto& [id, ptr, e, c] : kernels) { auto kn_ptr - = (void (*)(T*, const T*, const T*, - const typename geom_type::value_type*, - const int*, const std::uint8_t*, void*))ptr; + = (void (*)(T*, const T*, const T*, const U*, const int*, + const std::uint8_t*, void*))ptr; _integrals.insert( {{type, id, 0}, {kn_ptr, @@ -843,14 +836,20 @@ void declare_form(nb::module_& m, std::string type) sd.insert({itg, std::move(x)}); } - std::vector> ps; - ps.reserve(forms.size()); - for (auto form : forms) - ps.push_back(*(reinterpret_cast(form))); - new (fp) - dolfinx::fem::Form(dolfinx::fem::create_form_factory( - ps, spaces, coefficients, constants, sd, - ptr_to_ref_wrapper_vec(entity_maps), mesh)); + if constexpr (std::is_same_v>) + { + std::vector> ps; + ps.reserve(forms.size()); + for (auto form : forms) + ps.push_back(*(reinterpret_cast(form))); + new (fp) + dolfinx::fem::Form(dolfinx::fem::create_form_factory( + ps, spaces, coefficients, constants, sd, + ptr_to_ref_wrapper_vec(entity_maps), mesh)); + } + else + throw std::runtime_error( + "create_form requires geometry type scalar_value_t."); }, nb::arg("form"), nb::arg("spaces"), nb::arg("coefficients"), nb::arg("constants"), nb::arg("subdomains"), nb::arg("entity_maps"), @@ -924,9 +923,15 @@ void declare_form(nb::module_& m, std::string type) sd.insert({itg, std::move(x)}); } - ufcx_form* p = reinterpret_cast(form); - return dolfinx::fem::create_form_factory( - {*p}, spaces, coefficients, constants, sd, {}, std::move(mesh)); + if constexpr (std::is_same_v>) + { + ufcx_form* p = reinterpret_cast(form); + return dolfinx::fem::create_form_factory( + {*p}, spaces, coefficients, constants, sd, {}, std::move(mesh)); + } + else + throw std::runtime_error( + "create_form requires geometry type scalar_value_t."); }, nb::arg("form"), nb::arg("spaces"), nb::arg("coefficients"), nb::arg("constants"), nb::arg("subdomains"), nb::arg("mesh"), @@ -962,10 +967,16 @@ void declare_form(nb::module_& m, std::string type) sd.insert({itg, std::move(x)}); } - ufcx_form* p = reinterpret_cast(form); - return dolfinx::fem::create_form( - *p, spaces, coefficients, constants, sd, - ptr_to_ref_wrapper_vec(entity_maps), std::move(mesh)); + if constexpr (std::is_same_v>) + { + ufcx_form* p = reinterpret_cast(form); + return dolfinx::fem::create_form( + *p, spaces, coefficients, constants, sd, + ptr_to_ref_wrapper_vec(entity_maps), std::move(mesh)); + } + else + throw std::runtime_error( + "create_form requires geometry type scalar_value_t."); }, // NOLINTEND(performance-no-int-to-ptr) nb::arg("form"), nb::arg("spaces"), nb::arg("coefficients"), diff --git a/python/dolfinx/wrappers/fem.cpp b/python/dolfinx/wrappers/fem.cpp index 8cbc07c99a..6323eead50 100644 --- a/python/dolfinx/wrappers/fem.cpp +++ b/python/dolfinx/wrappers/fem.cpp @@ -160,6 +160,11 @@ void fem(nb::module_& m) .value("vertex", dolfinx::fem::IntegralType::vertex, "vertex integral") .value("ridge", dolfinx::fem::IntegralType::ridge, "ridge integral"); + declare_constant(m, "float32"); + declare_constant(m, "float64"); + declare_constant>(m, "complex64"); + declare_constant>(m, "complex128"); + declare_objects(m, "float32"); declare_objects(m, "float64"); declare_objects>(m, "complex64");