Fix create_submesh by generalising sub index map creation

cpp/dolfinx/common/IndexMap.h

@@ -54,6 +54,27 @@ stack_index_maps(
     const std::vector<
         std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
 
+/// @brief Create a new index map from a subset of indices in an
+/// existing index map.
+///
+/// @param[in] imap Parent map to create a new sub-map from.
+/// @param[in] indices Local indices in `imap` (owned and ghost) to
+/// include in the new index map.
+/// @param[in] allow_owner_change If `true`, indices that are not
+/// included in `indices` by their owning process can be included in
+/// `indices` by processes that ghost the indices to be included in the
+/// new submap. These indices will be owned by one of the sharing
+/// processes in the submap. If `false`, and exception is raised if an
+/// index is included by a sharing process and not by the owning
+/// process.
+/// @return The (i) new index map and (ii) a map from local indices in
+/// the submap to local indices in the original (this) map.
+/// @pre `indices` must be sorted and must not contain duplicates.
+std::pair<IndexMap, std::vector<std::int32_t>>
+create_sub_index_map(const IndexMap& imap,
+                     std::span<const std::int32_t> indices,
+                     bool allow_owner_change = false);
+
 /// This class represents the distribution index arrays across
 /// processes. An index array is a contiguous collection of N+1 indices
 /// [0, 1, . . ., N] that are distributed across M processes. On a given
@@ -173,22 +194,6 @@ class IndexMap
   /// index is `owners()[i]`.
   const std::vector<int>& owners() const { return _owners; }
 
-  /// @brief Create new index map from a subset of indices in this index
-  /// map.
-  ///
-  /// The order of the owned indices is preserved, with new map
-  /// effectively a 'compressed' map.
-  ///
-  /// @param[in] indices Local indices in the map that should appear in
-  /// the new index map. All indices must be owned, i.e. indices must be
-  /// less than `this->size_local()`.
-  /// @pre `indices` must be sorted and contain no duplicates.
-  /// @return The (i) new index map and (ii) a map from the ghost
-  /// position in the new map to the ghost position in the original
-  /// (this) map
-  std::pair<IndexMap, std::vector<std::int32_t>>
-  create_submap(std::span<const std::int32_t> indices) const;
-
   /// @todo Aim to remove this function?
   ///
   /// @brief Compute map from each local (owned) index to the set of

cpp/dolfinx/fem/DofMap.cpp

@@ -50,59 +50,43 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   dofs_view.erase(std::unique(dofs_view.begin(), dofs_view.end()),
                   dofs_view.end());
 
-  // Compute sizes
-  const std::int32_t num_owned_view = dofmap_view.index_map->size_local();
-
   // Get block size
   int bs_view = dofmap_view.index_map_bs();
 
-  const auto it = std::lower_bound(dofs_view.begin(), dofs_view.end(),
-                                   bs_view * num_owned_view);
-
   // Create sub-index map
   std::shared_ptr<common::IndexMap> index_map;
-  std::vector<std::int32_t> ghost_new_to_old;
+  // Map from indices in the sub-index map to indices in the original
+  // index map
+  std::vector<std::int32_t> sub_imap_to_imap;
   if (bs_view == 1)
   {
-    std::span<std::int32_t> indices(dofs_view.data(),
-                                    std::distance(dofs_view.begin(), it));
-    auto [_index_map, gmap] = dofmap_view.index_map->create_submap(indices);
+    auto [_index_map, _sub_imap_to_imap] = dolfinx::common::create_sub_index_map(
+        *dofmap_view.index_map, dofs_view);
     index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
-    ghost_new_to_old = std::move(gmap);
+    sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }
   else
   {
     std::vector<std::int32_t> indices;
     indices.reserve(dofs_view.size());
-    std::transform(dofs_view.begin(), it, std::back_inserter(indices),
+    std::transform(dofs_view.begin(), dofs_view.end(),
+                   std::back_inserter(indices),
                    [bs_view](auto idx) { return idx / bs_view; });
-    indices.erase(std::unique(indices.begin(), indices.end()), indices.end());
-    auto [_index_map, gmap] = dofmap_view.index_map->create_submap(indices);
+    auto [_index_map, _sub_imap_to_imap]
+        = dolfinx::common::create_sub_index_map(*dofmap_view.index_map, indices);
     index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
-    ghost_new_to_old = std::move(gmap);
+    sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }
 
-  // Create map from dof in view to new dof index
+  // Create a map from old dofs to new dofs
   std::vector<std::int32_t> old_to_new(dofs_view.back() + bs_view, -1);
+  for (std::size_t new_idx = 0; new_idx < sub_imap_to_imap.size(); ++new_idx)
   {
-    // old-to-new map for owned dofs
-    std::int32_t count = 0;
-    for (auto dof_old = dofs_view.begin(); dof_old != it; ++dof_old)
-      old_to_new[*dof_old] = count++;
-
-    // old-to-new map for ghost dofs
-    const std::int32_t local_size_new = index_map->size_local();
-    for (auto itp_old = ghost_new_to_old.begin();
-         itp_old != ghost_new_to_old.end(); ++itp_old)
+    for (int k = 0; k < bs_view; ++k)
     {
-      std::int32_t map_pos_new
-          = local_size_new + std::distance(ghost_new_to_old.begin(), itp_old);
-      std::int32_t idx = bs_view * (num_owned_view + *itp_old);
-      for (int k = 0; k < bs_view; ++k)
-      {
-        assert(idx + k < (int)old_to_new.size());
-        old_to_new[idx + k] = map_pos_new;
-      }
+      std::int32_t old_idx = sub_imap_to_imap[new_idx] * bs_view + k;
+      assert(old_idx < (int)old_to_new.size());
+      old_to_new[old_idx] = new_idx;
     }
   }
 

cpp/dolfinx/mesh/Geometry.h

@@ -328,22 +328,15 @@ create_subgeometry(const Topology& topology, const Geometry<T>& geometry,
   // Get the sub-geometry dofs owned by this process
   auto x_index_map = geometry.index_map();
   assert(x_index_map);
-  std::vector<std::int32_t> subx_to_x_dofmap
-      = common::compute_owned_indices(sub_x_dofs, *x_index_map);
+
   std::shared_ptr<common::IndexMap> sub_x_dof_index_map;
+  std::vector<std::int32_t> subx_to_x_dofmap;
   {
     std::pair<common::IndexMap, std::vector<int32_t>> map_data
-        = x_index_map->create_submap(subx_to_x_dofmap);
+        = dolfinx::common::create_sub_index_map(*x_index_map, sub_x_dofs, true);
     sub_x_dof_index_map
         = std::make_shared<common::IndexMap>(std::move(map_data.first));
-
-    // Create a map from the dofs in the sub-geometry to the geometry
-    subx_to_x_dofmap.reserve(sub_x_dof_index_map->size_local()
-                             + sub_x_dof_index_map->num_ghosts());
-    std::transform(map_data.second.begin(), map_data.second.end(),
-                   std::back_inserter(subx_to_x_dofmap),
-                   [offset = x_index_map->size_local()](auto x_dof_index)
-                   { return offset + x_dof_index; });
+    subx_to_x_dofmap = std::move(map_data.second);
   }
 
   // Create sub-geometry coordinates

cpp/dolfinx/mesh/Topology.cpp

@@ -1160,32 +1160,23 @@ mesh::create_subtopology(const Topology& topology, int dim,
   // TODO Call common::get_owned_indices here? Do we want to
   // support `entities` possibly having a ghost on one process that is
   // not in `entities` on the owning process?
-  // TODO: Should entities still be ghosted in the sub-topology even if
-  // they are not in the `entities` list? If this is not desirable,
-  // create_submap needs to be changed
 
   // Create a map from an entity in the sub-topology to the
   // corresponding entity in the topology, and create an index map
-  std::vector<int32_t> subentities;
   std::shared_ptr<common::IndexMap> submap;
+  std::vector<int32_t> subentities;
   {
-    // Entities in the sub-topology that are owned by this process
-    auto entity_map = topology.index_map(dim);
-    assert(entity_map);
-    std::copy_if(
-        entities.begin(), entities.end(), std::back_inserter(subentities),
-        [size = entity_map->size_local()](std::int32_t e) { return e < size; });
-
+    // FIXME Make this an input requirement?
+    std::vector<std::int32_t> unique_entities(entities.begin(), entities.end());
+    std::sort(unique_entities.begin(), unique_entities.end());
+    unique_entities.erase(
+        std::unique(unique_entities.begin(), unique_entities.end()),
+        unique_entities.end());
     std::pair<common::IndexMap, std::vector<int32_t>> map_data
-        = entity_map->create_submap(subentities);
+        = dolfinx::common::create_sub_index_map(*topology.index_map(dim),
+                                              unique_entities);
     submap = std::make_shared<common::IndexMap>(std::move(map_data.first));
-
-    // Add ghost entities to subentities
-    subentities.reserve(submap->size_local() + submap->num_ghosts());
-    std::transform(map_data.second.begin(), map_data.second.end(),
-                   std::back_inserter(subentities),
-                   [offset = entity_map->size_local()](auto entity_index)
-                   { return offset + entity_index; });
+    subentities = std::move(map_data.second);
   }
 
   // Get the vertices in the sub-topology. Use subentities
@@ -1195,26 +1186,17 @@ mesh::create_subtopology(const Topology& topology, int dim,
   // Get the vertices in the sub-topology owned by this process
   auto map0 = topology.index_map(0);
   assert(map0);
-  std::vector<std::int32_t> indices
-      = compute_incident_entities(topology, subentities, dim, 0);
-  std::sort(indices.begin(), indices.end());
-  std::vector<std::int32_t> subvertices0
-      = common::compute_owned_indices(indices, *map0);
 
   // Create map from the vertices in the sub-topology to the vertices in the
   // parent topology, and an index map
   std::shared_ptr<common::IndexMap> submap0;
+  std::vector<int32_t> subvertices0;
   {
     std::pair<common::IndexMap, std::vector<int32_t>> map_data
-        = map0->create_submap(subvertices0);
+        = dolfinx::common::create_sub_index_map(
+            *map0, compute_incident_entities(topology, subentities, dim, 0), true);
     submap0 = std::make_shared<common::IndexMap>(std::move(map_data.first));
-
-    // Add ghost vertices to the map
-    subvertices0.reserve(submap0->size_local() + submap0->num_ghosts());
-    std::transform(map_data.second.begin(), map_data.second.end(),
-                   std::back_inserter(subvertices0),
-                   [offset = map0->size_local()](std::int32_t vertex_index)
-                   { return offset + vertex_index; });
+    subvertices0 = std::move(map_data.second);
   }
 
   // Sub-topology vertex-to-vertex connectivity (identity)

python/dolfinx/wrappers/common.cpp

@@ -133,19 +133,7 @@ void common(nb::module_& m)
                 std::span<std::int64_t>(global.data(), global.size()));
             return global;
           },
-          nb::arg("local"))
-      .def(
-          "create_submap",
-          [](const dolfinx::common::IndexMap& self,
-             nb::ndarray<const std::int32_t, nb::ndim<1>, nb::c_contig>
-                 entities)
-          {
-            auto [map, ghosts] = self.create_submap(
-                std::span(entities.data(), entities.size()));
-            return std::pair(std::move(map),
-                             dolfinx_wrappers::as_nbarray(std::move(ghosts)));
-          },
-          nb::arg("entities"));
+          nb::arg("local"));
 
   // dolfinx::common::Timer
   nb::class_<dolfinx::common::Timer>(m, "Timer", "Timer class")
@@ -184,5 +172,18 @@ void common(nb::module_& m)
         dolfinx::init_logging(args.size(), argv.data());
       },
       nb::arg("args"));
+
+  m.def(
+      "create_sub_index_map",
+      [](const dolfinx::common::IndexMap& imap,
+         nb::ndarray<const std::int32_t, nb::ndim<1>, nb::c_contig> indices,
+         bool allow_owner_change)
+      {
+        auto [map, submap_to_map] = dolfinx::common::create_sub_index_map(
+            imap, std::span(indices.data(), indices.size()), allow_owner_change);
+        return std::pair(std::move(map), dolfinx_wrappers::as_nbarray(
+                                             std::move(submap_to_map)));
+      },
+      nb::arg("index_map"), nb::arg("indices"), nb::arg("allow_owner_change"));
 }
 } // namespace dolfinx_wrappers