Add option to preserve sub-IndexMap ghost ordering (#3091)

* Work on index map order preservation

* Doc fix

* Tidy up names

* Doc update

* Remove stray statement

* Move re-ordering

---------

Co-authored-by: Joe Dean <jpd62@cam.ac.uk>

cpp/dolfinx/common/IndexMap.cpp

@@ -164,6 +164,8 @@ communicate_ghosts_to_owners(MPI_Comm comm, std::span<const int> src,
 /// @param[in] imap An index map.
 /// @param[in] indices List of entity indices (indices local to the
 /// process).
+/// @param[in] order Control the order in which ghost indices appear in
+/// the new map.
 /// @param[in] allow_owner_change Allows indices that are not included
 /// by their owning process but included on sharing processes to be
 /// included in the submap. These indices will be owned by one of the
@@ -176,7 +178,7 @@ std::tuple<std::vector<std::int32_t>, std::vector<std::int32_t>,
            std::vector<int>, std::vector<int>, std::vector<int>>
 compute_submap_indices(const IndexMap& imap,
                        std::span<const std::int32_t> indices,
-                       bool allow_owner_change)
+                       IndexMapOrder order, bool allow_owner_change)
 {
   std::span<const int> src = imap.src();
   std::span<const int> dest = imap.dest();
@@ -350,6 +352,31 @@ compute_submap_indices(const IndexMap& imap,
                    submap_src.end());
   submap_src.shrink_to_fit();
 
+  // If required, preserve the order of the ghost indices
+  if (order == IndexMapOrder::preserve)
+  {
+    // Build (old postion, new position) list for ghosts and sort
+    std::vector<std::pair<std::int32_t, std::int32_t>> pos;
+    pos.reserve(submap_ghost.size());
+    for (std::int32_t idx : submap_ghost)
+      pos.emplace_back(idx, pos.size());
+    std::sort(pos.begin(), pos.end());
+
+    // Order ghosts in the sub-map by their position in the parent map
+    std::vector<int> submap_ghost_owners1;
+    submap_ghost_owners1.reserve(submap_ghost_owners.size());
+    std::vector<std::int32_t> submap_ghost1;
+    submap_ghost1.reserve(submap_ghost.size());
+    for (auto [_, idx] : pos)
+    {
+      submap_ghost_owners1.push_back(submap_ghost_owners[idx]);
+      submap_ghost1.push_back(submap_ghost[idx]);
+    }
+
+    submap_ghost_owners = std::move(submap_ghost_owners1);
+    submap_ghost = std::move(submap_ghost1);
+  }
+
   return {std::move(submap_owned), std::move(submap_ghost),
           std::move(submap_ghost_owners), std::move(submap_src),
           std::move(submap_dest)};
@@ -725,14 +752,14 @@ common::stack_index_maps(
 std::pair<IndexMap, std::vector<std::int32_t>>
 common::create_sub_index_map(const IndexMap& imap,
                              std::span<const std::int32_t> indices,
-                             bool allow_owner_change)
+                             IndexMapOrder order, bool allow_owner_change)
 {
   // Compute the owned, ghost, and ghost owners of submap indices.
   // NOTE: All indices are local and numbered w.r.t. the original (imap)
   // index map
-  const auto [submap_owned, submap_ghost, submap_ghost_owners, submap_src,
-              submap_dest]
-      = compute_submap_indices(imap, indices, allow_owner_change);
+  auto [submap_owned, submap_ghost, submap_ghost_owners, submap_src,
+        submap_dest]
+      = compute_submap_indices(imap, indices, order, allow_owner_change);
 
   // Compute submap offset for this rank
   std::int64_t submap_local_size = submap_owned.size();

cpp/dolfinx/common/IndexMap.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2015-2022 Chris Richardson, Garth N. Wells and Igor Baratta
+// Copyright (C) 2015-2024 Chris Richardson, Garth N. Wells and Igor Baratta
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -19,6 +19,14 @@ namespace dolfinx::common
 // Forward declaration
 class IndexMap;
 
+/// Enum to control preservation of ghost index ordering in
+/// sub-IndexMaps.
+enum class IndexMapOrder : bool
+{
+  preserve = true, ///< Preserve the ordering of ghost indices
+  any = false      ///< Allow arbitrary ordering of ghost indices in sub-maps
+};
+
 /// @brief Given a sorted vector of indices (local numbering, owned or
 /// ghost) and an index map, this function returns the indices owned by
 /// this process, including indices that might have been in the list of
@@ -60,6 +68,8 @@ stack_index_maps(
 /// @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] order Control the order in which ghost indices appear in
+/// the new 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
@@ -70,10 +80,9 @@ stack_index_maps(
 /// @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);
+std::pair<IndexMap, std::vector<std::int32_t>> create_sub_index_map(
+    const IndexMap& imap, std::span<const std::int32_t> indices,
+    IndexMapOrder order = IndexMapOrder::any, bool allow_owner_change = false);
 
 /// This class represents the distribution index arrays across
 /// processes. An index array is a contiguous collection of `N+1`
@@ -227,18 +236,19 @@ class IndexMap
 
   /// @brief Returns the imbalance of the current IndexMap.
   ///
-  /// The imbalance is a measure of load balancing across all processes, defined
-  /// as the maximum number of indices on any process divided by the average
-  /// number of indices per process. This function calculates the imbalance
-  /// separately for owned indices and ghost indices and returns them as a
-  /// std::array<double, 2>. If the total number of owned or ghost indices is
-  /// zero, the respective entry in the array is set to -1.
+  /// The imbalance is a measure of load balancing across all processes,
+  /// defined as the maximum number of indices on any process divided by
+  /// the average number of indices per process. This function
+  /// calculates the imbalance separately for owned indices and ghost
+  /// indices and returns them as a std::array<double, 2>. If the total
+  /// number of owned or ghost indices is zero, the respective entry in
+  /// the array is set to -1.
   ///
-  /// @note This is a collective operation and must be called by all processes
-  /// in the communicator associated with the IndexMap.
+  /// @note This is a collective operation and must be called by all
+  /// processes in the communicator associated with the IndexMap.
   ///
-  /// @return An array containing the imbalance in owned indices
-  /// (first element) and the imbalance in ghost indices (second element).
+  /// @return An array containing the imbalance in owned indices (first
+  /// element) and the imbalance in ghost indices (second element).
   std::array<double, 2> imbalance() const;
 
 private:

cpp/dolfinx/fem/DofMap.cpp

@@ -60,9 +60,8 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   std::vector<std::int32_t> sub_imap_to_imap;
   if (bs_view == 1)
   {
-    auto [_index_map, _sub_imap_to_imap]
-        = dolfinx::common::create_sub_index_map(*dofmap_view.index_map,
-                                                dofs_view);
+    auto [_index_map, _sub_imap_to_imap] = common::create_sub_index_map(
+        *dofmap_view.index_map, dofs_view, common::IndexMapOrder::preserve);
     index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
     sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }
@@ -73,9 +72,8 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
     std::transform(dofs_view.begin(), dofs_view.end(),
                    std::back_inserter(indices),
                    [bs_view](auto idx) { return idx / bs_view; });
-    auto [_index_map, _sub_imap_to_imap]
-        = dolfinx::common::create_sub_index_map(*dofmap_view.index_map,
-                                                indices);
+    auto [_index_map, _sub_imap_to_imap] = common::create_sub_index_map(
+        *dofmap_view.index_map, indices, common::IndexMapOrder::preserve);
     index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
     sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }

cpp/dolfinx/mesh/Geometry.h

@@ -505,8 +505,8 @@ create_subgeometry(const Topology& topology, const Geometry<T>& geometry,
   std::shared_ptr<common::IndexMap> sub_x_dof_index_map;
   std::vector<std::int32_t> subx_to_x_dofmap;
   {
-    auto [map, new_to_old]
-        = common::create_sub_index_map(*x_index_map, sub_x_dofs, true);
+    auto [map, new_to_old] = common::create_sub_index_map(
+        *x_index_map, sub_x_dofs, common::IndexMapOrder::any, true);
     sub_x_dof_index_map = std::make_shared<common::IndexMap>(std::move(map));
     subx_to_x_dofmap = std::move(new_to_old);
   }

cpp/dolfinx/mesh/Topology.cpp

@@ -1376,7 +1376,7 @@ mesh::create_subtopology(const Topology& topology, int dim,
     std::pair<common::IndexMap, std::vector<int32_t>> map_data
         = common::create_sub_index_map(
             *map0, compute_incident_entities(topology, subentities, dim, 0),
-            true);
+            common::IndexMapOrder::any, true);
     submap0 = std::make_shared<common::IndexMap>(std::move(map_data.first));
     subvertices0 = std::move(map_data.second);
   }

python/dolfinx/wrappers/common.cpp

@@ -189,7 +189,7 @@ void common(nb::module_& m)
       {
         auto [map, submap_to_map] = dolfinx::common::create_sub_index_map(
             imap, std::span(indices.data(), indices.size()),
-            allow_owner_change);
+            dolfinx::common::IndexMapOrder::any, allow_owner_change);
         return std::pair(std::move(map), dolfinx_wrappers::as_nbarray(
                                              std::move(submap_to_map)));
       },

python/test/unit/fem/test_function_space.py

@@ -168,6 +168,8 @@ def test_collapse(W, V):
         Function(W.sub(1))
 
     Ws = [W.sub(i).collapse() for i in range(W.num_sub_spaces)]
+    for Wi, _ in Ws:
+        assert np.allclose(Wi.dofmap.index_map.ghosts, W.dofmap.index_map.ghosts)
 
     msh = W.mesh
     cell_imap = msh.topology.index_map(msh.topology.dim)
@@ -181,9 +183,9 @@ def test_collapse(W, V):
                 new_to_old = Ws[k][1]
                 assert dof * bs + k == new_to_old[new_dof]
 
-    f_0 = Function(Ws[0][0])
-    f_1 = Function(V)
-    assert f_0.x.index_map.size_global == f_1.x.index_map.size_global
+    f0 = Function(Ws[0][0])
+    f1 = Function(V)
+    assert f0.x.index_map.size_global == f1.x.index_map.size_global
 
 
 def test_argument_equality(mesh, V, V2, W, W2):