From c7c6205920f2a45e6a555942f0fdad86908733ce Mon Sep 17 00:00:00 2001 From: Joe Dean Date: Tue, 27 Feb 2024 11:37:55 +0000 Subject: [PATCH] Add maps for mixed-domain assembly (#3056) * Add ent maps to form * Add maps to create_form_factory * Pass entity maps * Add maps * Add additional domain function * Update demo * Update test * Add new sparsity build function * Handle case when same mesh is used * Call mixed-domain sparsity * Add mixed-domain assembler * Call mixed domain assembler and add test * Add some docs * Docs * Use transform * Add braces * Add braces * Clarify docs * Change numering style * Pass cells as array * Remove default arg * Simplify * Use hack to avoid passing shared pointer * Remove second Form constructor * Doc updates * Doc updates * Remove duplicate function. * Remove PETSc dependency in test * Small simplification * Remove duplicated function * Cluster dof data * Tidy up * Small simplifications * Remove include * Small doc fixes * Doc fixes * Improve docs * More docs * Simplification * Small update * Simplification * Simplification --------- Co-authored-by: Garth N. Wells --- cpp/demo/custom_kernel/main.cpp | 15 +- cpp/dolfinx/fem/ElementDofLayout.h | 8 +- cpp/dolfinx/fem/FiniteElement.h | 7 +- cpp/dolfinx/fem/Form.h | 192 ++++++++++++------ cpp/dolfinx/fem/assemble_matrix_impl.h | 159 ++++++++++----- cpp/dolfinx/fem/assemble_vector_impl.h | 17 +- cpp/dolfinx/fem/assembler.h | 2 +- cpp/dolfinx/fem/sparsitybuild.cpp | 27 ++- cpp/dolfinx/fem/sparsitybuild.h | 32 +-- cpp/dolfinx/fem/traits.h | 7 + cpp/dolfinx/fem/utils.h | 47 +++-- cpp/dolfinx/mesh/utils.h | 4 +- python/demo/demo_static-condensation.py | 2 +- python/dolfinx/fem/forms.py | 12 +- python/dolfinx/wrappers/assemble.cpp | 2 +- python/dolfinx/wrappers/fem.cpp | 39 ++-- python/dolfinx/wrappers/petsc.cpp | 4 +- python/test/unit/fem/test_assemble_submesh.py | 70 ++++++- .../test/unit/fem/test_custom_jit_kernels.py | 12 +- 19 files changed, 445 insertions(+), 213 deletions(-) diff --git a/cpp/demo/custom_kernel/main.cpp b/cpp/demo/custom_kernel/main.cpp index 77c8ebda697..5dfa4a10ca0 100644 --- a/cpp/demo/custom_kernel/main.cpp +++ b/cpp/demo/custom_kernel/main.cpp @@ -76,12 +76,12 @@ double assemble_matrix0(std::shared_ptr> V, auto kernel, // Associate kernel with cells (as opposed to facets, etc) std::map integrals{std::pair{fem::IntegralType::cell, kernel_data}}; - fem::Form a({V, V}, integrals, {}, {}, false, V->mesh()); + fem::Form a({V, V}, integrals, {}, {}, false, {}, V->mesh()); auto dofmap = V->dofmap(); auto sp = la::SparsityPattern( V->mesh()->comm(), {dofmap->index_map, dofmap->index_map}, {dofmap->index_map_bs(), dofmap->index_map_bs()}); - fem::sparsitybuild::cells(sp, cells, {*dofmap, *dofmap}); + fem::sparsitybuild::cells(sp, {cells, cells}, {*dofmap, *dofmap}); sp.finalize(); la::MatrixCSR A(sp); common::Timer timer("Assembler0 std::function (matrix)"); @@ -103,7 +103,7 @@ double assemble_vector0(std::shared_ptr> V, auto kernel, auto mesh = V->mesh(); std::vector kernal_data{fem::integral_data(-1, kernel, cells)}; std::map integrals{std::pair{fem::IntegralType::cell, kernal_data}}; - fem::Form L({V}, integrals, {}, {}, false, mesh); + fem::Form L({V}, integrals, {}, {}, false, {}, mesh); auto dofmap = V->dofmap(); la::Vector b(dofmap->index_map, 1); common::Timer timer("Assembler0 std::function (vector)"); @@ -130,14 +130,15 @@ double assemble_matrix1(const mesh::Geometry& g, const fem::DofMap& dofmap, auto sp = la::SparsityPattern(dofmap.index_map->comm(), {dofmap.index_map, dofmap.index_map}, {dofmap.index_map_bs(), dofmap.index_map_bs()}); - fem::sparsitybuild::cells(sp, cells, {dofmap, dofmap}); + fem::sparsitybuild::cells(sp, {cells, cells}, {dofmap, dofmap}); sp.finalize(); la::MatrixCSR A(sp); auto ident = [](auto, auto, auto, auto) {}; // DOF permutation not required common::Timer timer("Assembler1 lambda (matrix)"); - fem::impl::assemble_cells(A.mat_add_values(), g.dofmap(), g.x(), cells, ident, - dofmap.map(), 1, ident, dofmap.map(), 1, {}, {}, - kernel, std::span(), 0, {}, {}); + fem::impl::assemble_cells(A.mat_add_values(), g.dofmap(), g.x(), cells, + {dofmap.map(), 1, cells}, ident, + {dofmap.map(), 1, cells}, ident, {}, {}, kernel, + std::span(), 0, {}, {}, {}); A.scatter_rev(); return A.squared_norm(); } diff --git a/cpp/dolfinx/fem/ElementDofLayout.h b/cpp/dolfinx/fem/ElementDofLayout.h index 9aca21261f3..f7f12aa13c5 100644 --- a/cpp/dolfinx/fem/ElementDofLayout.h +++ b/cpp/dolfinx/fem/ElementDofLayout.h @@ -18,15 +18,15 @@ enum class CellType; namespace dolfinx::fem { -/// The class represents the degree-of-freedom (dofs) for an element. -/// Dofs are associated with a mesh entity. This class also handles -/// sub-space dofs, which are views into the parent dofs. // TODO: For this class/concept to be robust, the topology of the // reference cell needs to be defined. - +// // TODO: Handle block dofmaps properly +/// The class represents the degree-of-freedom (dofs) for an element. +/// Dofs are associated with a mesh entity. This class also handles +/// sub-space dofs, which are views into the parent dofs. class ElementDofLayout { public: diff --git a/cpp/dolfinx/fem/FiniteElement.h b/cpp/dolfinx/fem/FiniteElement.h index dc357677d1e..1f1f9decdb9 100644 --- a/cpp/dolfinx/fem/FiniteElement.h +++ b/cpp/dolfinx/fem/FiniteElement.h @@ -6,6 +6,7 @@ #pragma once +#include "traits.h" #include #include #include @@ -21,12 +22,6 @@ struct ufcx_finite_element; namespace dolfinx::fem { -/// @brief DOF transform kernel concept. -template -concept DofTransformKernel - = std::is_invocable_v, std::span, - std::int32_t, int>; - /// Finite Element, containing the dof layout on a reference element, /// and various methods for evaluating and transforming the basis. template diff --git a/cpp/dolfinx/fem/Form.h b/cpp/dolfinx/fem/Form.h index b6d41939752..1b90ede7b8c 100644 --- a/cpp/dolfinx/fem/Form.h +++ b/cpp/dolfinx/fem/Form.h @@ -40,22 +40,24 @@ enum class IntegralType : std::int8_t /// @brief Represents integral data, containing the integral ID, the /// kernel, and a list of entities to integrate over. -template Kern = std::function> +template > struct integral_data { - /// @brief Kernel type - using kern_t = Kern; - /// @brief Create a structure to hold integral data. - /// @tparam U `std::vector` holding entity indices. + /// @tparam U Entity indices type. /// @param id Domain ID. /// @param kernel Integration kernel. /// @param entities Entities to integrate over. - template - integral_data(int id, kern_t kernel, U&& entities) - : id(id), kernel(kernel), entities(std::forward(entities)) + template + requires std::is_convertible_v< + std::remove_cvref_t, + std::function> + and std::is_convertible_v, + std::vector> + integral_data(int id, K&& kernel, V&& entities) + : id(id), kernel(std::forward(kernel)), + entities(std::forward(entities)) { } @@ -63,8 +65,13 @@ struct integral_data /// @param id Domain ID /// @param kernel Integration kernel. /// @param e Entities to integrate over. - integral_data(int id, kern_t kernel, std::span e) - : id(id), kernel(kernel), entities(e.begin(), e.end()) + template + requires std::is_convertible_v< + std::remove_cvref_t, + std::function> + integral_data(int id, K&& kernel, std::span e) + : id(id), kernel(std::forward(kernel)), entities(e.begin(), e.end()) { } @@ -72,7 +79,9 @@ struct integral_data int id; /// The integration kernel - kern_t kernel; + std::function + kernel; /// The entities to integrate over std::vector entities; @@ -101,37 +110,50 @@ struct integral_data /// (the variable `function_spaces` in the constructors below), the list /// of spaces should start with space number 0 (the test space) and then /// space number 1 (the trial space). -template < - dolfinx::scalar T, std::floating_point U = dolfinx::scalar_value_type_t, - FEkernel Kern - = std::function*, - const int*, const std::uint8_t*)>> +/// +/// @tparam T Scalar type in the form. +/// @tparam U Float (real) type used for the finite element and geometry. +/// @tparam Kern Element kernel. +template > class Form { - using kern_t = Kern; - public: /// Scalar type using scalar_type = T; /// @brief Create a finite element form. /// - /// @note User applications will normally call a builder function + /// @note User applications will normally call a factory function /// rather using this interface directly. /// - /// @param[in] V Function spaces for the form arguments - /// @param[in] integrals The integrals in the form. For each - /// integral type, there is a list of integral data - /// @param[in] coefficients - /// @param[in] constants Constants in the Form - /// @param[in] needs_facet_permutations Set to true is any of the - /// integration kernels require cell permutation data + /// @param[in] V Function spaces for the form arguments. + /// @param[in] integrals The integrals in the form. For each integral + /// type, there is a list of integral data. + /// @param[in] coefficients Coefficients in the form. + /// @param[in] constants Constants in the form. + /// @param[in] needs_facet_permutations Set to `true` is any of the + /// integration kernels require cell permutation data. + /// @param[in] entity_maps If any trial functions, test functions, or + /// coefficients in the form are not defined over the same mesh as the + /// integration domain, `entity_maps` must be supplied. For each key + /// (a mesh, different to the integration domain mesh) a map should be + /// provided relating the entities in the integration domain mesh to + /// the entities in the key mesh e.g. for a pair (msh, emap) in + /// `entity_maps`, `emap[i]` is the entity in `msh` corresponding to + /// entity `i` in the integration domain mesh. /// @param[in] mesh Mesh of the domain. This is required when there /// are no argument functions from which the mesh can be extracted, /// e.g. for functionals. /// + /// @note For the single domain case, pass an empty `entity_maps`. + /// /// @pre The integral data in integrals must be sorted by domain + /// (domain id). template + requires std::is_convertible_v< + std::remove_cvref_t, + std::map>>> Form(const std::vector>>& V, X&& integrals, const std::vector>>& @@ -139,6 +161,8 @@ class Form const std::vector>>& constants, bool needs_facet_permutations, + const std::map>, + std::span>& entity_maps, std::shared_ptr> mesh = nullptr) : _function_spaces(V), _coefficients(coefficients), _constants(constants), _mesh(mesh), _needs_facet_permutations(needs_facet_permutations) @@ -147,8 +171,14 @@ class Form if (!_mesh and !V.empty()) _mesh = V[0]->mesh(); for (auto& space : V) - if (_mesh != space->mesh()) - throw std::runtime_error("Incompatible mesh"); + { + if (_mesh != space->mesh() + and entity_maps.find(space->mesh()) == entity_maps.end()) + { + throw std::runtime_error( + "Incompatible mesh. entity_maps must be provided."); + } + } if (!_mesh) throw std::runtime_error("No mesh could be associated with the Form."); @@ -161,11 +191,15 @@ class Form throw std::runtime_error("Integral IDs not sorted"); } - std::vector>& itg + std::vector>& itg = _integrals[static_cast(domain_type)]; for (auto&& [id, kern, e] : data) itg.emplace_back(id, kern, std::move(e)); } + + // Store entity maps + for (auto [msh, map] : entity_maps) + _entity_maps.insert({msh, std::vector(map.begin(), map.end())}); } /// Copy constructor @@ -177,29 +211,33 @@ class Form /// Destructor virtual ~Form() = default; - /// Rank of the form (bilinear form = 2, linear form = 1, functional = - /// 0, etc) + /// @brief Rank of the form. + /// + /// bilinear form = 2, linear form = 1, functional = 0, etc. + /// /// @return The rank of the form int rank() const { return _function_spaces.size(); } - /// Extract common mesh for the form - /// @return The mesh + /// @brief Extract common mesh for the form. + /// @return The mesh. std::shared_ptr> mesh() const { return _mesh; } - /// Return function spaces for all arguments - /// @return Function spaces + /// @brief Function spaces for all arguments. + /// @return Function spaces. const std::vector>>& function_spaces() const { return _function_spaces; } - /// @brief Get the kernel function for integral i on given domain + /// @brief Get the kernel function for integral `i` on given domain /// type. /// @param[in] type Integral type /// @param[in] i Domain identifier (index) /// @return Function to call for tabulate_tensor - kern_t kernel(IntegralType type, int i) const + std::function + kernel(IntegralType type, int i) const { const auto& integrals = _integrals[static_cast(type)]; auto it = std::lower_bound(integrals.begin(), integrals.end(), i, @@ -233,12 +271,13 @@ class Form return _integrals[static_cast(type)].size(); } - /// Get the IDs for integrals (kernels) for given integral type. The - /// IDs correspond to the domain IDs which the integrals are defined - /// for in the form. ID=-1 is the default integral over the whole - /// domain. - /// @param[in] type Integral type - /// @return List of IDs for given integral type + /// @brief Get the IDs for integrals (kernels) for given integral type. + /// + /// The IDs correspond to the domain IDs which the integrals are + /// defined for in the form. `ID=-1` is the default integral over the + /// whole domain. + /// @param[in] type Integral type. + /// @return List of IDs for given integral type. std::vector integral_ids(IntegralType type) const { std::vector ids; @@ -248,8 +287,8 @@ class Form return ids; } - /// @brief Get the list of cell indices for the ith integral (kernel) - /// for the cell domain type. + /// @brief Get the list of mesh entity indices for the ith integral + /// (kernel) of a given type. /// /// For IntegralType::cell, returns a list of cell indices. /// @@ -262,9 +301,9 @@ class Form /// local_facet_index_1)`. Data is flattened with row-major layout, /// `shape=(num_facets, 4)`. /// - /// @param[in] type Integral domain type - /// @param[in] i Integral ID, i.e. (sub)domain index - /// @return List of active cell entities for the given integral (kernel) + /// @param[in] type Integral domain type. + /// @param[in] i Integral ID, i.e. (sub)domain index. + /// @return List of active entities for the given integral (kernel). std::span domain(IntegralType type, int i) const { const auto& integrals = _integrals[static_cast(type)]; @@ -277,20 +316,51 @@ class Form throw std::runtime_error("No mesh entities for requested domain index."); } - /// Access coefficients + /// @brief Compute the list of entity indices in `mesh` for the + /// ith integral (kernel) of a given type (i.e. cell, exterior facet, or + /// interior facet). + /// + /// @param type Integral type. + /// @param i Integral ID, i.e. the (sub)domain index. + /// @param mesh The mesh the entities are numbered with respect to. + /// @return List of active entities in `mesh` for the given integral. + std::vector domain(IntegralType type, int i, + const mesh::Mesh& mesh) const + { + // Hack to avoid passing shared pointer to this function + std::shared_ptr> msh_ptr(&mesh, + [](const mesh::Mesh*) {}); + + std::span entities = domain(type, i); + if (msh_ptr == _mesh) + return std::vector(entities.begin(), entities.end()); + else + { + std::span entity_map = _entity_maps.at(msh_ptr); + std::vector mapped_entities; + mapped_entities.reserve(entities.size()); + std::transform(entities.begin(), entities.end(), + std::back_inserter(mapped_entities), + [&entity_map](auto e) { return entity_map[e]; }); + return mapped_entities; + } + } + + /// @brief Access coefficients. const std::vector>>& coefficients() const { return _coefficients; } - /// Get bool indicating whether permutation data needs to be passed - /// into these integrals + /// @brief Get bool indicating whether permutation data needs to be + /// passed into these integrals. /// @return True if cell permutation data is required bool needs_facet_permutations() const { return _needs_facet_permutations; } - /// Offset for each coefficient expansion array on a cell. Used to - /// pack data for multiple coefficients in a flat array. The last - /// entry is the size required to store all coefficients. + /// @brief Offset for each coefficient expansion array on a cell. + /// + /// Used to pack data for multiple coefficients in a flat array. The + /// last entry is the size required to store all coefficients. std::vector coefficient_offsets() const { std::vector n = {0}; @@ -303,7 +373,7 @@ class Form return n; } - /// Access constants + /// @brief Access constants. const std::vector>>& constants() const { return _constants; @@ -315,7 +385,7 @@ class Form // Integrals. Array index is // static_cast>, 4> _integrals; + std::array>, 4> _integrals; // Form coefficients std::vector>> _coefficients; @@ -328,5 +398,9 @@ class Form // True if permutation data needs to be passed into these integrals bool _needs_facet_permutations; -}; + + // Entity maps (see Form documentation) + std::map>, std::vector> + _entity_maps; +}; // namespace dolfinx::fem } // namespace dolfinx::fem diff --git a/cpp/dolfinx/fem/assemble_matrix_impl.h b/cpp/dolfinx/fem/assemble_matrix_impl.h index 61a4d6235ee..5b15fef8001 100644 --- a/cpp/dolfinx/fem/assemble_matrix_impl.h +++ b/cpp/dolfinx/fem/assemble_matrix_impl.h @@ -19,6 +19,7 @@ #include #include #include +#include #include namespace dolfinx::fem::impl @@ -28,27 +29,57 @@ using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan< const std::int32_t, MDSPAN_IMPL_STANDARD_NAMESPACE::dextents>; -/// Execute kernel over cells and accumulate result in matrix +/// @brief Execute kernel over cells and accumulate result in matrix. +/// @tparam T Matrix/form scalar type. +/// @param mat_set Function that accumulates computed entries into a +/// matrix. +/// @param x_dofmap Dofmap for the mesh geometry. +/// @param x Mesh geometry (coordinates). +/// @param cells Cell indices (in the integration domain mesh) to execute +/// the kernel over. These are the indices into the geometry dofmap. +/// @param dofmap0 Test function (row) degree-of-freedom data holding +/// the (0) dofmap, (1) dofmap block size and (2) dofmap cell indices. +/// @param P0 Function that applies transformation P_0 A in-place to +/// transform trial degrees-of-freedom. +/// @param dofmap1 Trial function (column) degree-of-freedom data +/// holding the (0) dofmap, (1) dofmap block size and (2) dofmap cell +/// indices. +/// @param P1T Function that applies transformation A P_1^T in-place to +/// transform trial degrees-of-freedom. +/// @param bc0 Marker for rows with Dirichlet boundary conditions applied +/// @param bc1 Marker for columns with Dirichlet boundary conditions applied +/// @param kernel Kernel function to execute over each cell. +/// @param coeffs The coefficient data array of shape (cells.size(), cstride), +/// flattened into row-major format. +/// @param cstride The coefficient stride +/// @param constants The constant data +/// @param cell_info0 The cell permutation information for the test function +/// mesh +/// @param cell_info1 The cell permutation information for the trial function +/// mesh template -void assemble_cells(la::MatSet auto mat_set, mdspan2_t x_dofmap, - std::span> x, - std::span cells, - fem::DofTransformKernel auto pre_dof_transform, - mdspan2_t dofmap0, int bs0, - fem::DofTransformKernel auto post_dof_transform, - mdspan2_t dofmap1, int bs1, - std::span bc0, - std::span bc1, FEkernel auto kernel, - std::span coeffs, int cstride, - std::span constants, - std::span cell_info) +void assemble_cells( + la::MatSet auto mat_set, mdspan2_t x_dofmap, + std::span> 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 coeffs, int cstride, std::span constants, + std::span cell_info0, + std::span cell_info1) { if (cells.empty()) return; + const auto [dmap0, bs0, cells0] = dofmap0; + const auto [dmap1, bs1, cells1] = dofmap1; + // Iterate over active cells - const int num_dofs0 = dofmap0.extent(1); - const int num_dofs1 = dofmap1.extent(1); + const int num_dofs0 = dmap0.extent(1); + const int num_dofs1 = dmap1.extent(1); const int ndim0 = bs0 * num_dofs0; const int ndim1 = bs1 * num_dofs1; std::vector Ae(ndim0 * ndim1); @@ -56,9 +87,15 @@ void assemble_cells(la::MatSet auto mat_set, mdspan2_t x_dofmap, std::vector> coordinate_dofs(3 * x_dofmap.extent(1)); // Iterate over active cells + assert(cells0.size() == cells.size()); + assert(cells1.size() == cells.size()); for (std::size_t index = 0; index < cells.size(); ++index) { + // Cell index in integration domain mesh (c), test function mesh + // (c0) and trial function mesh (c1) std::int32_t c = cells[index]; + std::int32_t c0 = cells0[index]; + std::int32_t c1 = cells1[index]; // Get cell coordinates/geometry auto x_dofs @@ -75,12 +112,13 @@ void assemble_cells(la::MatSet auto mat_set, mdspan2_t x_dofmap, kernel(Ae.data(), coeffs.data() + index * cstride, constants.data(), coordinate_dofs.data(), nullptr, nullptr); - pre_dof_transform(_Ae, cell_info, c, ndim1); - post_dof_transform(_Ae, cell_info, c, ndim0); + // Compute A = P_0 \tilde{A} P_1^T (dof transformation) + P0(_Ae, cell_info1, c1, ndim1); // B = P0 \tilde{A} + P1T(_Ae, cell_info0, c0, ndim0); // A = B P1_T // Zero rows/columns for essential bcs - auto dofs0 = std::span(dofmap0.data_handle() + c * num_dofs0, num_dofs0); - auto dofs1 = std::span(dofmap1.data_handle() + c * num_dofs1, num_dofs1); + auto dofs0 = std::span(dmap0.data_handle() + c0 * num_dofs0, num_dofs0); + auto dofs1 = std::span(dmap1.data_handle() + c1 * num_dofs1, num_dofs1); if (!bc0.empty()) { @@ -92,7 +130,7 @@ void assemble_cells(la::MatSet auto mat_set, mdspan2_t x_dofmap, { // Zero row bs0 * i + k const int row = bs0 * i + k; - std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0.0); + std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0); } } } @@ -109,7 +147,7 @@ void assemble_cells(la::MatSet auto mat_set, mdspan2_t x_dofmap, // Zero column bs1 * j + k const int col = bs1 * j + k; for (int row = 0; row < ndim0; ++row) - Ae[row * ndim1 + col] = 0.0; + Ae[row * ndim1 + col] = 0; } } } @@ -124,9 +162,8 @@ template void assemble_exterior_facets( la::MatSet auto mat_set, mdspan2_t x_dofmap, std::span> x, - std::span facets, - fem::DofTransformKernel auto pre_dof_transform, mdspan2_t dofmap0, - int bs0, fem::DofTransformKernel auto post_dof_transform, + std::span facets, fem::DofTransformKernel auto P0, + mdspan2_t dofmap0, int bs0, fem::DofTransformKernel auto P1T, mdspan2_t dofmap1, int bs1, std::span bc0, std::span bc1, FEkernel auto kernel, std::span coeffs, int cstride, std::span constants, @@ -164,8 +201,8 @@ void assemble_exterior_facets( kernel(Ae.data(), coeffs.data() + index / 2 * cstride, constants.data(), coordinate_dofs.data(), &local_facet, nullptr); - pre_dof_transform(_Ae, cell_info, cell, ndim1); - post_dof_transform(_Ae, cell_info, cell, ndim0); + P0(_Ae, cell_info, cell, ndim1); + P1T(_Ae, cell_info, cell, ndim0); // Zero rows/columns for essential bcs auto dofs0 = std::span(dofmap0.data_handle() + cell * num_dofs0, num_dofs0); @@ -180,7 +217,7 @@ void assemble_exterior_facets( { // Zero row bs0 * i + k const int row = bs0 * i + k; - std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0.0); + std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0); } } } @@ -196,7 +233,7 @@ void assemble_exterior_facets( // Zero column bs1 * j + k const int col = bs1 * j + k; for (int row = 0; row < ndim0; ++row) - Ae[row * ndim1 + col] = 0.0; + Ae[row * ndim1 + col] = 0; } } } @@ -211,9 +248,8 @@ template void assemble_interior_facets( la::MatSet auto mat_set, mdspan2_t x_dofmap, std::span> x, int num_cell_facets, - std::span facets, - fem::DofTransformKernel auto pre_dof_transform, const DofMap& dofmap0, - int bs0, fem::DofTransformKernel auto post_dof_transform, + std::span facets, fem::DofTransformKernel auto P0, + const DofMap& dofmap0, int bs0, fem::DofTransformKernel auto P1T, const DofMap& dofmap1, int bs1, std::span bc0, std::span bc1, FEkernel auto kernel, std::span coeffs, int cstride, std::span offsets, @@ -300,9 +336,9 @@ void assemble_interior_facets( std::span sub_Ae0 = _Ae.subspan(bs0 * dmap0_cell0.size() * num_cols, bs0 * dmap0_cell1.size() * num_cols); - pre_dof_transform(_Ae, cell_info, cells[0], num_cols); - pre_dof_transform(sub_Ae0, cell_info, cells[1], num_cols); - post_dof_transform(_Ae, cell_info, cells[0], num_rows); + P0(_Ae, cell_info, cells[0], num_cols); + P0(sub_Ae0, cell_info, cells[1], num_cols); + P1T(_Ae, cell_info, cells[0], num_rows); for (int row = 0; row < num_rows; ++row) { @@ -310,7 +346,7 @@ void assemble_interior_facets( // the block matrix, so each row needs a separate span access std::span sub_Ae1 = _Ae.subspan( row * num_cols + bs1 * dmap1_cell0.size(), bs1 * dmap1_cell1.size()); - post_dof_transform(sub_Ae1, cell_info, cells[1], 1); + P1T(sub_Ae1, cell_info, cells[1], 1); } // Zero rows/columns for essential bcs @@ -324,7 +360,7 @@ void assemble_interior_facets( { // Zero row bs0 * i + k std::fill_n(std::next(Ae.begin(), num_cols * (bs0 * i + k)), - num_cols, 0.0); + num_cols, 0); } } } @@ -339,7 +375,7 @@ void assemble_interior_facets( { // Zero column bs1 * j + k for (int m = 0; m < num_rows; ++m) - Ae[m * num_cols + bs1 * j + k] = 0.0; + Ae[m * num_cols + bs1 * j + k] = 0; } } } @@ -352,7 +388,7 @@ void assemble_interior_facets( /// The matrix A must already be initialised. The matrix may be a proxy, /// i.e. a view into a larger matrix, and assembly is performed using /// local indices. Rows (bc0) and columns (bc1) with Dirichlet -/// conditions are zeroed. Markers (bc0 and bc1) can be empty if not bcs +/// conditions are zeroed. Markers (bc0 and bc1) can be empty if no bcs /// are applied. Matrix is not finalised. template void assemble_matrix( @@ -362,8 +398,15 @@ void assemble_matrix( std::pair, int>>& coefficients, std::span bc0, std::span bc1) { + // Integration domain mesh std::shared_ptr> mesh = a.mesh(); assert(mesh); + // Test function mesh + auto mesh0 = a.function_spaces().at(0)->mesh(); + assert(mesh0); + // Trial function mesh + auto mesh1 = a.function_spaces().at(1)->mesh(); + assert(mesh1); // Get dofmap data std::shared_ptr dofmap0 @@ -381,18 +424,21 @@ void assemble_matrix( assert(element0); auto element1 = a.function_spaces().at(1)->element(); assert(element1); - fem::DofTransformKernel auto pre_dof_transform + fem::DofTransformKernel auto P0 = element0->template get_pre_dof_transformation_function(); - fem::DofTransformKernel auto post_dof_transform + fem::DofTransformKernel auto P1T = element1->template get_post_dof_transformation_function( FiniteElement::doftransform::transpose); - std::span cell_info; + std::span cell_info0; + std::span cell_info1; if (element0->needs_dof_transformations() or element1->needs_dof_transformations() or a.needs_facet_permutations()) { - mesh->topology_mutable()->create_entity_permutations(); - cell_info = std::span(mesh->topology()->get_cell_permutation_info()); + mesh0->topology_mutable()->create_entity_permutations(); + mesh1->topology_mutable()->create_entity_permutations(); + cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info()); + cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info()); } for (int i : a.integral_ids(IntegralType::cell)) @@ -400,10 +446,11 @@ void assemble_matrix( auto fn = a.kernel(IntegralType::cell, i); assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i}); - impl::assemble_cells(mat_set, x_dofmap, x, a.domain(IntegralType::cell, i), - pre_dof_transform, dofs0, bs0, post_dof_transform, - dofs1, bs1, bc0, bc1, fn, coeffs, cstride, constants, - cell_info); + impl::assemble_cells( + mat_set, x_dofmap, x, a.domain(IntegralType::cell, i), + {dofs0, bs0, a.domain(IntegralType::cell, i, *mesh0)}, P0, + {dofs1, bs1, a.domain(IntegralType::cell, i, *mesh1)}, P1T, bc0, bc1, + fn, coeffs, cstride, constants, cell_info0, cell_info1); } for (int i : a.integral_ids(IntegralType::exterior_facet)) @@ -412,10 +459,10 @@ void assemble_matrix( assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::exterior_facet, i}); - impl::assemble_exterior_facets( - mat_set, x_dofmap, x, a.domain(IntegralType::exterior_facet, i), - pre_dof_transform, dofs0, bs0, post_dof_transform, dofs1, bs1, bc0, bc1, - fn, coeffs, cstride, constants, cell_info); + impl::assemble_exterior_facets(mat_set, x_dofmap, x, + a.domain(IntegralType::exterior_facet, i), + P0, dofs0, bs0, P1T, dofs1, bs1, bc0, bc1, + fn, coeffs, cstride, constants, cell_info0); } if (a.num_integrals(IntegralType::interior_facet) > 0) @@ -441,11 +488,11 @@ void assemble_matrix( assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::interior_facet, i}); - impl::assemble_interior_facets( - mat_set, x_dofmap, x, num_cell_facets, - a.domain(IntegralType::interior_facet, i), pre_dof_transform, - *dofmap0, bs0, post_dof_transform, *dofmap1, bs1, bc0, bc1, fn, - coeffs, cstride, c_offsets, constants, cell_info, get_perm); + impl::assemble_interior_facets(mat_set, x_dofmap, x, num_cell_facets, + a.domain(IntegralType::interior_facet, i), + P0, *dofmap0, bs0, P1T, *dofmap1, bs1, bc0, + bc1, fn, coeffs, cstride, c_offsets, + constants, cell_info0, get_perm); } } } diff --git a/cpp/dolfinx/fem/assemble_vector_impl.h b/cpp/dolfinx/fem/assemble_vector_impl.h index 9cd2f3f55da..cf6653a20a8 100644 --- a/cpp/dolfinx/fem/assemble_vector_impl.h +++ b/cpp/dolfinx/fem/assemble_vector_impl.h @@ -11,10 +11,10 @@ #include "DofMap.h" #include "Form.h" #include "FunctionSpace.h" +#include "traits.h" #include "utils.h" #include #include -#include #include #include #include @@ -144,8 +144,8 @@ void _lift_bc_cells( if (bc_markers1[jj]) { const T bc = bc_values1[jj]; - const T _x0 = x0.empty() ? 0.0 : x0[jj]; - // const T _x0 = 0.0; + const T _x0 = x0.empty() ? 0 : x0[jj]; + // const T _x0 = 0; // be -= Ae.col(bs1 * j + k) * scale * (bc - _x0); for (int m = 0; m < num_rows; ++m) be[m] -= Ae[m * num_cols + _bs1 * j + k] * scale * (bc - _x0); @@ -161,7 +161,7 @@ void _lift_bc_cells( if (bc_markers1[jj]) { const T bc = bc_values1[jj]; - const T _x0 = x0.empty() ? 0.0 : x0[jj]; + const T _x0 = x0.empty() ? 0 : x0[jj]; // be -= Ae.col(bs1 * j + k) * scale * (bc - _x0); for (int m = 0; m < num_rows; ++m) be[m] -= Ae[m * num_cols + bs1 * j + k] * scale * (bc - _x0); @@ -275,7 +275,7 @@ void _lift_bc_exterior_facets( if (bc_markers1[jj]) { const T bc = bc_values1[jj]; - const T _x0 = x0.empty() ? 0.0 : x0[jj]; + const T _x0 = x0.empty() ? 0 : x0[jj]; // be -= Ae.col(bs1 * j + k) * scale * (bc - _x0); for (int m = 0; m < num_rows; ++m) be[m] -= Ae[m * num_cols + bs1 * j + k] * scale * (bc - _x0); @@ -442,7 +442,7 @@ void _lift_bc_interior_facets( if (bc_markers1[jj]) { const T bc = bc_values1[jj]; - const T _x0 = x0.empty() ? 0.0 : x0[jj]; + const T _x0 = x0.empty() ? 0 : x0[jj]; // be -= Ae.col(bs1 * j + k) * scale * (bc - _x0); for (int m = 0; m < num_rows; ++m) be[m] -= Ae[m * num_cols + bs1 * j + k] * scale * (bc - _x0); @@ -460,7 +460,7 @@ void _lift_bc_interior_facets( if (bc_markers1[jj]) { const T bc = bc_values1[jj]; - const T _x0 = x0.empty() ? 0.0 : x0[jj]; + const T _x0 = x0.empty() ? 0 : x0[jj]; // be -= Ae.col(offset + bs1 * j + k) * scale * (bc - x0[jj]); for (int m = 0; m < num_rows; ++m) { @@ -497,7 +497,6 @@ void assemble_cells(fem::DofTransformKernel auto dof_transform, int cstride, std::span cell_info) { assert(_bs < 0 or _bs == bs); - if (cells.empty()) return; @@ -891,7 +890,7 @@ void apply_lifting( assert(map1); const int crange = bs1 * (map1->size_local() + map1->num_ghosts()); bc_markers1.assign(crange, false); - bc_values1.assign(crange, 0.0); + bc_values1.assign(crange, 0); for (const std::shared_ptr>& bc : bcs1[j]) { bc->mark_dofs(bc_markers1); diff --git a/cpp/dolfinx/fem/assembler.h b/cpp/dolfinx/fem/assembler.h index 7a98e419bea..72aaeadb553 100644 --- a/cpp/dolfinx/fem/assembler.h +++ b/cpp/dolfinx/fem/assembler.h @@ -21,7 +21,7 @@ namespace dolfinx::fem { template class DirichletBC; -template Kern> +template class Form; template class FunctionSpace; diff --git a/cpp/dolfinx/fem/sparsitybuild.cpp b/cpp/dolfinx/fem/sparsitybuild.cpp index 71f7648c1ba..a1d5e95441d 100644 --- a/cpp/dolfinx/fem/sparsitybuild.cpp +++ b/cpp/dolfinx/fem/sparsitybuild.cpp @@ -6,23 +6,22 @@ #include "sparsitybuild.h" #include "DofMap.h" -#include -#include #include -#include using namespace dolfinx; using namespace dolfinx::fem; //----------------------------------------------------------------------------- void sparsitybuild::cells( - la::SparsityPattern& pattern, std::span cells, + la::SparsityPattern& pattern, + std::array, 2> cells, std::array, 2> dofmaps) { + assert(cells[0].size() == cells[1].size()); const DofMap& map0 = dofmaps[0].get(); const DofMap& map1 = dofmaps[1].get(); - for (auto c : cells) - pattern.insert(map0.cell_dofs(c), map1.cell_dofs(c)); + for (std::size_t i = 0; i < cells[0].size(); ++i) + pattern.insert(map0.cell_dofs(cells[0][i]), map1.cell_dofs(cells[1][i])); } //----------------------------------------------------------------------------- void sparsitybuild::interior_facets( @@ -32,16 +31,16 @@ void sparsitybuild::interior_facets( std::array, 2> macro_dofs; for (std::size_t index = 0; index < facets.size(); index += 2) { - int cell_0 = facets[index]; - int cell_1 = facets[index + 1]; + std::int32_t cell0 = facets[index]; + std::int32_t cell1 = facets[index + 1]; for (std::size_t i = 0; i < 2; ++i) { - auto cell_dofs_0 = dofmaps[i].get().cell_dofs(cell_0); - auto cell_dofs_1 = dofmaps[i].get().cell_dofs(cell_1); - macro_dofs[i].resize(cell_dofs_0.size() + cell_dofs_1.size()); - std::copy(cell_dofs_0.begin(), cell_dofs_0.end(), macro_dofs[i].begin()); - std::copy(cell_dofs_1.begin(), cell_dofs_1.end(), - std::next(macro_dofs[i].begin(), cell_dofs_0.size())); + auto cell_dofs0 = dofmaps[i].get().cell_dofs(cell0); + auto cell_dofs1 = dofmaps[i].get().cell_dofs(cell1); + macro_dofs[i].resize(cell_dofs0.size() + cell_dofs1.size()); + std::copy(cell_dofs0.begin(), cell_dofs0.end(), macro_dofs[i].begin()); + std::copy(cell_dofs1.begin(), cell_dofs1.end(), + std::next(macro_dofs[i].begin(), cell_dofs0.size())); } pattern.insert(macro_dofs[0], macro_dofs[1]); } diff --git a/cpp/dolfinx/fem/sparsitybuild.h b/cpp/dolfinx/fem/sparsitybuild.h index 57422164cde..14d53b04d8e 100644 --- a/cpp/dolfinx/fem/sparsitybuild.h +++ b/cpp/dolfinx/fem/sparsitybuild.h @@ -16,11 +16,6 @@ namespace dolfinx::la class SparsityPattern; } -namespace dolfinx::mesh -{ -class Topology; -} - namespace dolfinx::fem { class DofMap; @@ -30,20 +25,31 @@ namespace sparsitybuild { /// @brief Iterate over cells and insert entries into sparsity pattern. /// -/// @param[in,out] pattern The sparsity pattern to insert into -/// @param[in] cells The cell indices -/// @param[in] dofmaps Dofmaps to used in building the sparsity pattern -/// @note The sparsity pattern is not finalised -void cells(la::SparsityPattern& pattern, std::span cells, +/// Inserts the rectangular blocks of indices `dofmap[0][cells[0][i]] x +/// dofmap[1][cells[1][i]]` into the sparsity pattern, i.e. entries +/// (dofmap[0][cells[0][i]][k0], dofmap[0][cells[0][i]][k1])` will +/// appear in the sparsity pattern. +/// +/// @param pattern Sparsity pattern to insert into. +/// @param cells Lists of cells to iterate over. `cells[0]` and +/// `cells[1]` must have the same size. +/// @param dofmaps Dofmaps to used in building the sparsity pattern. +/// @note The sparsity pattern is not finalised. +void cells(la::SparsityPattern& pattern, + std::array, 2> cells, std::array, 2> dofmaps); /// @brief Iterate over interior facets and insert entries into sparsity /// pattern. /// +/// Inserts the rectangular blocks of indices `[dofmap[0][cell0], +/// dofmap[0][cell1]] x [dofmap[1][cell0] + dofmap[1][cell1]]` where +/// `cell0` and `cell1` are the two cells attached to a facet. +/// /// @param[in,out] pattern Sparsity pattern to insert into -/// @param[in] facets Facets as `(cell0, cell1)` pairs for each facet. -/// @param[in] dofmaps The dofmap to use in building the sparsity -/// pattern. +/// @param[in] facets Facets as `(cell0, cell1)` pairs (row-major) for +/// each facet. +/// @param[in] dofmaps Dofmaps to use in building the sparsity pattern. /// @note The sparsity pattern is not finalised. void interior_facets( la::SparsityPattern& pattern, std::span facets, diff --git a/cpp/dolfinx/fem/traits.h b/cpp/dolfinx/fem/traits.h index 52ee1718d9d..daca676c0e9 100644 --- a/cpp/dolfinx/fem/traits.h +++ b/cpp/dolfinx/fem/traits.h @@ -8,11 +8,18 @@ #include #include #include +#include #include namespace dolfinx::fem { +/// @brief DOF transform kernel concept. +template +concept DofTransformKernel + = std::is_invocable_v, std::span, + std::int32_t, int>; + /// @brief Finite element cell kernel concept. /// /// Kernel functions that can be passed to an assembler for execution diff --git a/cpp/dolfinx/fem/utils.h b/cpp/dolfinx/fem/utils.h index 66c0c6d77f7..b1140b1fdb7 100644 --- a/cpp/dolfinx/fem/utils.h +++ b/cpp/dolfinx/fem/utils.h @@ -162,6 +162,11 @@ la::SparsityPattern create_sparsity_pattern(const Form& a) std::shared_ptr mesh = a.mesh(); assert(mesh); + std::shared_ptr mesh0 = a.function_spaces().at(0)->mesh(); + assert(mesh0); + std::shared_ptr mesh1 = a.function_spaces().at(1)->mesh(); + assert(mesh1); + const std::set types = a.integral_types(); if (types.find(IntegralType::interior_facet) != types.end() or types.find(IntegralType::exterior_facet) != types.end()) @@ -190,8 +195,9 @@ la::SparsityPattern create_sparsity_pattern(const Form& a) case IntegralType::cell: for (int id : ids) { - sparsitybuild::cells(pattern, a.domain(type, id), - {{dofmaps[0], dofmaps[1]}}); + sparsitybuild::cells( + pattern, {a.domain(type, id, *mesh0), a.domain(type, id, *mesh1)}, + {{dofmaps[0], dofmaps[1]}}); } break; case IntegralType::interior_facet: @@ -213,7 +219,8 @@ la::SparsityPattern create_sparsity_pattern(const Form& a) cells.reserve(facets.size() / 2); for (std::size_t i = 0; i < facets.size(); i += 2) cells.push_back(facets[i]); - sparsitybuild::cells(pattern, cells, {{dofmaps[0], dofmaps[1]}}); + sparsitybuild::cells(pattern, {cells, cells}, + {{dofmaps[0], dofmaps[1]}}); } break; default: @@ -268,10 +275,12 @@ std::vector get_constant_names(const ufcx_form& ufcx_form); /// @param[in] coefficients Coefficient fields in the form. /// @param[in] constants Spatial constants in the form. /// @param[in] subdomains Subdomain markers. -/// @param[in] mesh The mesh of the domain +/// @param[in] entity_maps The entity maps for the form. Empty for +/// single domain problems. +/// @param[in] mesh The mesh of the domain. /// /// @pre Each value in `subdomains` must be sorted by domain id. -template > +template > Form create_form_factory( const ufcx_form& ufcx_form, const std::vector>>& spaces, @@ -281,6 +290,8 @@ Form create_form_factory( IntegralType, std::vector>>>& subdomains, + const std::map>, + std::span>& entity_maps, std::shared_ptr> mesh = nullptr) { if (ufcx_form.rank != (int)spaces.size()) @@ -317,8 +328,9 @@ Form create_form_factory( mesh = spaces[0]->mesh(); for (auto& V : spaces) { - if (mesh != V->mesh()) - throw std::runtime_error("Incompatible mesh"); + if (mesh != V->mesh() and entity_maps.find(V->mesh()) == entity_maps.end()) + throw std::runtime_error( + "Incompatible mesh. entity_maps must be provided."); } if (!mesh) throw std::runtime_error("No mesh could be associated with the Form."); @@ -343,10 +355,9 @@ Form create_form_factory( // Get list of integral IDs, and load tabulate tensor into memory for // each - using kern_t = std::function::value_type*, - const int*, const std::uint8_t*)>; - std::map>> integrals; + using kern_t = std::function; + std::map>> integrals; // Attach cell kernels bool needs_facet_permutations = false; @@ -569,7 +580,7 @@ Form create_form_factory( } return Form(spaces, integrals, coefficients, constants, - needs_facet_permutations, mesh); + needs_facet_permutations, entity_maps, mesh); } /// @brief Create a Form from UFC input with coefficients and constants @@ -583,7 +594,7 @@ Form create_form_factory( /// @param[in] mesh Mesh of the domain. This is required if the form has /// no arguments, e.g. a functional. /// @return A Form -template > +template > Form create_form( const ufcx_form& ufcx_form, const std::vector>>& spaces, @@ -620,7 +631,7 @@ Form create_form( } return create_form_factory(ufcx_form, spaces, coeff_map, const_map, - subdomains, mesh); + subdomains, {}, mesh); } /// @brief Create a Form using a factory function that returns a pointer @@ -638,7 +649,7 @@ Form create_form( /// @param[in] mesh Mesh of the domain. This is required if the form has /// no arguments, e.g. a functional. /// @return A Form -template > +template > Form create_form( ufcx_form* (*fptr)(), const std::vector>>& spaces, @@ -1067,7 +1078,7 @@ void pack_coefficients(const Form& form, IntegralType integral_type, } /// @brief Create Expression from UFC -template > +template > Expression create_expression( const ufcx_expression& e, const std::vector>>& coefficients, @@ -1117,7 +1128,7 @@ Expression create_expression( /// @brief Create Expression from UFC input (with named coefficients and /// constants). -template > +template > Expression create_expression( const ufcx_expression& e, const std::map>>& @@ -1216,7 +1227,7 @@ pack_coefficients(const Expression& e, return {std::move(c), cstride}; } -/// @brief Pack constants of u into a sigle array ready for assembly. +/// @brief Pack constants of u into a single array ready for assembly. /// @warning This function is subject to change. template std::vector pack_constants(const U& u) diff --git a/cpp/dolfinx/mesh/utils.h b/cpp/dolfinx/mesh/utils.h index fcc0fce7320..663303c5f9b 100644 --- a/cpp/dolfinx/mesh/utils.h +++ b/cpp/dolfinx/mesh/utils.h @@ -533,8 +533,8 @@ std::vector locate_entities(const Mesh& mesh, int dim, /// owned boundary facet and evaluate to true for the provided geometric /// marking function. /// -/// An entity is considered marked if the marker function evaluates to true -/// for all of its vertices. +/// An entity is considered marked if the marker function evaluates to +/// true for all of its vertices. /// /// @note For vertices and edges, in parallel this function will not /// necessarily mark all entities that are on the exterior boundary. For diff --git a/python/demo/demo_static-condensation.py b/python/demo/demo_static-condensation.py index 8fa99af5957..7bcaf269a53 100644 --- a/python/demo/demo_static-condensation.py +++ b/python/demo/demo_static-condensation.py @@ -151,7 +151,7 @@ def tabulate_A(A_, w_, c_, coords_, entity_local_index, permutation=ffi.NULL): formtype = form_cpp_class(PETSc.ScalarType) # type: ignore cells = np.arange(msh.topology.index_map(msh.topology.dim).size_local) integrals = {IntegralType.cell: [(-1, tabulate_A.address, cells)]} -a_cond = Form(formtype([U._cpp_object, U._cpp_object], integrals, [], [], False, None)) +a_cond = Form(formtype([U._cpp_object, U._cpp_object], integrals, [], [], False, {}, None)) A_cond = assemble_matrix(a_cond, bcs=[bc]) A_cond.assemble() diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 47b20cc5509..be24713946c 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -136,6 +136,7 @@ def form( dtype: npt.DTypeLike = default_scalar_type, form_compiler_options: typing.Optional[dict] = None, jit_options: typing.Optional[dict] = None, + entity_maps: dict[_cpp.mesh.Mesh, np.typing.NDArray[np.int32]] = {}, ): """Create a Form or an array of Forms. @@ -144,6 +145,15 @@ def form( dtype: Scalar type to use for the compiled form. form_compiler_options: See :func:`ffcx_jit ` jit_options: See :func:`ffcx_jit `. + entity_maps: If any trial functions, test functions, or + coefficients in the form are not defined over the same mesh + as the integration domain, `entity_maps` must be supplied. + For each key (a mesh, different to the integration domain + mesh) a map should be provided relating the entities in the + integration domain mesh to the entities in the key mesh e.g. + for a key-value pair (msh, emap) in `entity_maps`, `emap[i]` + is the entity in `msh` corresponding to entity `i` in the + integration domain mesh. Returns: Compiled finite element Form. @@ -154,7 +164,6 @@ def form( data to the underlying C++ form. It dynamically create a :class:`Form` instance with an appropriate base class for the scalar type, e.g. :func:`_cpp.fem.Form_float64`. - """ if form_compiler_options is None: form_compiler_options = dict() @@ -223,6 +232,7 @@ def get_integration_domains(integral_type, subdomain): coeffs, constants, subdomains, + entity_maps, mesh, ) return Form(f, ufcx_form, code) diff --git a/python/dolfinx/wrappers/assemble.cpp b/python/dolfinx/wrappers/assemble.cpp index 6a55705f933..fdf2312b73f 100644 --- a/python/dolfinx/wrappers/assemble.cpp +++ b/python/dolfinx/wrappers/assemble.cpp @@ -75,7 +75,7 @@ void declare_discrete_operators(nb::module_& m) assert(map); std::vector c(map->size_local(), 0); std::iota(c.begin(), c.end(), 0); - dolfinx::fem::sparsitybuild::cells(sp, c, {*dofmap1, *dofmap0}); + dolfinx::fem::sparsitybuild::cells(sp, {c, c}, {*dofmap1, *dofmap0}); sp.finalize(); // Build operator diff --git a/python/dolfinx/wrappers/fem.cpp b/python/dolfinx/wrappers/fem.cpp index 471a43da1b8..5344cb5e19e 100644 --- a/python/dolfinx/wrappers/fem.cpp +++ b/python/dolfinx/wrappers/fem.cpp @@ -536,14 +536,13 @@ void declare_form(nb::module_& m, std::string type) const std::vector< std::shared_ptr>>& constants, bool needs_permutation_data, + const std::map>, + nb::ndarray, + nb::c_contig>>& entity_maps, std::shared_ptr> mesh) { - using kern_t - = std::function::value_type*, - const int*, const std::uint8_t*)>; std::map>> + std::vector>> _integrals; // Loop over kernel for each entity type @@ -561,13 +560,18 @@ void declare_form(nb::module_& m, std::string type) } } - new (fp) dolfinx::fem::Form(spaces, std::move(_integrals), - coefficients, constants, - needs_permutation_data, mesh); + std::map>, + std::span> + _entity_maps; + for (auto& [msh, map] : entity_maps) + _entity_maps.emplace(msh, std::span(map.data(), map.size())); + new (fp) dolfinx::fem::Form( + spaces, std::move(_integrals), coefficients, constants, + needs_permutation_data, _entity_maps, mesh); }, nb::arg("spaces"), nb::arg("integrals"), nb::arg("coefficients"), nb::arg("constants"), nb::arg("need_permutation_data"), - nb::arg("mesh").none()) + nb::arg("entity_maps"), nb::arg("mesh").none()) .def( "__init__", [](dolfinx::fem::Form* fp, std::uintptr_t form, @@ -582,6 +586,9 @@ void declare_form(nb::module_& m, std::string type) std::vector, nb::c_contig>>>>& subdomains, + const std::map>, + nb::ndarray, + nb::c_contig>>& entity_maps, std::shared_ptr> mesh) { std::map>, + std::span> + _entity_maps; + for (auto& [msh, map] : entity_maps) + _entity_maps.emplace(msh, std::span(map.data(), map.size())); ufcx_form* p = reinterpret_cast(form); new (fp) dolfinx::fem::Form(dolfinx::fem::create_form_factory( - *p, spaces, coefficients, constants, sd, mesh)); + *p, spaces, coefficients, constants, sd, _entity_maps, + mesh)); }, nb::arg("form"), nb::arg("spaces"), nb::arg("coefficients"), - nb::arg("constants"), nb::arg("subdomains"), nb::arg("mesh").none(), - "Create a Form from a pointer to a ufcx_form") + nb::arg("constants"), nb::arg("subdomains"), nb::arg("entity_maps"), + nb::arg("mesh").none(), "Create a Form from a pointer to a ufcx_form") .def_prop_ro("dtype", [](const dolfinx::fem::Form&) { return dolfinx_wrappers::numpy_dtype(); }) .def_prop_ro("coefficients", &dolfinx::fem::Form::coefficients) @@ -676,7 +689,7 @@ void declare_form(nb::module_& m, std::string type) ufcx_form* p = reinterpret_cast(form); return dolfinx::fem::create_form_factory(*p, spaces, coefficients, - constants, sd, mesh); + constants, sd, {}, mesh); }, nb::arg("form"), nb::arg("spaces"), nb::arg("coefficients"), nb::arg("constants"), nb::arg("subdomains"), nb::arg("mesh"), diff --git a/python/dolfinx/wrappers/petsc.cpp b/python/dolfinx/wrappers/petsc.cpp index a7f16d4f82e..9167a539880 100644 --- a/python/dolfinx/wrappers/petsc.cpp +++ b/python/dolfinx/wrappers/petsc.cpp @@ -70,7 +70,7 @@ void declare_petsc_discrete_operators(nb::module_& m) assert(map); std::vector c(map->size_local(), 0); std::iota(c.begin(), c.end(), 0); - dolfinx::fem::sparsitybuild::cells(sp, c, {*dofmap1, *dofmap0}); + dolfinx::fem::sparsitybuild::cells(sp, {c, c}, {*dofmap1, *dofmap0}); sp.finalize(); // Build operator @@ -111,7 +111,7 @@ void declare_petsc_discrete_operators(nb::module_& m) assert(map); std::vector c(map->size_local(), 0); std::iota(c.begin(), c.end(), 0); - dolfinx::fem::sparsitybuild::cells(sp, c, {*dofmap1, *dofmap0}); + dolfinx::fem::sparsitybuild::cells(sp, {c, c}, {*dofmap1, *dofmap0}); sp.finalize(); // Build operator diff --git a/python/test/unit/fem/test_assemble_submesh.py b/python/test/unit/fem/test_assemble_submesh.py index 6dcb0ed110b..0c80827336d 100644 --- a/python/test/unit/fem/test_assemble_submesh.py +++ b/python/test/unit/fem/test_assemble_submesh.py @@ -22,6 +22,7 @@ create_unit_square, locate_entities, locate_entities_boundary, + meshtags, ) @@ -40,7 +41,6 @@ def assemble(mesh, space, k): bc_func = fem.Function(V) bc_func.interpolate(lambda x: np.sin(x[0])) - bc = fem.dirichletbc(bc_func, dofs) A = fem.assemble_matrix(a, bcs=[bc]) @@ -116,3 +116,71 @@ def test_submesh_facet_assembly(n, k, space, ghost_mode): ) assert b_submesh.norm() == pytest.approx(b_square_mesh.norm()) assert np.isclose(s_submesh, s_square_mesh) + + +@pytest.mark.parametrize("n", [2, 6]) +@pytest.mark.parametrize("k", [1, 4]) +@pytest.mark.parametrize("space", ["Lagrange", "Discontinuous Lagrange"]) +@pytest.mark.parametrize("ghost_mode", [GhostMode.none, GhostMode.shared_facet]) +def test_mixed_dom_codim_0(n, k, space, ghost_mode): + """Test assembling a form where the trial and test functions + are defined on different meshes""" + + msh = create_rectangle( + MPI.COMM_WORLD, ((0.0, 0.0), (2.0, 1.0)), (2 * n, n), ghost_mode=ghost_mode + ) + + # Locate cells in left half of mesh and create mesh tags + tdim = msh.topology.dim + cells = locate_entities(msh, tdim, lambda x: x[0] <= 1.0) + perm = np.argsort(cells) + tag = 1 + values = np.full_like(cells, tag, dtype=np.intc) + ct = meshtags(msh, tdim, cells[perm], values[perm]) + + # Create integration measure on the mesh + dx_msh = ufl.Measure("dx", domain=msh, subdomain_data=ct) + + # Create a submesh of the left half of the mesh + smsh, smsh_to_msh = create_submesh(msh, tdim, cells)[:2] + + # Define function spaces over the mesh and submesh + V_msh = fem.functionspace(msh, (space, k)) + V_smsh = fem.functionspace(smsh, (space, k)) + + # Trial and test functions on the mesh + u, v = ufl.TrialFunction(V_msh), ufl.TestFunction(V_msh) + + # Test function on the submesh + w = ufl.TestFunction(V_smsh) + + def ufl_form(u, v, dx): + return ufl.inner(u, v) * dx + + # Define form to compare to and assemble + a = fem.form(ufl_form(u, v, dx_msh(tag))) + A = fem.assemble_matrix(a) + A.scatter_reverse() + + # Assemble a mixed-domain form, taking smsh to be the integration domain + # Entity maps must map cells in smsh (the integration domain mesh) to + # cells in msh + entity_maps = {msh._cpp_object: np.array(smsh_to_msh, dtype=np.int32)} + a0 = fem.form(ufl_form(u, w, ufl.dx(smsh)), entity_maps=entity_maps) + A0 = fem.assemble_matrix(a0) + A0.scatter_reverse() + assert np.isclose(A0.squared_norm(), A.squared_norm()) + + # Now assemble a mixed-domain form using msh as integration domain + # Entity maps must map cells in msh (the integration domain mesh) to + # cells in smsh + cell_imap = msh.topology.index_map(tdim) + num_cells = cell_imap.size_local + cell_imap.num_ghosts + msh_to_smsh = np.full(num_cells, -1) + msh_to_smsh[smsh_to_msh] = np.arange(len(smsh_to_msh)) + entity_maps = {smsh._cpp_object: np.array(msh_to_smsh, dtype=np.int32)} + + a1 = fem.form(ufl_form(u, w, dx_msh(tag)), entity_maps=entity_maps) + A1 = fem.assemble_matrix(a1) + A1.scatter_reverse() + assert np.isclose(A1.squared_norm(), A.squared_norm()) diff --git a/python/test/unit/fem/test_custom_jit_kernels.py b/python/test/unit/fem/test_custom_jit_kernels.py index 6b341348e67..12c8748da04 100644 --- a/python/test/unit/fem/test_custom_jit_kernels.py +++ b/python/test/unit/fem/test_custom_jit_kernels.py @@ -98,9 +98,9 @@ def test_numba_assembly(dtype): ] } formtype = form_cpp_class(dtype) - a = Form(formtype([V._cpp_object, V._cpp_object], integrals, [], [], False, None)) + a = Form(formtype([V._cpp_object, V._cpp_object], integrals, [], [], False, {}, None)) integrals = {IntegralType.cell: [(-1, k1.address, cells)]} - L = Form(formtype([V._cpp_object], integrals, [], [], False, None)) + L = Form(formtype([V._cpp_object], integrals, [], [], False, {}, None)) A = dolfinx.fem.assemble_matrix(a) A.scatter_reverse() @@ -130,7 +130,7 @@ def test_coefficient(dtype): num_cells = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts integrals = {IntegralType.cell: [(1, k1.address, np.arange(num_cells, dtype=np.int32))]} formtype = form_cpp_class(dtype) - L = Form(formtype([V._cpp_object], integrals, [vals._cpp_object], [], False, None)) + L = Form(formtype([V._cpp_object], integrals, [vals._cpp_object], [], False, {}, None)) b = dolfinx.fem.assemble_vector(L) b.scatter_reverse(la.InsertMode.add) @@ -259,11 +259,13 @@ def test_cffi_assembly(): ptrA = ffi.cast("intptr_t", ffi.addressof(lib, "tabulate_tensor_poissonA")) integrals = {IntegralType.cell: [(-1, ptrA, cells)]} - a = Form(_cpp.fem.Form_float64([V._cpp_object, V._cpp_object], integrals, [], [], False, None)) + a = Form( + _cpp.fem.Form_float64([V._cpp_object, V._cpp_object], integrals, [], [], False, {}, None) + ) ptrL = ffi.cast("intptr_t", ffi.addressof(lib, "tabulate_tensor_poissonL")) integrals = {IntegralType.cell: [(-1, ptrL, cells)]} - L = Form(_cpp.fem.Form_float64([V._cpp_object], integrals, [], [], False, None)) + L = Form(_cpp.fem.Form_float64([V._cpp_object], integrals, [], [], False, {}, None)) A = fem.assemble_matrix(a) A.scatter_reverse()