Merge branch 'main' into release

Author: jhale

.github/workflows/oneapi.yml

@@ -109,4 +109,4 @@ jobs:
         run: |
           . /opt/intel/oneapi/setvars.sh
           cd python/test/unit
-          mpiexec -n 2 pytest .
+          mpiexec -n 2 pytest -vs .

cpp/dolfinx/common/IndexMap.cpp

@@ -947,3 +947,35 @@ const std::vector<int>& IndexMap::dest() const noexcept { return _dest; }
 //-----------------------------------------------------------------------------
 bool IndexMap::overlapped() const noexcept { return _overlapping; }
 //-----------------------------------------------------------------------------
+std::array<double, 2> IndexMap::imbalance() const
+{
+  std::array<double, 2> imbalance{-1., -1.};
+  std::array<std::int32_t, 2> max_count;
+  std::array<std::int32_t, 2> local_sizes
+      = {static_cast<std::int32_t>(_local_range[1] - _local_range[0]),
+         static_cast<std::int32_t>(_ghosts.size())};
+
+  // Find the maximum number of owned indices and the maximum number of ghost
+  // indices across all processes.
+  MPI_Allreduce(local_sizes.data(), max_count.data(), 2,
+                dolfinx::MPI::mpi_type<std::int32_t>(), MPI_MAX, _comm.comm());
+
+  std::int32_t total_num_ghosts = 0;
+  MPI_Allreduce(&local_sizes[1], &total_num_ghosts, 1,
+                dolfinx::MPI::mpi_type<std::int32_t>(), MPI_SUM, _comm.comm());
+
+  // Compute the average number of owned and ghost indices per process.
+  int comm_size = dolfinx::MPI::size(_comm.comm());
+  double avg_owned = static_cast<double>(_size_global) / comm_size;
+  double avg_ghosts = static_cast<double>(total_num_ghosts) / comm_size;
+
+  // Compute the imbalance by dividing the maximum number of indices by the
+  // corresponding average.
+  if (avg_owned > 0)
+    imbalance[0] = max_count[0] / avg_owned;
+  if (avg_ghosts > 0)
+    imbalance[1] = max_count[1] / avg_ghosts;
+
+  return imbalance;
+}
+//-----------------------------------------------------------------------------

cpp/dolfinx/common/IndexMap.h

@@ -228,6 +228,22 @@ class IndexMap
   /// false.
   bool overlapped() const noexcept;
 
+  /// @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.
+  ///
+  /// @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).
+  std::array<double, 2> imbalance() const;
+
 private:
   // Range of indices (global) owned by this process
   std::array<std::int64_t, 2> _local_range;

cpp/dolfinx/geometry/BoundingBoxTree.h

@@ -328,7 +328,11 @@ class BoundingBoxTree
     const int mpi_size = dolfinx::MPI::size(comm);
 
     // Send root node coordinates to all processes
-    std::array<T, 6> send_bbox = {0, 0, 0, 0, 0, 0};
+    // This is to counteract the fact that a process might have 0 bounding box
+    // causing false positives on process collisions around (0,0,0)
+    constexpr T max_val = std::numeric_limits<T>::max();
+    std::array<T, 6> send_bbox
+        = {max_val, max_val, max_val, max_val, max_val, max_val};
     if (num_bboxes() > 0)
       std::copy_n(std::prev(_bbox_coordinates.end(), 6), 6, send_bbox.begin());
     std::vector<T> recv_bbox(mpi_size * 6);

cpp/dolfinx/geometry/utils.h

@@ -48,6 +48,9 @@ std::vector<T> shortest_vector(const mesh::Mesh<T>& mesh, int dim,
   {
     for (std::size_t e = 0; e < entities.size(); e++)
     {
+      // Check that we have sent in valid entities, i.e. that they exist in the
+      // local dofmap. One gets a cryptical memory segfault if entities is -1
+      assert(entities[e] >= 0);
       auto dofs
           = MDSPAN_IMPL_STANDARD_NAMESPACE::MDSPAN_IMPL_PROPOSED_NAMESPACE::
               submdspan(x_dofmap, entities[e],
@@ -674,6 +677,7 @@ determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points)
   std::vector<std::int32_t> cells(num_cells, 0);
   std::iota(cells.begin(), cells.end(), 0);
   BoundingBoxTree bb(mesh, tdim, cells, padding);
+  BoundingBoxTree midpoint_tree = create_midpoint_tree(mesh, tdim, cells);
   BoundingBoxTree global_bbtree = bb.create_global_tree(comm);
 
   // Compute collisions:
@@ -751,20 +755,15 @@ determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points)
       dolfinx::MPI::mpi_type<T>(), received_points.data(), recv_sizes.data(),
       recv_offsets.data(), dolfinx::MPI::mpi_type<T>(), forward_comm);
 
-  // Each process checks which points collides with a cell on the process
+  // Each process checks which local cell is closest and computes the squared
+  // distance to the cell
   const int rank = dolfinx::MPI::rank(comm);
-  std::vector<std::int32_t> cell_indicator(received_points.size() / 3);
-  std::vector<std::int32_t> colliding_cells(received_points.size() / 3);
-  for (std::size_t p = 0; p < received_points.size(); p += 3)
-  {
-    std::array<T, 3> point;
-    std::copy(std::next(received_points.begin(), p),
-              std::next(received_points.begin(), p + 3), point.begin());
-    const int colliding_cell
-        = geometry::compute_first_colliding_cell(mesh, bb, point);
-    cell_indicator[p / 3] = (colliding_cell >= 0) ? rank : -1;
-    colliding_cells[p / 3] = colliding_cell;
-  }
+  const std::vector<std::int32_t> closest_cells = compute_closest_entity(
+      bb, midpoint_tree, mesh,
+      std::span<const T>(received_points.data(), received_points.size()));
+  const std::vector<T> squared_distances = squared_distance(
+      mesh, tdim, closest_cells,
+      std::span<const T>(received_points.data(), received_points.size()));
 
   // Create neighborhood communicator in the reverse direction: send
   // back col to requesting processes
@@ -791,20 +790,28 @@ determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points)
     std::swap(recv_offsets, send_offsets);
   }
 
-  std::vector<std::int32_t> recv_ranks(recv_offsets.back());
+  // Get distances from closest entity of points that were on the other process
+  std::vector<T> recv_distances(recv_offsets.back());
   MPI_Neighbor_alltoallv(
-      cell_indicator.data(), send_sizes.data(), send_offsets.data(),
-      dolfinx::MPI::mpi_type<std::int32_t>(), recv_ranks.data(),
-      recv_sizes.data(), recv_offsets.data(),
-      dolfinx::MPI::mpi_type<std::int32_t>(), reverse_comm);
+      squared_distances.data(), send_sizes.data(), send_offsets.data(),
+      dolfinx::MPI::mpi_type<T>(), recv_distances.data(), recv_sizes.data(),
+      recv_offsets.data(), dolfinx::MPI::mpi_type<T>(), reverse_comm);
 
   std::vector<std::int32_t> point_owners(points.size() / 3, -1);
-  for (std::size_t i = 0; i < unpack_map.size(); i++)
+  std::vector<T> closest_distance(points.size() / 3, -1);
+  for (std::size_t i = 0; i < out_ranks.size(); i++)
   {
-    const std::int32_t pos = unpack_map[i];
-    // Only insert new owner if no owner has previously been found
-    if ((recv_ranks[i] >= 0) && (point_owners[pos] == -1))
-      point_owners[pos] = recv_ranks[i];
+    for (std::int32_t j = recv_offsets[i]; j < recv_offsets[i + 1]; j++)
+    {
+      const std::int32_t pos = unpack_map[j];
+      // If point has not been found yet distance is negative
+      // If new received distance smaller than current distance choose owner
+      if (auto d = closest_distance[pos]; d < 0 or d > recv_distances[j])
+      {
+        point_owners[pos] = out_ranks[i];
+        closest_distance[pos] = recv_distances[j];
+      }
+    }
   }
 
   // Communication is reversed again to send dest ranks to all processes
@@ -847,7 +854,7 @@ determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points)
         owned_recv_points.insert(
             owned_recv_points.end(), std::next(received_points.cbegin(), 3 * j),
             std::next(received_points.cbegin(), 3 * (j + 1)));
-        owned_recv_cells.push_back(colliding_cells[j]);
+        owned_recv_cells.push_back(closest_cells[j]);
       }
     }
   }

python/dolfinx/wrappers/common.cpp

@@ -99,6 +99,8 @@ void common(py::module& m)
       .def_property_readonly("local_range",
                              &dolfinx::common::IndexMap::local_range,
                              "Range of indices owned by this map")
+      .def_property_readonly("imbalance", &dolfinx::common::IndexMap::imbalance,
+                             "Imbalance of the current IndexMap.")
       .def_property_readonly("index_to_dest_ranks",
                              &dolfinx::common::IndexMap::index_to_dest_ranks)
       .def_property_readonly(

python/test/unit/fem/test_interpolation.py

@@ -282,29 +282,31 @@ def f(x):
         u, v = Function(V), Function(V)
         u.interpolate(U.sub(i))
         v.interpolate(f)
-        assert np.allclose(u.vector.array, v.vector.array)
+        assert np.allclose(u.x.array, v.x.array)
 
         # Same map, different elements
         gdim = mesh.geometry.dim
         V = FunctionSpace(mesh, ("Lagrange", 1, (gdim,)))
         u, v = Function(V), Function(V)
         u.interpolate(U.sub(i))
         v.interpolate(f)
-        assert np.allclose(u.vector.array, v.vector.array)
+        assert np.allclose(u.x.array, v.x.array)
 
         # Different maps (0)
         V = FunctionSpace(mesh, ("N1curl", 1))
         u, v = Function(V), Function(V)
         u.interpolate(U.sub(i))
         v.interpolate(f)
-        assert np.allclose(u.vector.array, v.vector.array, atol=1.0e-6)
+        atol = 5 * np.finfo(u.x.array.dtype).resolution
+        assert np.allclose(u.x.array, v.x.array, atol=atol)
 
         # Different maps (1)
         V = FunctionSpace(mesh, ("RT", 2))
         u, v = Function(V), Function(V)
         u.interpolate(U.sub(i))
         v.interpolate(f)
-        assert np.allclose(u.vector.array, v.vector.array, atol=1.0e-6)
+        atol = 5 * np.finfo(u.x.array.dtype).resolution
+        assert np.allclose(u.x.array, v.x.array, atol=atol)
 
         # Test with wrong shape
         V0 = FunctionSpace(mesh, P.sub_elements()[0])
@@ -778,8 +780,8 @@ def test_nonmatching_mesh_single_cell_overlap_interpolation(xtype):
     mesh2 = create_rectangle(MPI.COMM_WORLD, [[0.0, 0.0], [p0_mesh2, p0_mesh2]], [n_mesh2, n_mesh2],
                              cell_type=CellType.triangle, dtype=xtype)
 
-    u1 = Function(FunctionSpace(mesh1, ("CG", 1)), name="u1", dtype=xtype)
-    u2 = Function(FunctionSpace(mesh2, ("CG", 1)), name="u2", dtype=xtype)
+    u1 = Function(FunctionSpace(mesh1, ("Lagrange", 1)), name="u1", dtype=xtype)
+    u2 = Function(FunctionSpace(mesh2, ("Lagrange", 1)), name="u2", dtype=xtype)
 
     def f_test1(x):
         return 1.0 - x[0] * x[1]