diff --git a/cpp/demo/hyperelasticity/main.cpp b/cpp/demo/hyperelasticity/main.cpp index b202498844..35da5a1889 100644 --- a/cpp/demo/hyperelasticity/main.cpp +++ b/cpp/demo/hyperelasticity/main.cpp @@ -264,7 +264,7 @@ int main(int argc, char* argv[]) auto sigma = fem::Function(S); sigma.name = "cauchy_stress"; - sigma.interpolate(sigma_expression); + sigma.interpolate(sigma_expression, *mesh); // Save solution in VTK format io::VTKFile file_u(mesh->comm(), "u.pvd", "w"); diff --git a/cpp/demo/interpolation_different_meshes/main.cpp b/cpp/demo/interpolation_different_meshes/main.cpp index 5dbcb4d895..09651b6da4 100644 --- a/cpp/demo/interpolation_different_meshes/main.cpp +++ b/cpp/demo/interpolation_different_meshes/main.cpp @@ -74,7 +74,8 @@ int main(int argc, char* argv[]) *u_hex->function_space()->mesh(), *u_hex->function_space()->element(), *u_tet->function_space()->mesh(), padding); - u_hex->interpolate(*u_tet, nmm_interpolation_data); + constexpr std::span cell_map; + u_hex->interpolate(*u_tet, cell_map, nmm_interpolation_data); #ifdef HAS_ADIOS2 io::VTXWriter write_tet(mesh_tet->comm(), "u_tet.bp", {u_tet}); diff --git a/cpp/dolfinx/fem/Function.h b/cpp/dolfinx/fem/Function.h index 1927da921a..050b4892db 100644 --- a/cpp/dolfinx/fem/Function.h +++ b/cpp/dolfinx/fem/Function.h @@ -153,28 +153,34 @@ class Function /// @brief Interpolate a provided Function. /// @param[in] v The function to be interpolated /// @param[in] cells The cells to interpolate on + /// @param[in] cell_map A map from cells in the mesh associated with `this` + /// function to cells in mesh associated with `v` /// @param[in] nmm_interpolation_data Auxiliary data to interpolate on /// nonmatching meshes. This data can be generated with /// generate_nonmatching_meshes_interpolation_data (optional). void interpolate( const Function& v, std::span cells, + std::span cell_map, const std::tuple, std::span, std::span, std::span>& nmm_interpolation_data = {}) { - fem::interpolate(*this, v, cells, nmm_interpolation_data); + fem::interpolate(*this, v, cells, cell_map, nmm_interpolation_data); } /// @brief Interpolate a provided Function. /// @param[in] v The function to be interpolated + /// @param[in] cell_map Map from cells in self to cell indices in `v` /// @param[in] nmm_interpolation_data Auxiliary data to interpolate on /// nonmatching meshes. This data can be generated with /// generate_nonmatching_meshes_interpolation_data (optional). void interpolate( const Function& v, + std::span cell_map = std::span() + = {}, const std::tuple, std::span, std::span, @@ -184,12 +190,12 @@ class Function assert(_function_space); assert(_function_space->mesh()); int tdim = _function_space->mesh()->topology()->dim(); - auto cell_map = _function_space->mesh()->topology()->index_map(tdim); - assert(cell_map); - std::int32_t num_cells = cell_map->size_local() + cell_map->num_ghosts(); + auto cell_imap = _function_space->mesh()->topology()->index_map(tdim); + assert(cell_imap); + std::int32_t num_cells = cell_imap->size_local() + cell_imap->num_ghosts(); std::vector cells(num_cells, 0); std::iota(cells.begin(), cells.end(), 0); - interpolate(v, cells, nmm_interpolation_data); + interpolate(v, cells, cell_map, nmm_interpolation_data); } /// Interpolate an expression function on a list of cells @@ -283,8 +289,13 @@ class Function /// `FiniteElement::interpolation_points()` for the element associated /// with `u`. /// @param[in] cells The cells to interpolate on + /// @param[in] expr_mesh The mesh to evaluate the expression on + /// @param[in] cell_map Map from `cells` to cells in expression void interpolate(const Expression& e, - std::span cells) + std::span cells, + const dolfinx::mesh::Mesh& expr_mesh, + std::span cell_map + = std::span()) { // Check that spaces are compatible assert(_function_space); @@ -331,8 +342,25 @@ class Function // Evaluate Expression at points assert(_function_space->mesh()); - e.eval(*_function_space->mesh(), cells, fdata, - {num_cells, num_points * value_size}); + + std::vector cells_expr; + cells_expr.reserve(num_cells); + // Get mesh and check if mesh is shared + if (auto mesh_v = _function_space->mesh(); + expr_mesh.topology() == mesh_v->topology()) + { + cells_expr.insert(cells_expr.end(), cells.begin(), cells.end()); + } + // If meshes are different and input mapping is given + else if (!cell_map.empty()) + { + std::transform(cells.begin(), cells.end(), std::back_inserter(cells_expr), + [&cell_map](std::int32_t c) { return cell_map[c]; }); + } + else + std::runtime_error("Meshes are different and no cell map is provided"); + + e.eval(expr_mesh, cells_expr, fdata, {num_cells, num_points * value_size}); // Reshape evaluated data to fit interpolate // Expression returns matrix of shape (num_cells, num_points * @@ -356,17 +384,27 @@ class Function /// Interpolate an Expression (based on UFL) on all cells /// @param[in] e The function to be interpolated - void interpolate(const Expression& e) + /// @param[in] expr_mesh Mesh the Expression `e` is defined on. + /// @param[in] cell_map Map from `cells` to cells in expression if + /// receiving function is defined on a different mesh than the expression + void interpolate(const Expression& e, + const dolfinx::mesh::Mesh& expr_mesh, + std::span cell_map + = std::span() = {}) { assert(_function_space); assert(_function_space->mesh()); const int tdim = _function_space->mesh()->topology()->dim(); - auto cell_map = _function_space->mesh()->topology()->index_map(tdim); - assert(cell_map); - std::int32_t num_cells = cell_map->size_local() + cell_map->num_ghosts(); + auto cell_imap = _function_space->mesh()->topology()->index_map(tdim); + assert(cell_imap); + std::int32_t num_cells = cell_imap->size_local() + cell_imap->num_ghosts(); std::vector cells(num_cells, 0); std::iota(cells.begin(), cells.end(), 0); - interpolate(e, cells); + + if (cell_map.size() == 0) + interpolate(e, cells, expr_mesh, cell_map); + else + interpolate(e, cells, expr_mesh, std::span(cells.data(), cells.size())); } /// @brief Evaluate the Function at points. diff --git a/cpp/dolfinx/fem/interpolate.h b/cpp/dolfinx/fem/interpolate.h index ffe21b6a0c..239c3384c6 100644 --- a/cpp/dolfinx/fem/interpolate.h +++ b/cpp/dolfinx/fem/interpolate.h @@ -317,38 +317,47 @@ void interpolation_apply(U&& Pi, V&& data, std::span coeffs, int bs) /// map. /// @param[out] u1 The function to interpolate to /// @param[in] u0 The function to interpolate from -/// @param[in] cells The cells to interpolate on +/// @param[in] cells1 The cells to interpolate on +/// @param[in] cells0 Equivalent cell in u0 for each cell in u1 /// @pre The functions `u1` and `u0` must share the same mesh and the /// elements must share the same basis function map. Neither is checked /// by the function. template void interpolate_same_map(Function& u1, const Function& u0, - std::span cells) + std::span cells1, + std::span cells0) { auto V0 = u0.function_space(); assert(V0); auto V1 = u1.function_space(); assert(V1); - auto mesh = V0->mesh(); - assert(mesh); + auto mesh0 = V0->mesh(); + assert(mesh0); + + auto mesh1 = V1->mesh(); + assert(mesh1); auto element0 = V0->element(); assert(element0); auto element1 = V1->element(); assert(element1); - const int tdim = mesh->topology()->dim(); - auto map = mesh->topology()->index_map(tdim); + assert(mesh0->topology()->dim()); + const int tdim = mesh0->topology()->dim(); + auto map = mesh0->topology()->index_map(tdim); assert(map); std::span u1_array = u1.x()->mutable_array(); std::span u0_array = u0.x()->array(); - std::span cell_info; + std::span cell_info0; + std::span cell_info1; if (element1->needs_dof_transformations() or element0->needs_dof_transformations()) { - mesh->topology_mutable()->create_entity_permutations(); - cell_info = std::span(mesh->topology()->get_cell_permutation_info()); + mesh0->topology_mutable()->create_entity_permutations(); + cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info()); + mesh1->topology_mutable()->create_entity_permutations(); + cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info()); } // Get dofmaps @@ -358,17 +367,13 @@ void interpolate_same_map(Function& u1, const Function& u0, // Get block sizes and dof transformation operators const int bs1 = dofmap1->bs(); const int bs0 = dofmap0->bs(); - - // Transform conforming dofs to reference dofs function auto apply_dof_transformation = element0->template dof_transformation_fn( doftransform::transpose, false); - - // Transform reference dofs to conforming dofs function auto apply_inverse_dof_transform = element1->template dof_transformation_fn( doftransform::inverse_transpose, false); - // Create working arrays + // Create working array std::vector local0(element0->space_dimension()); std::vector local1(element1->space_dimension()); @@ -378,27 +383,25 @@ void interpolate_same_map(Function& u1, const Function& u0, // Iterate over mesh and interpolate on each cell using X = typename dolfinx::scalar_value_type_t; - for (auto c : cells) + for (std::size_t c = 0; c < cells0.size(); c++) { // Pack and transform cell dofs to reference ordering - std::span dofs0 = dofmap0->cell_dofs(c); + std::span dofs0 = dofmap0->cell_dofs(cells0[c]); for (std::size_t i = 0; i < dofs0.size(); ++i) for (int k = 0; k < bs0; ++k) local0[bs0 * i + k] = u0_array[bs0 * dofs0[i] + k]; - apply_dof_transformation(local0, cell_info, c, 1); - // FIXME: Get compile-time ranges from Basix + apply_dof_transformation(local0, cell_info0, cells0[c], 1); + // FIXME: Get compile-time ranges from Basix // Apply interpolation operator std::fill(local1.begin(), local1.end(), 0); for (std::size_t i = 0; i < im_shape[0]; ++i) for (std::size_t j = 0; j < im_shape[1]; ++j) local1[i] += static_cast(i_m[im_shape[1] * i + j]) * local0[j]; - // Transform cells dofs back to conforming physical ordering and - // scatter - apply_inverse_dof_transform(local1, cell_info, c, 1); - std::span dofs1 = dofmap1->cell_dofs(c); + apply_inverse_dof_transform(local1, cell_info1, cells1[c], 1); + std::span dofs1 = dofmap1->cell_dofs(cells1[c]); for (std::size_t i = 0; i < dofs1.size(); ++i) for (int k = 0; k < bs1; ++k) u1_array[bs1 * dofs1[i] + k] = local1[bs1 * i + k]; @@ -412,37 +415,45 @@ void interpolate_same_map(Function& u1, const Function& u0, /// standard isoparametric mapping. /// @param[out] u1 The function to interpolate to /// @param[in] u0 The function to interpolate from -/// @param[in] cells The cells to interpolate on +/// @param[in] cells1 The cells to interpolate on +/// @param[in] cells0 Equivalent cell in u0 for each cell in u1 /// @pre The functions `u1` and `u0` must share the same mesh. This is /// not checked by the function. template void interpolate_nonmatching_maps(Function& u1, const Function& u0, - std::span cells) + std::span cells1, + std::span cells0) { // Get mesh auto V0 = u0.function_space(); assert(V0); - auto mesh = V0->mesh(); - assert(mesh); + auto mesh0 = V0->mesh(); + assert(mesh0); // Mesh dims - const int tdim = mesh->topology()->dim(); - const int gdim = mesh->geometry().dim(); + const int tdim = mesh0->topology()->dim(); + const int gdim = mesh0->geometry().dim(); // Get elements auto V1 = u1.function_space(); assert(V1); + auto mesh1 = V1->mesh(); + assert(mesh1); auto element0 = V0->element(); assert(element0); auto element1 = V1->element(); assert(element1); - std::span cell_info; + std::span cell_info0; + std::span cell_info1; + if (element1->needs_dof_transformations() or element0->needs_dof_transformations()) { - mesh->topology_mutable()->create_entity_permutations(); - cell_info = std::span(mesh->topology()->get_cell_permutation_info()); + mesh0->topology_mutable()->create_entity_permutations(); + cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info()); + mesh1->topology_mutable()->create_entity_permutations(); + cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info()); } // Get dofmaps @@ -454,14 +465,8 @@ void interpolate_nonmatching_maps(Function& u1, const Function& u0, // Get block sizes and dof transformation operators const int bs0 = element0->block_size(); const int bs1 = element1->block_size(); - - // Transform from reference basis function layout to conforming layout - // function auto apply_dof_transformation0 = element0->template dof_transformation_fn( doftransform::standard, false); - - // Transform from reference degree-of-freedom to conforming layout - // function auto apply_inverse_dof_transform1 = element1->template dof_transformation_fn( doftransform::inverse_transpose, false); @@ -471,10 +476,10 @@ void interpolate_nonmatching_maps(Function& u1, const Function& u0, const std::size_t value_size_ref0 = element0->reference_value_size() / bs0; const std::size_t value_size0 = V0->value_size() / bs0; - const CoordinateElement& cmap = mesh->geometry().cmap(); - auto x_dofmap = mesh->geometry().dofmap(); + const CoordinateElement& cmap = mesh0->geometry().cmap(); + auto x_dofmap = mesh0->geometry().dofmap(); const std::size_t num_dofs_g = cmap.dim(); - std::span x_g = mesh->geometry().x(); + std::span x_g = mesh0->geometry().x(); // Evaluate coordinate map basis at reference interpolation points const std::array phi_shape @@ -543,11 +548,11 @@ void interpolate_nonmatching_maps(Function& u1, const Function& u0, // Iterate over mesh and interpolate on each cell std::span array0 = u0.x()->array(); std::span array1 = u1.x()->mutable_array(); - for (auto c : cells) + for (std::size_t c = 0; c < cells0.size(); c++) { // Get cell geometry (coordinate dofs) auto x_dofs = MDSPAN_IMPL_STANDARD_NAMESPACE::submdspan( - x_dofmap, c, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent); + x_dofmap, cells0[c], MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent); for (std::size_t i = 0; i < num_dofs_g; ++i) { const int pos = 3 * x_dofs[i]; @@ -584,12 +589,10 @@ void interpolate_nonmatching_maps(Function& u1, const Function& u0, for (std::size_t p = 0; p < Xshape[0]; ++p) { - // Transform basis function data from reference layout to - // conforming layout apply_dof_transformation0( std::span(basis_reference0_b.data() + p * dim0 * value_size_ref0, dim0 * value_size_ref0), - cell_info, c, value_size_ref0); + cell_info0, cells0[c], value_size_ref0); } for (std::size_t i = 0; i < basis0.extent(0); ++i) @@ -611,7 +614,7 @@ void interpolate_nonmatching_maps(Function& u1, const Function& u0, // Copy expansion coefficients for v into local array const int dof_bs0 = dofmap0->bs(); - std::span dofs0 = dofmap0->cell_dofs(c); + std::span dofs0 = dofmap0->cell_dofs(cells0[c]); for (std::size_t i = 0; i < dofs0.size(); ++i) for (int k = 0; k < dof_bs0; ++k) coeffs0[dof_bs0 * i + k] = array0[dof_bs0 * dofs0[i] + k]; @@ -654,9 +657,7 @@ void interpolate_nonmatching_maps(Function& u1, const Function& u0, mapped_values0, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent, 0, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent); interpolation_apply(Pi_1, values, std::span(local1), bs1); - - // Transform from reference dof layout to conforming dof layoyut - apply_inverse_dof_transform1(local1, cell_info, c, 1); + apply_inverse_dof_transform1(local1, cell_info1, cells1[c], 1); // Copy local coefficients to the correct position in u dof array const int dof_bs1 = dofmap1->bs(); @@ -799,8 +800,6 @@ void interpolate(Function& u, std::span f, // Point evaluation element *and* the geometric map is the identity, // e.g. not Piola mapped - // Transform from reference degrees-of-freedom to conforming - // degrees-of-freedom function auto apply_inv_transpose_dof_transformation = element->template dof_transformation_fn( doftransform::inverse_transpose, true); @@ -845,8 +844,6 @@ void interpolate(Function& u, std::span f, const std::size_t num_interp_points = Pi.extent(1); assert(Pi.extent(0) == num_scalar_dofs); - // Transform reference degrees-of-freedom to conforming - // degrees-of-freedom functions auto apply_inv_transpose_dof_transformation = element->template dof_transformation_fn( doftransform::inverse_transpose, true); @@ -937,9 +934,9 @@ void interpolate(Function& u, std::span f, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent, 0); - // Transform from reference degrees-of-freedom to conforming - // degrees-of-freedom - auto apply_inverse_transpose_dof_transformation + const std::function, std::span, + std::int32_t, int)> + apply_inverse_transpose_dof_transformation = element->template dof_transformation_fn( doftransform::inverse_transpose); @@ -1110,111 +1107,127 @@ create_nonmatching_meshes_interpolation_data(const mesh::Mesh& mesh0, } /// @brief Interpolate from one finite element Function to another one. -/// @param[out] u The function to interpolate into -/// @param[in] v The function to be interpolated -/// @param[in] cells List of cell indices to interpolate on +/// @param[out] u1 The function to interpolate into +/// @param[in] u0 The function to be interpolated +/// @param[in] cells1 List of cell indices associated with the mesh of `u1` that +/// will be interpolated onto +/// @param[in] cell_map Mapping of cells in mesh associated with `u1` to cells +/// associated with `u0` /// @param[in] nmm_interpolation_data Auxiliary data to interpolate on /// nonmatching meshes. This data can be generated with /// create_nonmatching_meshes_interpolation_data (optional). template void interpolate( - Function& u, const Function& v, - std::span cells, + Function& u1, const Function& u0, + std::span cells1, + std::span cell_map, const std::tuple, std::span, std::span, std::span>& nmm_interpolation_data = {}) { - assert(u.function_space()); - assert(v.function_space()); - auto mesh = u.function_space()->mesh(); + assert(u1.function_space()); + assert(u0.function_space()); + auto mesh = u1.function_space()->mesh(); assert(mesh); auto cell_map0 = mesh->topology()->index_map(mesh->topology()->dim()); assert(cell_map0); std::size_t num_cells0 = cell_map0->size_local() + cell_map0->num_ghosts(); - if (u.function_space() == v.function_space() and cells.size() == num_cells0) + if (u1.function_space() == u0.function_space() + and cells1.size() == num_cells0) { // Same function spaces and on whole mesh - std::span u1_array = u.x()->mutable_array(); - std::span u0_array = v.x()->array(); + std::span u1_array = u1.x()->mutable_array(); + std::span u0_array = u0.x()->array(); std::copy(u0_array.begin(), u0_array.end(), u1_array.begin()); } else { + std::vector cells0; + cells0.reserve(cells1.size()); // Get mesh and check that functions share the same mesh - if (auto mesh_v = v.function_space()->mesh(); mesh != mesh_v) + if (auto mesh_v = u0.function_space()->mesh(); mesh == mesh_v) { - impl::interpolate_nonmatching_meshes(u, v, cells, nmm_interpolation_data); - return; + cells0.insert(cells0.end(), cells1.begin(), cells1.end()); } - else + // If meshes are different and input mapping is given + else if (cell_map.size() > 0) { + std::transform(cells1.begin(), cells1.end(), std::back_inserter(cells0), + [&cell_map](std::int32_t c) { return cell_map[c]; }); + } + // Non-matching meshes + if (cells0.empty()) + { + impl::interpolate_nonmatching_meshes(u1, u0, cells1, + nmm_interpolation_data); + return; + } - // Get elements and check value shape - auto fs0 = v.function_space(); - auto element0 = fs0->element(); - assert(element0); - auto fs1 = u.function_space(); - auto element1 = fs1->element(); - assert(element1); - if (fs0->value_shape().size() != fs1->value_shape().size() - or !std::equal(fs0->value_shape().begin(), fs0->value_shape().end(), - fs1->value_shape().begin())) - { - throw std::runtime_error( - "Interpolation: elements have different value dimensions"); - } + // Get elements and check value shape + auto fs0 = u0.function_space(); + auto element0 = fs0->element(); + assert(element0); + auto fs1 = u1.function_space(); + auto element1 = fs1->element(); + assert(element1); + if (fs0->value_shape().size() != fs1->value_shape().size() + or !std::equal(fs0->value_shape().begin(), fs0->value_shape().end(), + fs1->value_shape().begin())) + { + throw std::runtime_error( + "Interpolation: elements have different value dimensions"); + } - if (element1 == element0 or *element1 == *element0) - { - // Same element, different dofmaps (or just a subset of cells) + if (element1 == element0 or *element1 == *element0) + { + // Same element, different dofmaps (or just a subset of cells) - const int tdim = mesh->topology()->dim(); - auto cell_map = mesh->topology()->index_map(tdim); - assert(cell_map); + const int tdim = mesh->topology()->dim(); + auto cell_map1 = mesh->topology()->index_map(tdim); + assert(cell_map1); - assert(element1->block_size() == element0->block_size()); + assert(element1->block_size() == element0->block_size()); - // Get dofmaps - std::shared_ptr dofmap0 = v.function_space()->dofmap(); - assert(dofmap0); - std::shared_ptr dofmap1 = u.function_space()->dofmap(); - assert(dofmap1); + // Get dofmaps + std::shared_ptr dofmap0 = u0.function_space()->dofmap(); + assert(dofmap0); + std::shared_ptr dofmap1 = u1.function_space()->dofmap(); + assert(dofmap1); - std::span u1_array = u.x()->mutable_array(); - std::span u0_array = v.x()->array(); + std::span u1_array = u1.x()->mutable_array(); + std::span u0_array = u0.x()->array(); - // Iterate over mesh and interpolate on each cell - const int bs0 = dofmap0->bs(); - const int bs1 = dofmap1->bs(); - for (auto c : cells) + // Iterate over mesh and interpolate on each cell + const int bs0 = dofmap0->bs(); + const int bs1 = dofmap1->bs(); + for (std::size_t c = 0; c < cells1.size(); ++c) + { + std::span dofs0 = dofmap0->cell_dofs(cells0[c]); + std::span dofs1 = dofmap1->cell_dofs(cells1[c]); + assert(bs0 * dofs0.size() == bs1 * dofs1.size()); + for (std::size_t i = 0; i < dofs0.size(); ++i) { - std::span dofs0 = dofmap0->cell_dofs(c); - std::span dofs1 = dofmap1->cell_dofs(c); - assert(bs0 * dofs0.size() == bs1 * dofs1.size()); - for (std::size_t i = 0; i < dofs0.size(); ++i) + for (int k = 0; k < bs0; ++k) { - for (int k = 0; k < bs0; ++k) - { - int index = bs0 * i + k; - std::div_t dv1 = std::div(index, bs1); - u1_array[bs1 * dofs1[dv1.quot] + dv1.rem] - = u0_array[bs0 * dofs0[i] + k]; - } + int index = bs0 * i + k; + std::div_t dv1 = std::div(index, bs1); + u1_array[bs1 * dofs1[dv1.quot] + dv1.rem] + = u0_array[bs0 * dofs0[i] + k]; } } } - else if (element1->map_type() == element0->map_type()) - { - // Different elements, same basis function map type - impl::interpolate_same_map(u, v, cells); - } - else - { - // Different elements with different maps for basis functions - impl::interpolate_nonmatching_maps(u, v, cells); - } + } + else if (element1->map_type() == element0->map_type()) + { + // Different elements, same basis function map type + impl::interpolate_same_map(u1, u0, cells1, cells0); + } + else + { + // Different elements with different maps for basis functions + impl::interpolate_nonmatching_maps(u1, u0, cells1, cells0); } } } diff --git a/python/dolfinx/fem/function.py b/python/dolfinx/fem/function.py index 53680f23fc..25eecbbe7f 100644 --- a/python/dolfinx/fem/function.py +++ b/python/dolfinx/fem/function.py @@ -416,6 +416,8 @@ def interpolate( self, u: typing.Union[typing.Callable, Expression, Function], cells: typing.Optional[np.ndarray] = None, + cell_map: typing.Optional[np.ndarray] = None, + expr_mesh: typing.Optional[Mesh] = None, nmm_interpolation_data: typing.Optional[PointOwnershipData] = None, ) -> None: """Interpolate an expression @@ -424,6 +426,10 @@ def interpolate( u: The function, Expression or Function to interpolate. cells: The cells to interpolate over. If `None` then all cells are interpolated over. + cell_map: Mapping from `cells` to to cells in the mesh that `u` is defined over. + expr_mesh: If an Expression with coefficients or constants from another mesh + than the function is supplied, the mesh associated with this expression has + to be provided, along with `cell_map.` nmm_interpolation_data: Data needed to interpolate functions defined on other meshes """ if nmm_interpolation_data is None: @@ -435,6 +441,14 @@ def interpolate( dest_cells=np.empty(0, dtype=np.int32), ) + if cells is None: + mesh = self.function_space.mesh + map = mesh.topology.index_map(mesh.topology.dim) + cells = np.arange(map.size_local + map.num_ghosts, dtype=np.int32) + + if cell_map is None: + cell_map = np.empty(0, dtype=np.int32) + @singledispatch def _interpolate(u, cells: typing.Optional[np.ndarray] = None): """Interpolate a cpp.fem.Function""" @@ -443,7 +457,7 @@ def _interpolate(u, cells: typing.Optional[np.ndarray] = None): @_interpolate.register(Function) def _(u: Function, cells: typing.Optional[np.ndarray] = None): """Interpolate a fem.Function""" - self._cpp_object.interpolate(u._cpp_object, cells, nmm_interpolation_data) # type: ignore + self._cpp_object.interpolate(u._cpp_object, cells, cell_map, nmm_interpolation_data) # type: ignore @_interpolate.register(int) def _(u_ptr: int, cells: typing.Optional[np.ndarray] = None): @@ -452,13 +466,26 @@ def _(u_ptr: int, cells: typing.Optional[np.ndarray] = None): @_interpolate.register(Expression) def _(expr: Expression, cells: typing.Optional[np.ndarray] = None): - """Interpolate Expression for the set of cells""" - self._cpp_object.interpolate(expr._cpp_object, cells) # type: ignore - - if cells is None: - mesh = self.function_space.mesh - map = mesh.topology.index_map(mesh.topology.dim) - cells = np.arange(map.size_local + map.num_ghosts, dtype=np.int32) + """Interpolate Expression from a given mesh onto the set of cells + Args: + expr: Expression to interpolate + cells: The cells to interpolate over. If `None` then all + cells are interpolated over. + """ + assert cell_map is not None + if len(cell_map) == 0: + # If cell map is not provided create identity map + assert cells is not None + expr_cell_map = np.arange(len(cells), dtype=np.int32) + assert expr_mesh is None + mapping_mesh = self.function_space.mesh._cpp_object + else: + # If cell map is provided check that there is a mesh + # associated with the expression + expr_cell_map = cell_map + assert expr_mesh is not None + mapping_mesh = expr_mesh._cpp_object + self._cpp_object.interpolate(expr._cpp_object, cells, mapping_mesh, expr_cell_map) # type: ignore try: # u is a Function or Expression (or pointer to one) diff --git a/python/dolfinx/wrappers/fem.cpp b/python/dolfinx/wrappers/fem.cpp index 7e56313c97..1bb5bf6274 100644 --- a/python/dolfinx/wrappers/fem.cpp +++ b/python/dolfinx/wrappers/fem.cpp @@ -387,6 +387,8 @@ void declare_objects(nb::module_& m, const std::string& type) [](dolfinx::fem::Function& self, dolfinx::fem::Function& u, nb::ndarray, nb::c_contig> cells, + nb::ndarray, nb::c_contig> + cell_map, const std::tuple< nb::ndarray, nb::c_contig>, nb::ndarray, nb::c_contig>, @@ -410,9 +412,11 @@ void declare_objects(nb::module_& m, const std::string& type) std::get<3>(interpolation_data).data(), std::get<3>(interpolation_data).size())); self.interpolate(u, std::span(cells.data(), cells.size()), + std::span(cell_map.data(), cell_map.size()), _interpolation_data); }, - nb::arg("u"), nb::arg("cells"), nb::arg("nmm_interpolation_data"), + nb::arg("u"), nb::arg("cells"), nb::arg("cell_map"), + nb::arg("nmm_interpolation_data"), "Interpolate a finite element function") .def( "interpolate_ptr", @@ -449,10 +453,16 @@ void declare_objects(nb::module_& m, const std::string& type) "interpolate", [](dolfinx::fem::Function& self, const dolfinx::fem::Expression& expr, - nb::ndarray cells) - { self.interpolate(expr, std::span(cells.data(), cells.size())); }, - nb::arg("expr"), nb::arg("cells"), - "Interpolate an Expression on a set of cells") + nb::ndarray cells, + const dolfinx::mesh::Mesh& expr_mesh, + nb::ndarray cell_map) + { + self.interpolate(expr, std::span(cells.data(), cells.size()), + expr_mesh, + std::span(cell_map.data(), cell_map.size())); + }, + nb::arg("expr"), nb::arg("cells"), nb::arg("expr_mesh"), + nb::arg("cell_map"), "Interpolate an Expression on a set of cells") .def_prop_ro( "x", nb::overload_cast<>(&dolfinx::fem::Function::x), "Return the vector associated with the finite element Function") diff --git a/python/dolfinx/wrappers/mesh.cpp b/python/dolfinx/wrappers/mesh.cpp index e6f1b8c54d..e5c0861c72 100644 --- a/python/dolfinx/wrappers/mesh.cpp +++ b/python/dolfinx/wrappers/mesh.cpp @@ -299,8 +299,13 @@ void declare_mesh(nb::module_& m, std::string type) [](const dolfinx::mesh::Mesh& mesh, int dim, nb::ndarray, nb::c_contig> entities) { - return dolfinx::mesh::create_submesh( + auto submesh = dolfinx::mesh::create_submesh( mesh, dim, std::span(entities.data(), entities.size())); + auto _e_map = as_nbarray(std::move(std::get<1>(submesh))); + auto _v_map = as_nbarray(std::move(std::get<2>(submesh))); + auto _g_map = as_nbarray(std::move(std::get<3>(submesh))); + return std::tuple(std::move(std::get<0>(submesh)), _e_map, _v_map, + _g_map); }, nb::arg("mesh"), nb::arg("dim"), nb::arg("entities")); diff --git a/python/test/unit/fem/test_interpolation.py b/python/test/unit/fem/test_interpolation.py index 1834e6ce96..5a9d93205c 100644 --- a/python/test/unit/fem/test_interpolation.py +++ b/python/test/unit/fem/test_interpolation.py @@ -15,7 +15,7 @@ import basix import ufl from basix.ufl import blocked_element, custom_element, element, enriched_element, mixed_element -from dolfinx import default_real_type +from dolfinx import default_real_type, default_scalar_type from dolfinx.fem import ( Expression, Function, @@ -29,6 +29,7 @@ CellType, create_mesh, create_rectangle, + create_submesh, create_unit_cube, create_unit_square, locate_entities, @@ -1025,3 +1026,114 @@ def f_test2(x): l2_error = assemble_scalar(form((u1 - u1_exact) ** 2 * dx_cell(cell_label), dtype=xtype)) assert np.isclose(l2_error, 0.0, rtol=np.finfo(xtype).eps, atol=np.finfo(xtype).eps) + + +def test_submesh_interpolation(): + """Test interpolation of a function between a submesh and its parent""" + mesh = create_unit_square(MPI.COMM_WORLD, 6, 7) + + def left_locator(x): + return x[0] <= 0.5 + 1e-14 + + def ref_func(x): + return x[0] + x[1] ** 2 + + tdim = mesh.topology.dim + cells = locate_entities(mesh, tdim, left_locator) + submesh, sub_to_parent, _, _ = create_submesh(mesh, tdim, cells) + + V = functionspace(mesh, ("Lagrange", 2)) + u = Function(V) + u.interpolate(ref_func) + + V_sub = functionspace(submesh, ("DG", 3)) + u_sub = Function(V_sub) + + # Map from parent to sub mesh + u_sub.interpolate(u, cell_map=sub_to_parent) + + u_sub_exact = Function(V_sub) + u_sub_exact.interpolate(ref_func) + atol = 5 * np.finfo(default_scalar_type).resolution + np.testing.assert_allclose(u_sub_exact.x.array, u_sub.x.array, atol=atol) + + # Map from sub to parent + W = functionspace(mesh, ("DG", 4)) + w = Function(W) + + cell_imap = mesh.topology.index_map(tdim) + num_cells = cell_imap.size_local + cell_imap.num_ghosts + parent_to_sub = np.full(num_cells, -1, dtype=np.int32) + parent_to_sub[sub_to_parent] = np.arange(len(sub_to_parent)) + + # Mapping back needs to be restricted to the subset of cells in the submesh + w.interpolate(u_sub_exact, cells=sub_to_parent, cell_map=parent_to_sub) + + w_exact = Function(W) + w_exact.interpolate(ref_func, cells=cells) + + np.testing.assert_allclose(w.x.array, w_exact.x.array, atol=atol) + + +def test_submesh_expression_interpolation(): + """Test interpolation of an expression between a submesh and its parent""" + mesh = create_unit_square(MPI.COMM_WORLD, 10, 8, cell_type=CellType.quadrilateral) + + def left_locator(x): + return x[0] <= 0.5 + 1e-14 + + def ref_func(x): + return -3 * x[0] ** 2 + x[1] ** 2 + + def grad_ref_func(x): + values = np.zeros((2, x.shape[1]), dtype=default_scalar_type) + values[0] = -6 * x[0] + values[1] = 2 * x[1] + return values + + def modified_grad(x): + grad = grad_ref_func(x) + return -0.2 * grad[1], 0.1 * grad[0] + + tdim = mesh.topology.dim + cells = locate_entities(mesh, tdim, left_locator) + submesh, sub_to_parent, _, _ = create_submesh(mesh, tdim, cells) + + V = functionspace(mesh, ("Lagrange", 2)) + u = Function(V) + u.interpolate(ref_func) + + V_sub = functionspace(submesh, ("N2curl", 1)) + u_sub = Function(V_sub) + + parent_expr = Expression(ufl.grad(u), V_sub.element.interpolation_points()) + + # Map from parent to sub mesh + + u_sub.interpolate(parent_expr, expr_mesh=mesh, cell_map=sub_to_parent) + + u_sub_exact = Function(V_sub) + u_sub_exact.interpolate(grad_ref_func) + atol = 10 * np.finfo(default_scalar_type).resolution + np.testing.assert_allclose(u_sub_exact.x.array, u_sub.x.array, atol=atol) + + # Map from sub to parent + W = functionspace(mesh, ("DQ", 1, (mesh.geometry.dim,))) + w = Function(W) + + cell_imap = mesh.topology.index_map(tdim) + num_cells = cell_imap.size_local + cell_imap.num_ghosts + parent_to_sub = np.full(num_cells, -1, dtype=np.int32) + parent_to_sub[sub_to_parent] = np.arange(len(sub_to_parent)) + + # Map exact solution (based on quadrature points) back to parent mesh + sub_vec = ufl.as_vector((-0.2 * u_sub_exact[1], 0.1 * u_sub_exact[0])) + sub_expr = Expression(sub_vec, W.element.interpolation_points()) + + # Mapping back needs to be restricted to the subset of cells in the submesh + w.interpolate(sub_expr, cells=sub_to_parent, expr_mesh=submesh, cell_map=parent_to_sub) + + w_exact = Function(W) + w_exact.interpolate(modified_grad, cells=cells) + w_exact.x.scatter_forward() + np.testing.assert_allclose(w.x.array, w_exact.x.array, atol=atol)