Fix sub index map destination rank computation (#3099)

* Fix dest computation

* Add test

* Fix test

* Use pytest skip

* Ruff

cpp/dolfinx/common/IndexMap.cpp

@@ -205,7 +205,7 @@ compute_submap_indices(const IndexMap& imap,
   // `recv_indices` will necessarily be in `indices` on this process, and thus
   // other processes must own them in the submap.
 
-  std::vector<int> recv_owners(send_disp.back()), submap_dest;
+  std::vector<int> recv_owners(send_disp.back());
   const int rank = dolfinx::MPI::rank(imap.comm());
   {
     // Flag to track if the owner of any indices have changed in the submap
@@ -233,11 +233,7 @@ compute_submap_indices(const IndexMap& imap,
           // this process remains its owner in the submap. Otherwise,
           // add the process that sent it to the list of possible owners.
           if (is_in_submap[idx_local])
-          {
             global_idx_to_possible_owner.push_back({idx, rank});
-            // Add the sending process as a destination rank in the submap
-            submap_dest.push_back(dest[i]);
-          }
           else
           {
             owners_changed = true;
@@ -252,10 +248,6 @@ compute_submap_indices(const IndexMap& imap,
 
     std::sort(global_idx_to_possible_owner.begin(),
               global_idx_to_possible_owner.end());
-    std::sort(submap_dest.begin(), submap_dest.end());
-    submap_dest.erase(std::unique(submap_dest.begin(), submap_dest.end()),
-                      submap_dest.end());
-    submap_dest.shrink_to_fit();
 
     // Choose the submap owner for each index in `recv_indices`
     std::vector<int> send_owners;
@@ -377,6 +369,12 @@ compute_submap_indices(const IndexMap& imap,
     submap_ghost = std::move(submap_ghost1);
   }
 
+  // Compute submap destination ranks
+  // FIXME Remove call to NBX
+  std::vector<int> submap_dest
+      = dolfinx::MPI::compute_graph_edges_nbx(imap.comm(), submap_src);
+  std::sort(submap_dest.begin(), submap_dest.end());
+
   return {std::move(submap_owned), std::move(submap_ghost),
           std::move(submap_ghost_owners), std::move(submap_src),
           std::move(submap_dest)};

python/test/unit/common/test_index_map.py

@@ -9,6 +9,7 @@
 from mpi4py import MPI
 
 import numpy as np
+import pytest
 
 import dolfinx
 from dolfinx import cpp as _cpp
@@ -192,3 +193,54 @@ def test_create_submap_owner_change():
         np.arange(sub_imap.size_local + sub_imap.num_ghosts, dtype=np.int32)
     )
     assert np.array_equal(global_indices, np.arange(comm.rank * 2, comm.rank * 2 + 3))
+
+
+def test_sub_index_map_multiple_possible_owners():
+    """Check that creating a submap doesn't crash when an index need to change owner and
+    there are multiple possible new owners"""
+    comm = MPI.COMM_WORLD
+
+    if comm.size < 3:
+        pytest.skip("Test requires 3 or more processes")
+
+    # Create an index map with an index on process 2 that is ghosted by processes 0 and 1
+    if comm.rank == 0:
+        local_size = 1
+        ghosts = np.array([2], dtype=np.int64)
+        owners = np.array([2], dtype=np.int32)
+        submap_indices = np.array([0, 1], dtype=np.int32)
+        # NOTE: This assumes that the lowest ranking process takes ownership of the index
+        # in the submap
+        submap_size_local_expected = 2
+        submap_num_ghosts_expected = 0
+    elif comm.rank == 1:
+        local_size = 1
+        ghosts = np.array([2], dtype=np.int64)
+        owners = np.array([2], dtype=np.int32)
+        submap_indices = np.array([0, 1], dtype=np.int32)
+        submap_size_local_expected = 1
+        submap_num_ghosts_expected = 1
+    elif comm.rank == 2:
+        local_size = 1
+        ghosts = np.array([], dtype=np.int64)
+        owners = np.array([], dtype=np.int32)
+        submap_indices = np.array([], dtype=np.int32)
+        submap_size_local_expected = 0
+        submap_num_ghosts_expected = 0
+    else:
+        local_size = 0
+        ghosts = np.array([], dtype=np.int64)
+        owners = np.array([], dtype=np.int32)
+        submap_indices = np.array([], dtype=np.int32)
+        submap_size_local_expected = 0
+        submap_num_ghosts_expected = 0
+
+    imap = dolfinx.common.IndexMap(comm, local_size, ghosts, owners)
+
+    # Create a submap where both processes 0 and 1 include the index on process 2,
+    # but process 2 does not include it
+    sub_imap = _cpp.common.create_sub_index_map(imap, submap_indices, True)[0]
+
+    assert sub_imap.size_global == 3
+    assert sub_imap.size_local == submap_size_local_expected
+    assert sub_imap.num_ghosts == submap_num_ghosts_expected