utils.h

// Copyright (C) 2019-2021 Garth N. Wells and Jørgen S. Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "BoundingBoxTree.h"
#include "gjk.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <concepts>
#include <cstdint>
#include <dolfinx/fem/CoordinateElement.h>
#include <dolfinx/fem/ElementDofLayout.h>
#include <dolfinx/graph/AdjacencyList.h>
#include <dolfinx/mesh/Mesh.h>
#include <limits>
#include <map>
#include <numeric>
#include <optional>
#include <span>
#include <utility>
#include <vector>
 
namespace dolfinx::geometry
{
template <std::floating_point T>
struct PointOwnershipData
{
  std::vector<int> src_owner; 
  std::vector<int>
      dest_owners; 
  std::vector<T> dest_points; 
  std::vector<std::int32_t>
      dest_cells; 
};

template <std::floating_point T>
std::vector<T> shortest_vector(const mesh::Mesh<T>& mesh, int dim,
                               std::span<const std::int32_t> entities,
                               std::span<const T> points)
{
  const int tdim = mesh.topology()->dim();
  const mesh::Geometry<T>& geometry = mesh.geometry();
 
  std::span<const T> geom_dofs = geometry.x();
  auto x_dofmap = geometry.dofmaps().front();
  std::vector<T> shortest_vectors;
  shortest_vectors.reserve(3 * entities.size());
  if (dim == tdim)
  {
    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 = md::submdspan(x_dofmap, entities[e], md::full_extent);
      std::vector<T> nodes(3 * dofs.size());
      for (std::size_t i = 0; i < dofs.size(); ++i)
      {
        const std::int32_t pos = 3 * dofs[i];
        for (std::size_t j = 0; j < 3; ++j)
          nodes[3 * i + j] = geom_dofs[pos + j];
      }
 
      std::array<T, 3> d
          = compute_distance_gjk<T>(points.subspan(3 * e, 3), nodes);
      shortest_vectors.insert(shortest_vectors.end(), d.begin(), d.end());
    }
  }
  else
  {
    mesh.topology_mutable()->create_connectivity(dim, tdim);
    mesh.topology_mutable()->create_connectivity(tdim, dim);
    auto e_to_c = mesh.topology()->connectivity(dim, tdim);
    assert(e_to_c);
    auto c_to_e = mesh.topology_mutable()->connectivity(tdim, dim);
    assert(c_to_e);
    for (std::size_t e = 0; e < entities.size(); e++)
    {
      const std::int32_t index = entities[e];
 
      // Find attached cell
      assert(e_to_c->num_links(index) > 0);
      const std::int32_t c = e_to_c->links(index)[0];
 
      // Find local number of entity wrt cell
      auto cell_entities = c_to_e->links(c);
      auto it0 = std::find(cell_entities.begin(), cell_entities.end(), index);
      assert(it0 != cell_entities.end());
      const int local_cell_entity
          = std::ranges::distance(cell_entities.begin(), it0);
 
      // Tabulate geometry dofs for the entity
      auto dofs = md::submdspan(x_dofmap, c, md::full_extent);
      const std::vector<int> entity_dofs
          = geometry.cmaps().front().create_dof_layout().entity_closure_dofs(
              dim, local_cell_entity);
      std::vector<T> nodes(3 * entity_dofs.size());
      for (std::size_t i = 0; i < entity_dofs.size(); i++)
      {
        const std::int32_t pos = 3 * dofs[entity_dofs[i]];
        for (std::size_t j = 0; j < 3; ++j)
          nodes[3 * i + j] = geom_dofs[pos + j];
      }
 
      std::array<T, 3> d
          = compute_distance_gjk<T>(points.subspan(3 * e, 3), nodes);
      shortest_vectors.insert(shortest_vectors.end(), d.begin(), d.end());
    }
  }
 
  return shortest_vectors;
}

template <std::floating_point T>
T compute_squared_distance_bbox(std::span<const T, 6> b,
                                std::span<const T, 3> x)
{
  auto b0 = b.template subspan<0, 3>();
  auto b1 = b.template subspan<3, 3>();
  return std::transform_reduce(x.begin(), x.end(), b0.begin(), 0.0,
                               std::plus<>{},
                               [](auto x, auto b)
                               {
                                 auto dx = x - b;
                                 return dx > 0 ? 0 : dx * dx;
                               })
         + std::transform_reduce(x.begin(), x.end(), b1.begin(), 0.0,
                                 std::plus<>{},
                                 [](auto x, auto b)
                                 {
                                   auto dx = x - b;
                                   return dx < 0 ? 0 : dx * dx;
                                 });
}

template <std::floating_point T>
std::vector<T> squared_distance(const mesh::Mesh<T>& mesh, int dim,
                                std::span<const std::int32_t> entities,
                                std::span<const T> points)
{
  std::vector<T> v = shortest_vector(mesh, dim, entities, points);
  std::vector<T> d(v.size() / 3, 0);
  for (std::size_t i = 0; i < d.size(); ++i)
    for (std::size_t j = 0; j < 3; ++j)
      d[i] += v[3 * i + j] * v[3 * i + j];
  return d;
}
 
namespace impl
{
constexpr bool is_leaf(std::array<std::int32_t, 2> bbox)
{
  // Leaf nodes are marked by setting child_0 equal to child_1
  return bbox[0] == bbox[1];
}

template <std::floating_point T>
constexpr bool point_in_bbox(std::span<const T, 6> b, std::span<const T, 3> x)
{
  constexpr T rtol = 1e-14;
  bool in = true;
  for (std::size_t i = 0; i < 3; i++)
  {
    T eps = rtol * (b[i + 3] - b[i]);
    in &= (x[i] >= (b[i] - eps)) && (x[i] <= (b[i + 3] + eps));
    if (!in)
      break;
  }
 
  return in;
}

template <std::floating_point T>
constexpr bool bbox_in_bbox(std::span<const T, 6> a, std::span<const T, 6> b)
{
  constexpr T rtol = 1e-14;
  auto a0 = a.template subspan<0, 3>();
  auto a1 = a.template subspan<3, 3>();
  auto b0 = b.template subspan<0, 3>();
  auto b1 = b.template subspan<3, 3>();
 
  bool in = true;
  for (std::size_t i = 0; i < 3; i++)
  {
    T eps = rtol * (b1[i] - b0[i]);
    in &= a1[i] >= (b0[i] - eps);
    in &= a0[i] <= (b1[i] + eps);
  }
 
  return in;
}

template <std::floating_point T>
struct ClosestEntityScratch
{
  std::vector<T> nodes; 
  const fem::ElementDofLayout&
      cmap_dof_layout; 
};

template <std::floating_point T>
T squared_distance_entity(const mesh::Mesh<T>& mesh, int dim,
                          std::int32_t entity, std::span<const T, 3> point,
                          ClosestEntityScratch<T>& scratch)
{
  const int tdim = mesh.topology()->dim();
  const mesh::Geometry<T>& geometry = mesh.geometry();
  std::span<const T> geom_dofs = geometry.x();
  auto x_dofmap = geometry.dofmaps().front();
 
  std::array<T, 3> d;
  if (dim == tdim)
  {
    assert(entity >= 0);
    auto dofs = md::submdspan(x_dofmap, entity, md::full_extent);
    assert(3 * dofs.size() <= scratch.nodes.size());
    for (std::size_t i = 0; i < dofs.size(); ++i)
    {
      const std::int32_t pos = 3 * dofs[i];
      for (std::size_t j = 0; j < 3; ++j)
        scratch.nodes[3 * i + j] = geom_dofs[pos + j];
    }
 
    d = compute_distance_gjk<T>(
        point, std::span<const T>(scratch.nodes.data(), 3 * dofs.size()));
  }
  else
  {
    mesh.topology_mutable()->create_connectivity(dim, tdim);
    mesh.topology_mutable()->create_connectivity(tdim, dim);
    auto e_to_c = mesh.topology()->connectivity(dim, tdim);
    assert(e_to_c);
    auto c_to_e = mesh.topology_mutable()->connectivity(tdim, dim);
    assert(c_to_e);
 
    // Find attached cell
    assert(e_to_c->num_links(entity) > 0);
    const std::int32_t c = e_to_c->links(entity)[0];
 
    // Find local number of entity wrt cell
    auto cell_entities = c_to_e->links(c);
    auto it0 = std::find(cell_entities.begin(), cell_entities.end(), entity);
    assert(it0 != cell_entities.end());
    const int local_cell_entity
        = std::ranges::distance(cell_entities.begin(), it0);
 
    // Tabulate geometry dofs for the entity
    auto dofs = md::submdspan(x_dofmap, c, md::full_extent);
    const std::vector<int>& entity_dofs
        = scratch.cmap_dof_layout.entity_closure_dofs(dim, local_cell_entity);
    assert(3 * entity_dofs.size() <= scratch.nodes.size());
    for (std::size_t i = 0; i < entity_dofs.size(); i++)
    {
      const std::int32_t pos = 3 * dofs[entity_dofs[i]];
      for (std::size_t j = 0; j < 3; ++j)
        scratch.nodes[3 * i + j] = geom_dofs[pos + j];
    }
 
    d = compute_distance_gjk<T>(
        point,
        std::span<const T>(scratch.nodes.data(), 3 * entity_dofs.size()));
  }
 
  return d[0] * d[0] + d[1] * d[1] + d[2] * d[2];
}

template <std::floating_point T>
std::pair<std::int32_t, T>
_compute_closest_entity(const geometry::BoundingBoxTree<T>& tree,
                        std::span<const T, 3> point, std::int32_t node,
                        const mesh::Mesh<T>& mesh, std::int32_t closest_entity,
                        T R2, ClosestEntityScratch<T>& scratch)
{
  // Get children of current bounding box node (child_1 denotes entity
  // index for leaves)
  const std::array<std::int32_t, 2> bbox = tree.bbox(node);
  T r2;
  if (is_leaf(bbox))
  {
    // If point cloud tree the exact distance is easy to compute
    if (tree.tdim() == 0)
    {
      std::array<T, 6> diff = tree.get_bbox(node);
      for (std::size_t k = 0; k < 3; ++k)
        diff[k] -= point[k];
      r2 = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2];
    }
    else
    {
      r2 = compute_squared_distance_bbox<T>(tree.get_bbox(node), point);
 
      // If bounding box closer than previous closest entity, use gjk to
      // obtain exact distance to the convex hull of the entity
      if (r2 <= R2)
        r2 = squared_distance_entity(mesh, tree.tdim(), bbox.back(), point,
                                     scratch);
    }
 
    // If entity is closer than best result so far, return it
    if (r2 <= R2)
    {
      closest_entity = bbox.back();
      R2 = r2;
    }
 
    return {closest_entity, R2};
  }
  else
  {
    // If bounding box is outside radius, then don't search further
    r2 = compute_squared_distance_bbox<T>(tree.get_bbox(node), point);
    if (r2 > R2)
      return {closest_entity, R2};
 
    // Check both children. We use R2 (as opposed to r2), as a bounding
    // box can be closer than the actual entity.
    std::pair<std::int32_t, T> p0 = _compute_closest_entity(
        tree, point, bbox.front(), mesh, closest_entity, R2, scratch);
    std::pair<std::int32_t, T> p1 = _compute_closest_entity(
        tree, point, bbox.back(), mesh, p0.first, p0.second, scratch);
    return p1;
  }
}

template <std::floating_point T>
void _compute_collisions_point(const geometry::BoundingBoxTree<T>& tree,
                               std::span<const T, 3> p,
                               std::vector<std::int32_t>& entities,
                               std::vector<std::int32_t>& stack)
{
  assert(stack.empty());
  std::int32_t next = tree.num_bboxes() - 1;
  std::span<const T> coords = tree.bbox_coordinates();
  auto view_bbox = [&coords](std::int32_t node)
  { return std::span<const T, 6>(coords.data() + 6 * node, 6); };
  while (next != -1)
  {
    if (std::array bbox = tree.bbox(next);
        is_leaf(bbox) and point_in_bbox(view_bbox(next), p))
    {
      // If box is a leaf node then add it to the list of colliding
      // entities
      entities.push_back(bbox[1]);
      next = -1;
    }
    else
    {
      // Check whether the point collides with child nodes (left and
      // right)
      bool left = point_in_bbox(view_bbox(bbox[0]), p);
      bool right = point_in_bbox(view_bbox(bbox[1]), p);
      if (left and right)
      {
        // If the point collides with both child nodes, add the right
        // node to the stack (for later visiting) and continue the tree
        // traversal with the left subtree
        stack.push_back(bbox[1]);
        next = bbox[0];
      }
      else if (left)
      {
        // Traverse the current node's left subtree
        next = bbox[0];
      }
      else if (right)
      {
        // Traverse the current node's right subtree
        next = bbox[1];
      }
      else
        next = -1;
    }
 
    // If tree traversal reaches a dead end (box is a leaf node or no
    // collision detected), check the stack for deferred subtrees
    if (next == -1 and !stack.empty())
    {
      next = stack.back();
      stack.pop_back();
    }
  }
}
 
// Compute collisions with tree (recursive)
template <std::floating_point T>
void _compute_collisions_tree(const geometry::BoundingBoxTree<T>& A,
                              const geometry::BoundingBoxTree<T>& B,
                              std::int32_t node_A, std::int32_t node_B,
                              std::vector<std::int32_t>& entities)
{
  // If bounding boxes don't collide, then don't search further
  if (!bbox_in_bbox<T>(A.get_bbox(node_A), B.get_bbox(node_B)))
    return;
 
  // Get bounding boxes for current nodes
  const std::array<std::int32_t, 2> bbox_A = A.bbox(node_A);
  const std::array<std::int32_t, 2> bbox_B = B.bbox(node_B);
 
  // Check whether we've reached a leaf in A or B
  const bool is_leaf_A = is_leaf(bbox_A);
  const bool is_leaf_B = is_leaf(bbox_B);
  if (is_leaf_A and is_leaf_B)
  {
    // If both boxes are leaves (which we know collide), then add them
    // child_1 denotes entity for leaves
    entities.push_back(bbox_A[1]);
    entities.push_back(bbox_B[1]);
  }
  else if (is_leaf_A)
  {
    // If we reached the leaf in A, then descend B
    _compute_collisions_tree(A, B, node_A, bbox_B[0], entities);
    _compute_collisions_tree(A, B, node_A, bbox_B[1], entities);
  }
  else if (is_leaf_B)
  {
    // If we reached the leaf in B, then descend A
    _compute_collisions_tree(A, B, bbox_A[0], node_B, entities);
    _compute_collisions_tree(A, B, bbox_A[1], node_B, entities);
  }
  else if (node_A > node_B)
  {
    // At this point, we know neither is a leaf so descend the largest
    // tree first. Note that nodes are added in reverse order with the
    // top bounding box at the end so the largest tree (the one with the
    // the most boxes left to traverse) has the largest node number.
    _compute_collisions_tree(A, B, bbox_A[0], node_B, entities);
    _compute_collisions_tree(A, B, bbox_A[1], node_B, entities);
  }
  else
  {
    _compute_collisions_tree(A, B, node_A, bbox_B[0], entities);
    _compute_collisions_tree(A, B, node_A, bbox_B[1], entities);
  }
 
  // Note that cases above can be collected in fewer cases but this way
  // the logic is easier to follow.
}
 
} // namespace impl

template <std::floating_point T>
BoundingBoxTree<T> create_midpoint_tree(const mesh::Mesh<T>& mesh, int tdim,
                                        std::span<const std::int32_t> entities)
{
  spdlog::info("Building point search tree to accelerate distance queries for "
               "a given topological dimension and subset of entities.");
 
  const std::vector<T> midpoints
      = mesh::compute_midpoints(mesh, tdim, entities);
  std::vector<std::pair<std::array<T, 3>, std::int32_t>> points(
      entities.size());
  for (std::size_t i = 0; i < points.size(); ++i)
  {
    for (std::size_t j = 0; j < 3; ++j)
      points[i].first[j] = midpoints[3 * i + j];
    points[i].second = entities[i];
  }
 
  // Build tree
  return BoundingBoxTree(points);
}

template <std::floating_point T>
std::vector<std::int32_t> compute_collisions(const BoundingBoxTree<T>& tree0,
                                             const BoundingBoxTree<T>& tree1)
{
  // Call recursive find function
  std::vector<std::int32_t> entities;
  if (tree0.num_bboxes() > 0 and tree1.num_bboxes() > 0)
  {
    impl::_compute_collisions_tree(tree0, tree1, tree0.num_bboxes() - 1,
                                   tree1.num_bboxes() - 1, entities);
  }
 
  return entities;
}

template <std::floating_point T>
graph::AdjacencyList<std::int32_t>
compute_collisions(const BoundingBoxTree<T>& tree, std::span<const T> points)
{
  if (tree.num_bboxes() > 0)
  {
    std::vector<std::int32_t> entities, offsets(points.size() / 3 + 1, 0);
    entities.reserve(points.size() / 3);
 
    // Scratch stack reused across points.
    std::vector<std::int32_t> stack;
    for (std::size_t p = 0; p < points.size() / 3; ++p)
    {
      impl::_compute_collisions_point(
          tree, std::span<const T, 3>(points.data() + 3 * p, 3), entities,
          stack);
      offsets[p + 1] = entities.size();
    }
 
    return graph::AdjacencyList(std::move(entities), std::move(offsets));
  }
  else
  {
    return graph::AdjacencyList(
        std::vector<std::int32_t>(),
        std::vector<std::int32_t>(points.size() / 3 + 1, 0));
  }
}

template <std::floating_point T>
std::int32_t compute_first_colliding_cell(const mesh::Mesh<T>& mesh,
                                          std::span<const std::int32_t> cells,
                                          std::array<T, 3> point, T tol,
                                          std::span<T> coordinate_dofs)
{
  if (cells.empty())
    return -1;
 
  const mesh::Geometry<T>& geometry = mesh.geometry();
  std::span<const T> geom_dofs = geometry.x();
  auto x_dofmap = geometry.dofmaps().front();
  const std::size_t num_nodes = x_dofmap.extent(1);
  assert(coordinate_dofs.size() >= num_nodes * 3);
  std::span<T> cell_dofs = coordinate_dofs.first(num_nodes * 3);
  for (auto cell : cells)
  {
    auto dofs = md::submdspan(x_dofmap, cell, md::full_extent);
    for (std::size_t i = 0; i < num_nodes; ++i)
    {
      std::copy_n(std::next(geom_dofs.begin(), 3 * dofs[i]), 3,
                  std::next(cell_dofs.begin(), 3 * i));
    }
 
    std::array<T, 3> shortest_vector
        = compute_distance_gjk<T>(point, cell_dofs);
    T d2 = std::reduce(shortest_vector.begin(), shortest_vector.end(), T(0),
                       [](auto d, auto e) { return d + e * e; });
    if (d2 < tol)
      return cell;
  }
 
  return -1;
}

template <std::floating_point T>
std::vector<std::int32_t>
compute_closest_entity(const BoundingBoxTree<T>& tree,
                       const BoundingBoxTree<T>& midpoint_tree,
                       const mesh::Mesh<T>& mesh, std::span<const T> points)
{
  if (tree.num_bboxes() == 0)
    return std::vector<std::int32_t>(points.size() / 3, -1);
 
  std::vector<std::int32_t> entities;
  entities.reserve(points.size() / 3);
 
  const fem::ElementDofLayout cmap_dof_layout
      = mesh.geometry().cmaps().front().create_dof_layout();
 
  // Scratch buffer sized for the largest coordinate element, to avoid
  // repeated resizing in the closest-entity search (mixed topology may
  // have more than one coordinate element).
  std::size_t max_nodes = 0;
  for (const fem::CoordinateElement<T>& cmap : mesh.geometry().cmaps())
    max_nodes = std::max(max_nodes, static_cast<std::size_t>(cmap.dim()));
  impl::ClosestEntityScratch<T> scratch{std::vector<T>(3 * max_nodes),
                                        cmap_dof_layout};
  for (std::size_t i = 0; i < points.size() / 3; ++i)
  {
    // Use midpoint tree to find initial closest entity to the point.
    // Start by using a leaf node as the initial guess for the input
    // entity
    std::array<std::int32_t, 2> leaf0 = midpoint_tree.bbox(0);
    assert(impl::is_leaf(leaf0));
    std::array<T, 6> diff = midpoint_tree.get_bbox(0);
    for (std::size_t k = 0; k < 3; ++k)
      diff[k] -= points[3 * i + k];
    T R2 = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2];
 
    // Use a recursive search through the bounding box tree
    // to find determine the entity with the closest midpoint.
    // As the midpoint tree only consist of points, the distance
    // queries are lightweight.
    const auto [m_index, m_distance2] = impl::_compute_closest_entity(
        midpoint_tree, std::span<const T, 3>(points.data() + 3 * i, 3),
        midpoint_tree.num_bboxes() - 1, mesh, leaf0[0], R2, scratch);
 
    // Use a recursives search through the bounding box tree to
    // determine which entity is actually closest.
    // Uses the entity with the closest midpoint as initial guess, and
    // the distance from the midpoint to the point of interest as the
    // initial search radius.
    const auto [index, distance2] = impl::_compute_closest_entity(
        tree, std::span<const T, 3>(points.data() + 3 * i, 3),
        tree.num_bboxes() - 1, mesh, m_index, m_distance2, scratch);
 
    entities.push_back(index);
  }
 
  return entities;
}

template <std::floating_point T>
graph::AdjacencyList<std::int32_t> compute_colliding_cells(
    const mesh::Mesh<T>& mesh,
    const graph::AdjacencyList<std::int32_t>& candidate_cells,
    std::span<const T> points)
{
  std::vector<std::int32_t> offsets = {0};
  offsets.reserve(candidate_cells.num_nodes() + 1);
  std::vector<std::int32_t> colliding_cells;
  constexpr T eps2 = 1e-12;
  const int tdim = mesh.topology()->dim();
  for (std::int32_t i = 0; i < candidate_cells.num_nodes(); i++)
  {
    auto cells = candidate_cells.links(i);
    std::vector<T> _point(3 * cells.size());
    for (std::size_t j = 0; j < cells.size(); ++j)
      for (std::size_t k = 0; k < 3; ++k)
        _point[3 * j + k] = points[3 * i + k];
 
    std::vector distances_sq = squared_distance<T>(mesh, tdim, cells, _point);
    for (std::size_t j = 0; j < cells.size(); j++)
      if (distances_sq[j] < eps2)
        colliding_cells.push_back(cells[j]);
 
    offsets.push_back(colliding_cells.size());
  }
 
  return graph::AdjacencyList(std::move(colliding_cells), std::move(offsets));
}

template <std::floating_point T>
PointOwnershipData<T>
determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points,
                          T padding,
                          std::optional<std::span<const std::int32_t>> cells,
                          bool find_closest_cell = true)
{
  MPI_Comm comm = mesh.comm();
 
  const int tdim = mesh.topology()->dim();
 
  std::vector<std::int32_t> local_cells;
  if (not(cells.has_value()))
  {
    auto cell_map = mesh.topology()->index_map(tdim);
    local_cells.resize(cell_map->size_local());
    std::iota(local_cells.begin(), local_cells.end(), 0);
    cells
        = std::span<const std::int32_t>(local_cells.data(), local_cells.size());
  }
  // Create a global bounding-box tree to find candidate processes with
  // cells that could collide with the points
  BoundingBoxTree bb(mesh, tdim, padding, cells.value());
  BoundingBoxTree global_bbtree = bb.create_global_tree(comm);
 
  // Compute collisions:
  // For each point in `points` get the processes it should be sent to
  graph::AdjacencyList collisions = compute_collisions(global_bbtree, points);
 
  // Get unique list of outgoing ranks
  std::vector<std::int32_t> out_ranks = collisions.array();
  std::ranges::sort(out_ranks);
  auto [unique_end, range_end] = std::ranges::unique(out_ranks);
  out_ranks.erase(unique_end, range_end);
 
  // Compute incoming edges (source processes)
  std::vector in_ranks = dolfinx::MPI::compute_graph_edges_nbx(comm, out_ranks);
  std::ranges::sort(in_ranks);
 
  // Create neighborhood communicator in forward direction
  MPI_Comm forward_comm;
  MPI_Dist_graph_create_adjacent(
      comm, in_ranks.size(), in_ranks.data(), MPI_UNWEIGHTED, out_ranks.size(),
      out_ranks.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &forward_comm);
 
  // Compute map from global mpi rank to neighbor rank, "collisions"
  // uses global rank
  std::map<std::int32_t, std::int32_t> rank_to_neighbor;
  for (std::size_t i = 0; i < out_ranks.size(); i++)
    rank_to_neighbor[out_ranks[i]] = i;
 
  // Count the number of points to send per neighbor process
  std::vector<std::int32_t> send_sizes(out_ranks.size());
  for (std::size_t i = 0; i < points.size() / 3; ++i)
    for (std::int32_t p : collisions.links(i))
      send_sizes[rank_to_neighbor[p]] += 3;
 
  // Compute receive sizes
  std::vector<std::int32_t> recv_sizes(in_ranks.size());
  send_sizes.reserve(1);
  recv_sizes.reserve(1);
  MPI_Request sizes_request;
  MPI_Ineighbor_alltoall(send_sizes.data(), 1, MPI_INT, recv_sizes.data(), 1,
                         MPI_INT, forward_comm, &sizes_request);
 
  // Compute sending offsets
  std::vector<std::int32_t> send_offsets(send_sizes.size() + 1, 0);
  std::partial_sum(send_sizes.begin(), send_sizes.end(),
                   std::next(send_offsets.begin(), 1));
 
  // Pack data to send and store unpack map
  std::vector<T> send_data(send_offsets.back());
  std::vector<std::int32_t> counter(send_sizes.size(), 0);
  // unpack map: [index in adj list][pos in x]
  std::vector<std::int32_t> unpack_map(send_offsets.back() / 3);
  for (std::size_t i = 0; i < points.size(); i += 3)
  {
    for (std::int32_t p : collisions.links(i / 3))
    {
      std::int32_t neighbor = rank_to_neighbor[p];
      std::int32_t pos = send_offsets[neighbor] + counter[neighbor];
      auto it = std::next(send_data.begin(), pos);
      std::copy_n(std::next(points.begin(), i), 3, it);
      unpack_map[pos / 3] = i / 3;
      counter[neighbor] += 3;
    }
  }
 
  MPI_Wait(&sizes_request, MPI_STATUS_IGNORE);
  std::vector<std::int32_t> recv_offsets(in_ranks.size() + 1, 0);
  std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
                   std::next(recv_offsets.begin(), 1));
 
  std::vector<T> received_points((std::size_t)recv_offsets.back());
  MPI_Neighbor_alltoallv(
      send_data.data(), send_sizes.data(), send_offsets.data(),
      dolfinx::MPI::mpi_t<T>, received_points.data(), recv_sizes.data(),
      recv_offsets.data(), dolfinx::MPI::mpi_t<T>, forward_comm);
 
  // Get mesh geometry for closest entity
  const mesh::Geometry<T>& geometry = mesh.geometry();
  std::span<const T> geom_dofs = geometry.x();
  auto x_dofmap = geometry.dofmaps().front();
 
  // Compute candidate cells for collisions (and extrapolation)
  const graph::AdjacencyList<std::int32_t> candidate_collisions
      = compute_collisions(bb, std::span<const T>(received_points.data(),
                                                  received_points.size()));
 
  // Scratch buffer reused across points below to avoid a per-point
  // allocation, both for the initial collision check and (further down)
  // for the closest-cell fallback.
  std::vector<T> coordinate_dofs(3 * x_dofmap.extent(1));
 
  // Each process checks which points collide with a cell on the process
  const int rank = dolfinx::MPI::rank(comm);
  std::vector<std::int32_t> cell_indicator(received_points.size() / 3);
  std::vector<std::int32_t> closest_cells(received_points.size() / 3);
  for (std::size_t p = 0; p < received_points.size(); p += 3)
  {
    std::array<T, 3> point;
    std::copy_n(std::next(received_points.begin(), p), 3, point.begin());
    // Find first colliding cell among the cells with colliding bounding boxes
    const std::int32_t colliding_cell = geometry::compute_first_colliding_cell(
        mesh, candidate_collisions.links(p / 3), point,
        10 * std::numeric_limits<T>::epsilon(), std::span<T>(coordinate_dofs));
    // If a collding cell is found, store the rank of the current process
    // which will be sent back to the owner of the point
    cell_indicator[p / 3] = (colliding_cell >= 0) ? rank : -1;
    // Store the cell index for lookup once the owning processes has determined
    // the ownership of the point
    closest_cells[p / 3] = colliding_cell;
  }
 
  // Create neighborhood communicator in the reverse direction: send
  // back col to requesting processes
  MPI_Comm reverse_comm;
  MPI_Dist_graph_create_adjacent(
      comm, out_ranks.size(), out_ranks.data(), MPI_UNWEIGHTED, in_ranks.size(),
      in_ranks.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &reverse_comm);
 
  // Reuse sizes and offsets from first communication set
  // but divide by three
  {
    auto rescale = [](auto& x)
    { std::ranges::transform(x, x.begin(), [](auto e) { return (e / 3); }); };
    rescale(recv_sizes);
    rescale(recv_offsets);
    rescale(send_sizes);
    rescale(send_offsets);
 
    // The communication is reversed, so swap recv to send offsets
    std::swap(recv_sizes, send_sizes);
    std::swap(recv_offsets, send_offsets);
  }
 
  std::vector<std::int32_t> recv_ranks(recv_offsets.back());
  MPI_Neighbor_alltoallv(cell_indicator.data(), send_sizes.data(),
                         send_offsets.data(), MPI_INT32_T, recv_ranks.data(),
                         recv_sizes.data(), recv_offsets.data(), MPI_INT32_T,
                         reverse_comm);
 
  std::vector<int> point_owners(points.size() / 3, -1);
  for (std::size_t i = 0; i < unpack_map.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];
  }
 
  if (find_closest_cell)
  {
    // Marker for those points already sent to another process that still
    // need a closest-cell fallback (no owner was found above)
    std::vector<std::uint8_t> send_needs_closest_cell(recv_offsets.back());
    for (std::int32_t i = 0; i < recv_offsets.back(); i++)
    {
      const std::int32_t pos = unpack_map[i];
      send_needs_closest_cell[i] = point_owners[pos] == -1;
    }
 
    // Swap communication direction, to send the marker to other processes
    std::swap(send_sizes, recv_sizes);
    std::swap(send_offsets, recv_offsets);
    std::vector<std::uint8_t> dest_needs_closest_cell(recv_offsets.back());
    MPI_Neighbor_alltoallv(send_needs_closest_cell.data(), send_sizes.data(),
                           send_offsets.data(), MPI_UINT8_T,
                           dest_needs_closest_cell.data(), recv_sizes.data(),
                           recv_offsets.data(), MPI_UINT8_T, forward_comm);
 
    std::vector<T> squared_distances(received_points.size() / 3, -1);
 
    for (std::size_t i = 0; i < dest_needs_closest_cell.size(); i++)
    {
      if (dest_needs_closest_cell[i] == 1)
      {
        assert(closest_cells[i] == -1);
        std::array<T, 3> point;
        std::copy_n(std::next(received_points.begin(), 3 * i), 3,
                    point.begin());
 
        // Find shortest distance among cells with colliding bounding box
        T shortest_distance = std::numeric_limits<T>::max();
        std::int32_t closest_cell = -1;
        for (auto cell : candidate_collisions.links(i))
        {
          auto dofs = md::submdspan(x_dofmap, cell, md::full_extent);
          for (std::size_t j = 0; j < dofs.size(); ++j)
          {
            const int pos = 3 * dofs[j];
            for (std::size_t k = 0; k < 3; ++k)
              coordinate_dofs[3 * j + k] = geom_dofs[pos + k];
          }
          const std::array<T, 3> d = compute_distance_gjk<T>(
              std::span<const T>(point.data(), point.size()),
              std::span<const T>(coordinate_dofs.data(), 3 * dofs.size()));
          if (T current_distance = d[0] * d[0] + d[1] * d[1] + d[2] * d[2];
              current_distance < shortest_distance)
          {
            shortest_distance = current_distance;
            closest_cell = cell;
          }
        }
        closest_cells[i] = closest_cell;
        squared_distances[i] = shortest_distance;
      }
    }
 
    std::swap(recv_sizes, send_sizes);
    std::swap(recv_offsets, send_offsets);
 
    // Get distances from closest entity of points that were on the other
    // process
    std::vector<T> recv_distances(recv_offsets.back());
    MPI_Neighbor_alltoallv(
        squared_distances.data(), send_sizes.data(), send_offsets.data(),
        dolfinx::MPI::mpi_t<T>, recv_distances.data(), recv_sizes.data(),
        recv_offsets.data(), dolfinx::MPI::mpi_t<T>, reverse_comm);
 
    // Update point ownership with extrapolation information
    std::vector<T> closest_distance(point_owners.size(),
                                    std::numeric_limits<T>::max());
    for (std::size_t i = 0; i < out_ranks.size(); i++)
    {
      for (std::int32_t j = recv_offsets[i]; j < recv_offsets[i + 1]; j++)
      {
        const std::int32_t pos = unpack_map[j];
        T current_dist = recv_distances[j];
        // Update if closer than previous guess and was found (current_dist
        // is -1, not a valid squared distance, if no candidate cell was
        // found on the sending rank)
        if (T d = closest_distance[pos];
            (current_dist >= 0) and (current_dist < d))
        {
          point_owners[pos] = out_ranks[i];
          closest_distance[pos] = current_dist;
        }
      }
    }
  }
 
  // Communication is reversed again to send dest ranks to all processes
  std::swap(send_sizes, recv_sizes);
  std::swap(send_offsets, recv_offsets);
 
  // Pack ownership data
  std::vector<std::int32_t> send_owners(send_offsets.back());
  std::ranges::fill(counter, 0);
  for (std::size_t i = 0; i < points.size() / 3; ++i)
  {
    for (std::int32_t p : collisions.links(i))
    {
      std::int32_t neighbor = rank_to_neighbor[p];
      send_owners[send_offsets[neighbor] + counter[neighbor]++]
          = point_owners[i];
    }
  }
 
  // Send ownership info
  std::vector<std::int32_t> dest_ranks(recv_offsets.back());
  MPI_Neighbor_alltoallv(send_owners.data(), send_sizes.data(),
                         send_offsets.data(), MPI_INT32_T, dest_ranks.data(),
                         recv_sizes.data(), recv_offsets.data(), MPI_INT32_T,
                         forward_comm);
 
  // Unpack dest ranks if point owner is this rank
  std::vector<int> owned_recv_ranks;
  owned_recv_ranks.reserve(recv_offsets.back());
  std::vector<T> owned_recv_points;
  std::vector<std::int32_t> owned_recv_cells;
  for (std::size_t i = 0; i < in_ranks.size(); i++)
  {
    for (std::int32_t j = recv_offsets[i]; j < recv_offsets[i + 1]; j++)
    {
      if (rank == dest_ranks[j])
      {
        owned_recv_ranks.push_back(in_ranks[i]);
        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(closest_cells[j]);
      }
    }
  }
 
  MPI_Comm_free(&forward_comm);
  MPI_Comm_free(&reverse_comm);
  return PointOwnershipData<T>{.src_owner = std::move(point_owners),
                               .dest_owners = std::move(owned_recv_ranks),
                               .dest_points = std::move(owned_recv_points),
                               .dest_cells = std::move(owned_recv_cells)};
}
 
} // namespace dolfinx::geometry