utils.h

Go to the documentation of this file.

// Copyright (C) 2019-2026 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 "EntityMap.h"
#include "Mesh.h"
#include "MeshTags.h"
#include "Topology.h"
#include "graphbuild.h"
#include "types.h"
#include <algorithm>
#include <array>
#include <basix/mdspan.hpp>
#include <boost/unordered/unordered_flat_map.hpp>
#include <cassert>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/MPI.h>
#include <dolfinx/common/Timer.h>
#include <dolfinx/common/log.h>
#include <dolfinx/common/sort.h>
#include <dolfinx/graph/AdjacencyList.h>
#include <dolfinx/graph/ordering.h>
#include <dolfinx/graph/partition.h>
#include <exception>
#include <format>
#include <iterator>
#include <memory>
#include <mpi.h>
#include <numeric>
#include <optional>
#include <ranges>
#include <span>
#include <stdexcept>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>

 
namespace dolfinx::fem
{
class ElementDofLayout;
}
 
namespace dolfinx::mesh
{
enum class CellType : std::int8_t;
 
namespace impl
{
template <typename T>
void reorder_list(std::span<T> list, std::span<const std::int32_t> nodemap)
{
  if (nodemap.empty())
    return;
 
  assert(list.size() % nodemap.size() == 0);
  std::size_t degree = list.size() / nodemap.size();
  const std::vector<T> orig(list.begin(), list.end());
  for (std::size_t n = 0; n < nodemap.size(); ++n)
  {
    std::span links_old(orig.data() + n * degree, degree);
    auto links_new = list.subspan(nodemap[n] * degree, degree);
    std::ranges::copy(links_old, links_new.begin());
  }
}

template <std::floating_point T>
std::tuple<std::vector<std::int32_t>, std::vector<T>, std::vector<std::int32_t>>
compute_vertex_coords_boundary(const mesh::Mesh<T>& mesh, int dim,
                               std::span<const std::int32_t> facets)
{
  auto topology = mesh.topology();
  assert(topology);
  const int tdim = topology->dim();
  if (dim == tdim)
  {
    throw std::invalid_argument(
        "Cannot use mesh::locate_entities_boundary (boundary) for cells.");
  }
 
  // Build set of vertices on boundary and set of boundary entities
  mesh.topology_mutable()->create_connectivity(tdim - 1, 0);
  mesh.topology_mutable()->create_connectivity(tdim - 1, dim);
  std::vector<std::int32_t> vertices, entities;
  {
    auto f_to_v = topology->connectivity(tdim - 1, 0);
    assert(f_to_v);
    auto f_to_e = topology->connectivity(tdim - 1, dim);
    assert(f_to_e);
    for (auto f : facets)
    {
      auto v = f_to_v->links(f);
      vertices.insert(vertices.end(), v.begin(), v.end());
      auto e = f_to_e->links(f);
      entities.insert(entities.end(), e.begin(), e.end());
    }
 
    // Build vector of boundary vertices
    {
      std::ranges::sort(vertices);
      auto [unique_end, range_end] = std::ranges::unique(vertices);
      vertices.erase(unique_end, range_end);
    }
 
    {
      std::ranges::sort(entities);
      auto [unique_end, range_end] = std::ranges::unique(entities);
      entities.erase(unique_end, range_end);
    }
  }
 
  // Get geometry data
  auto x_dofmap = mesh.geometry().dofmaps().front();
  std::span<const T> x_nodes = mesh.geometry().x();
 
  // Get all vertex 'node' indices
  mesh.topology_mutable()->create_connectivity(0, tdim);
  mesh.topology_mutable()->create_connectivity(tdim, 0);
  auto v_to_c = topology->connectivity(0, tdim);
  assert(v_to_c);
  auto c_to_v = topology->connectivity(tdim, 0);
  assert(c_to_v);
  std::vector<T> x_vertices(3 * vertices.size(), -1.0);
  std::vector<std::int32_t> vertex_to_pos(v_to_c->num_nodes(), -1);
  for (std::size_t i = 0; i < vertices.size(); ++i)
  {
    const std::int32_t v = vertices[i];
 
    // Get first cell and find position
    const std::int32_t c = v_to_c->links(v).front();
    auto cell_vertices = c_to_v->links(c);
    auto it = std::ranges::find(cell_vertices, v);
    assert(it != cell_vertices.end());
    const std::size_t local_pos
        = std::ranges::distance(cell_vertices.begin(), it);
 
    auto dofs = md::submdspan(x_dofmap, c, md::full_extent);
    for (std::size_t j = 0; j < 3; ++j)
      x_vertices[j * vertices.size() + i] = x_nodes[3 * dofs[local_pos] + j];
    vertex_to_pos[v] = i;
  }
 
  return {std::move(entities), std::move(x_vertices), std::move(vertex_to_pos)};
}
 
} // namespace impl

std::vector<std::int32_t> exterior_facet_indices(const Topology& topology,
                                                 int facet_type_idx);

std::vector<std::int32_t> exterior_facet_indices(const Topology& topology);
 
namespace impl
{
std::vector<std::int64_t> reorder_cells(
    const graph::Reorder& reorder_fn, std::span<const double> cell_centroids,
    int gdim, std::optional<std::int32_t> max_facet_to_cell_links,
    const std::vector<CellType>& celltypes,
    const std::vector<fem::ElementDofLayout>& doflayouts,
    const std::vector<std::vector<int>>& ghost_owners,
    std::vector<std::vector<std::int64_t>>& cells,
    std::vector<std::span<std::int64_t>>& cells_v,
    std::vector<std::vector<std::int64_t>>& original_idx, int num_threads);
} // namespace impl

std::vector<std::int64_t> extract_topology(CellType cell_type,
                                           const fem::ElementDofLayout& layout,
                                           std::span<const std::int64_t> cells);

bool is_vertex_dof_layout(CellType cell_type,
                          const fem::ElementDofLayout& layout);

template <std::floating_point T>
std::vector<T> h(const Mesh<T>& mesh, std::span<const std::int32_t> entities,
                 int dim)
{
  if (entities.empty())
    return std::vector<T>();
  if (dim == 0)
    return std::vector<T>(entities.size(), 0);
 
  // Get the geometry dofs for the vertices of each entity
  const auto [vertex_xdofs, xdof_shape]
      = entities_to_geometry(mesh, dim, entities, false);
 
  // Get the  geometry coordinate
  std::span<const T> x = mesh.geometry().x();
 
  // Function to compute the length of (p0 - p1)
  auto delta_norm = [](auto&& p0, auto&& p1)
  {
    T norm = 0;
    for (std::size_t i = 0; i < 3; ++i)
      norm += (p0[i] - p1[i]) * (p0[i] - p1[i]);
    return std::sqrt(norm);
  };
 
  // Compute greatest distance between any to vertices
  assert(dim > 0);
  std::vector<T> h(entities.size(), 0);
  for (std::size_t e = 0; e < entities.size(); ++e)
  {
    // Get geometry 'dof' for each vertex of entity e
    std::span<const std::int32_t> e_vertices(
        vertex_xdofs.data() + e * xdof_shape[1], xdof_shape[1]);
 
    // Compute maximum distance between any two vertices
    for (std::size_t i = 0; i < e_vertices.size(); ++i)
    {
      std::span<const T, 3> p0(x.data() + 3 * e_vertices[i], 3);
      for (std::size_t j = i + 1; j < e_vertices.size(); ++j)
      {
        std::span<const T, 3> p1(x.data() + 3 * e_vertices[j], 3);
        h[e] = std::max(h[e], delta_norm(p0, p1));
      }
    }
  }
 
  return h;
}

template <std::floating_point T>
std::vector<T> cell_normals(const Mesh<T>& mesh, int dim,
                            std::span<const std::int32_t> entities)
{
  if (entities.empty())
    return std::vector<T>();
 
  auto topology = mesh.topology();
  assert(topology);
  if (topology->cell_type() == CellType::prism and dim == 2)
  {
    throw std::invalid_argument(
        "Cell normal computation for prism cells not yet supported.");
  }
 
  const int gdim = mesh.geometry().dim();
  const CellType type = cell_entity_type(topology->cell_type(), dim, 0);
 
  // Find geometry nodes for topology entities
  std::span<const T> x = mesh.geometry().x();
  const auto [geometry_entities, eshape]
      = entities_to_geometry(mesh, dim, entities, false);
 
  std::vector<T> n(entities.size() * 3);
  switch (type)
  {
  case CellType::interval:
  {
    if (gdim > 2)
      throw std::invalid_argument("Interval cell normal undefined in 3D.");
    for (std::size_t i = 0; i < entities.size(); ++i)
    {
      // Get the two vertices as points
      std::array vertices{geometry_entities[i * eshape[1]],
                          geometry_entities[i * eshape[1] + 1]};
      std::array p = {std::span<const T, 3>(x.data() + 3 * vertices[0], 3),
                      std::span<const T, 3>(x.data() + 3 * vertices[1], 3)};
 
      // Define normal by rotating tangent counter-clockwise
      std::array<T, 3> t;
      std::ranges::transform(p[1], p[0], t.begin(),
                             [](auto x, auto y) { return x - y; });
 
      T norm = std::sqrt(t[0] * t[0] + t[1] * t[1]);
      std::span<T, 3> ni(n.data() + 3 * i, 3);
      ni[0] = -t[1] / norm;
      ni[1] = t[0] / norm;
      ni[2] = 0.0;
    }
    return n;
  }
  case CellType::triangle:
  {
    for (std::size_t i = 0; i < entities.size(); ++i)
    {
      // Get the three vertices as points
      std::array vertices = {geometry_entities[i * eshape[1] + 0],
                             geometry_entities[i * eshape[1] + 1],
                             geometry_entities[i * eshape[1] + 2]};
      std::array p = {std::span<const T, 3>(x.data() + 3 * vertices[0], 3),
                      std::span<const T, 3>(x.data() + 3 * vertices[1], 3),
                      std::span<const T, 3>(x.data() + 3 * vertices[2], 3)};
 
      // Compute (p1 - p0) and (p2 - p0)
      std::array<T, 3> dp1, dp2;
      std::ranges::transform(p[1], p[0], dp1.begin(),
                             [](auto x, auto y) { return x - y; });
      std::ranges::transform(p[2], p[0], dp2.begin(),
                             [](auto x, auto y) { return x - y; });
 
      // Define cell normal via cross product of first two edges
      std::array<T, 3> ni = math::cross(dp1, dp2);
      T norm = std::sqrt(ni[0] * ni[0] + ni[1] * ni[1] + ni[2] * ni[2]);
      std::ranges::transform(ni, std::next(n.begin(), 3 * i),
                             [norm](auto x) { return x / norm; });
    }
 
    return n;
  }
  case CellType::quadrilateral:
  {
    // TODO: check
    for (std::size_t i = 0; i < entities.size(); ++i)
    {
      // Get the three vertices as points
      std::array vertices = {geometry_entities[i * eshape[1] + 0],
                             geometry_entities[i * eshape[1] + 1],
                             geometry_entities[i * eshape[1] + 2]};
      std::array p = {std::span<const T, 3>(x.data() + 3 * vertices[0], 3),
                      std::span<const T, 3>(x.data() + 3 * vertices[1], 3),
                      std::span<const T, 3>(x.data() + 3 * vertices[2], 3)};
 
      // Compute (p1 - p0) and (p2 - p0)
      std::array<T, 3> dp1, dp2;
      std::ranges::transform(p[1], p[0], dp1.begin(),
                             [](auto x, auto y) { return x - y; });
      std::ranges::transform(p[2], p[0], dp2.begin(),
                             [](auto x, auto y) { return x - y; });
 
      // Define cell normal via cross product of first two edges
      std::array<T, 3> ni = math::cross(dp1, dp2);
      T norm = std::sqrt(ni[0] * ni[0] + ni[1] * ni[1] + ni[2] * ni[2]);
      std::ranges::transform(ni, std::next(n.begin(), 3 * i),
                             [norm](auto x) { return x / norm; });
    }
 
    return n;
  }
  default:
    throw std::invalid_argument(
        "cell_normal not supported for this cell type.");
  }
}

template <std::floating_point T>
std::vector<T> compute_midpoints(const Mesh<T>& mesh, int dim,
                                 std::span<const std::int32_t> entities)
{
  if (entities.empty())
    return std::vector<T>();
 
  std::span<const T> x = mesh.geometry().x();
 
  // Build map from entity -> geometry dof
  const auto [e_to_g, eshape]
      = entities_to_geometry(mesh, dim, entities, false);
 
  std::vector<T> x_mid(entities.size() * 3, 0);
  for (std::size_t e = 0; e < entities.size(); ++e)
  {
    std::span<T, 3> p(x_mid.data() + 3 * e, 3);
    std::span<const std::int32_t> rows(e_to_g.data() + e * eshape[1],
                                       eshape[1]);
    for (auto row : rows)
    {
      std::span<const T, 3> xg(x.data() + 3 * row, 3);
      std::ranges::transform(p, xg, p.begin(),
                             [size = rows.size()](auto x, auto y)
                             { return x + y / size; });
    }
  }
 
  return x_mid;
}
 
namespace impl
{
template <std::floating_point T>
std::pair<std::vector<T>, std::array<std::size_t, 2>>
compute_vertex_coords(const mesh::Mesh<T>& mesh)
{
  auto topology = mesh.topology();
  assert(topology);
  const int tdim = topology->dim();
 
  // Create entities and connectivities
 
  // Get all vertex 'node' indices
  const std::int32_t num_vertices = topology->index_map(0)->size_local()
                                    + topology->index_map(0)->num_ghosts();
 
  std::vector<std::int32_t> vertex_to_node(num_vertices);
  for (int cell_type_idx = 0,
           num_cell_types = topology->entity_types(tdim).size();
       cell_type_idx < num_cell_types; ++cell_type_idx)
  {
    auto x_dofmap = mesh.geometry().dofmaps().at(cell_type_idx);
    auto c_to_v = topology->connectivity({tdim, cell_type_idx}, {0, 0});
    assert(c_to_v);
    for (int c = 0; c < c_to_v->num_nodes(); ++c)
    {
      auto x_dofs = md::submdspan(x_dofmap, c, md::full_extent);
      auto vertices = c_to_v->links(c);
      for (std::size_t i = 0; i < vertices.size(); ++i)
        vertex_to_node[vertices[i]] = x_dofs[i];
    }
  }
 
  // Pack coordinates of vertices
  std::span<const T> x_nodes = mesh.geometry().x();
  std::vector<T> x_vertices(3 * vertex_to_node.size(), 0.0);
  for (std::size_t i = 0; i < vertex_to_node.size(); ++i)
  {
    std::int32_t pos = 3 * vertex_to_node[i];
    for (std::size_t j = 0; j < 3; ++j)
      x_vertices[j * vertex_to_node.size() + i] = x_nodes[pos + j];
  }
 
  return {std::move(x_vertices), {3, vertex_to_node.size()}};
}
 
} // namespace impl

template <typename Fn, typename T>
concept MarkerFn = std::is_invocable_r<
    std::vector<std::int8_t>, Fn,
    md::mdspan<const T,
               md::extents<std::size_t, 3, md::dynamic_extent>>>::value;

template <std::floating_point T, MarkerFn<T> U>
std::vector<std::int32_t> locate_entities(const Mesh<T>& mesh, int dim,
                                          U marker, int entity_type_idx)
{
 
  using cmdspan3x_t
      = md::mdspan<const T, md::extents<std::size_t, 3, md::dynamic_extent>>;
 
  // Run marker function on vertex coordinates
  const auto [xdata, xshape] = impl::compute_vertex_coords(mesh);
 
  cmdspan3x_t x(xdata.data(), xshape);
  const std::vector<std::int8_t> marked = marker(x);
  if (marked.size() != x.extent(1))
    throw std::invalid_argument("Length of array of markers is wrong.");
 
  auto topology = mesh.topology();
  assert(topology);
  const int tdim = topology->dim();
 
  if (entity_type_idx < 0
      or static_cast<std::size_t>(entity_type_idx)
             >= topology->entity_types(dim).size())
  {
    throw std::out_of_range(
        "entity_type_idx out of range for Topology::entity_types(dim).");
  }
 
  mesh.topology_mutable()->create_entities(dim);
  if (dim < tdim)
    mesh.topology_mutable()->create_connectivity(dim, 0);
 
  // Iterate over entities of dimension 'dim' to build vector of marked
  // entities
  auto e_to_v = topology->connectivity({dim, entity_type_idx}, {0, 0});
  assert(e_to_v);
  std::vector<std::int32_t> entities;
  for (int e = 0; e < e_to_v->num_nodes(); ++e)
  {
    // Iterate over entity vertices
    bool all_vertices_marked = true;
    for (std::int32_t v : e_to_v->links(e))
    {
      if (!marked[v])
      {
        all_vertices_marked = false;
        break;
      }
    }
 
    if (all_vertices_marked)
      entities.push_back(e);
  }
 
  return entities;
}

template <std::floating_point T, MarkerFn<T> U>
std::vector<std::int32_t> locate_entities(const Mesh<T>& mesh, int dim,
                                          U marker)
{
  const int num_entity_types = mesh.topology()->entity_types(dim).size();
  if (num_entity_types > 1)
  {
    throw std::runtime_error(
        "Multiple entity types of this dimension. Specify entity type index");
  }
  return locate_entities(mesh, dim, marker, 0);
}

template <std::floating_point T, MarkerFn<T> U>
std::vector<std::int32_t> locate_entities_boundary(const Mesh<T>& mesh, int dim,
                                                   U marker)
{
  // TODO Rewrite this function, it should be possible to simplify considerably
  auto topology = mesh.topology();
  assert(topology);
  int tdim = topology->dim();
  if (dim == tdim)
  {
    throw std::invalid_argument(
        "Cannot use mesh::locate_entities_boundary (boundary) for cells.");
  }
 
  // Compute list of boundary facets
  mesh.topology_mutable()->create_entities(tdim - 1);
  mesh.topology_mutable()->create_connectivity(tdim - 1, tdim);
  std::vector<std::int32_t> boundary_facets = exterior_facet_indices(*topology);
  mesh.topology_mutable()->create_entities(dim);
 
  using cmdspan3x_t
      = md::mdspan<const T, md::extents<std::size_t, 3, md::dynamic_extent>>;
 
  // Run marker function on the vertex coordinates
  auto [facet_entities, xdata, vertex_to_pos]
      = impl::compute_vertex_coords_boundary(mesh, dim, boundary_facets);
  cmdspan3x_t x(xdata.data(), 3, xdata.size() / 3);
  std::vector<std::int8_t> marked = marker(x);
  if (marked.size() != x.extent(1))
    throw std::invalid_argument("Length of array of markers is wrong.");
 
  // Loop over entities and check vertex markers
  auto e_to_v = topology->connectivity(dim, 0);
  assert(e_to_v);
  std::vector<std::int32_t> entities;
  for (auto e : facet_entities)
  {
    // Iterate over entity vertices
    bool all_vertices_marked = true;
    for (auto v : e_to_v->links(e))
    {
      const std::int32_t pos = vertex_to_pos[v];
      if (!marked[pos])
      {
        all_vertices_marked = false;
        break;
      }
    }
 
    // Mark facet with all vertices marked
    if (all_vertices_marked)
      entities.push_back(e);
  }
 
  return entities;
}

template <std::floating_point T>
std::pair<std::vector<std::int32_t>, std::array<std::size_t, 2>>
entities_to_geometry(const Mesh<T>& mesh, int dim,
                     std::span<const std::int32_t> entities,
                     bool permute = false)
{
  auto topology = mesh.topology();
  assert(topology);
  CellType cell_type = topology->cell_type();
  if ((cell_type == CellType::prism or cell_type == CellType::pyramid)
      and dim == 2)
  {
    throw std::invalid_argument("mesh::entities_to_geometry for prism/pyramid "
                                "cell facets not yet supported.");
  }
 
  const int tdim = topology->dim();
  const Geometry<T>& geometry = mesh.geometry();
  auto xdofs = geometry.dofmaps().front();
 
  // Get the DOF layout and the number of DOFs per entity
  const fem::CoordinateElement<T>& coord_ele = geometry.cmaps().front();
  if (dim < tdim and coord_ele.is_discontinuous())
  {
    throw std::invalid_argument(
        "mesh::entities_to_geometry for sub-entities of a cell is not "
        "supported for a discontinuous geometry, which has no coordinate "
        "degrees-of-freedom associated with sub-entities.");
  }
 
  const fem::ElementDofLayout layout = coord_ele.create_dof_layout();
  const std::size_t num_entity_dofs = layout.entity_closure_dofs(dim, 0).size();
  std::vector<std::int32_t> entity_xdofs;
  entity_xdofs.reserve(entities.size() * num_entity_dofs);
  std::array<std::size_t, 2> eshape{entities.size(), num_entity_dofs};
 
  // Get the element's closure DOFs
  const std::vector<std::vector<std::vector<int>>>& closure_dofs_all
      = layout.entity_closure_dofs_all();
 
  // Special case when dim == tdim (cells)
  if (dim == tdim)
  {
    for (std::int32_t c : entities)
    {
      // Extract degrees of freedom
      auto x_c = md::submdspan(xdofs, c, md::full_extent);
      for (std::int32_t entity_dof : closure_dofs_all[tdim][0])
        entity_xdofs.push_back(x_c[entity_dof]);
    }
 
    return {std::move(entity_xdofs), eshape};
  }
 
  assert(dim != tdim);
 
  auto e_to_c = topology->connectivity(dim, tdim);
  if (!e_to_c)
  {
    throw std::runtime_error(std::format(
        "Entity-to-cell connectivity has not been computed. Missing dims "
        "{}->{}",
        dim, tdim));
  }
 
  auto c_to_e = topology->connectivity(tdim, dim);
  if (!c_to_e)
  {
    throw std::runtime_error(std::format(
        "Cell-to-entity connectivity has not been computed. Missing dims "
        "{}->{}",
        tdim, dim));
  }
 
  // Get the cell info, which is needed to permute the closure dofs
  std::span<const std::uint32_t> cell_info;
  if (permute)
    cell_info = std::span(mesh.topology()->get_cell_permutation_info());
 
  // Closure DOF count is the same for every local entity of dimension
  // `dim` (the one case where it isn't, prism/pyramid facets, is
  // rejected above), so size once and reuse across entities.
  std::vector<std::int32_t> closure_dofs(num_entity_dofs);
  for (std::int32_t e : entities)
  {
    // Get a cell connected to the entity
    if (e_to_c->links(e).empty())
      throw std::runtime_error("No cell incident to entity.");
    std::int32_t c = e_to_c->links(e).front();
 
    // Get the local index of the entity
    std::span<const std::int32_t> cell_entities = c_to_e->links(c);
    auto it = std::find(cell_entities.begin(), cell_entities.end(), e);
    assert(it != cell_entities.end());
    std::size_t local_entity = std::ranges::distance(cell_entities.begin(), it);
 
    // Cell sub-entities must be permuted so that their local
    // orientation agrees with their global orientation
    const std::vector<int>& e_closure_dofs
        = closure_dofs_all[dim][local_entity];
    assert(e_closure_dofs.size() == closure_dofs.size());
    std::ranges::copy(e_closure_dofs, closure_dofs.begin());
    if (permute)
    {
      mesh::CellType entity_type
          = mesh::cell_entity_type(cell_type, dim, local_entity);
      coord_ele.permute_subentity_closure(closure_dofs, cell_info[c],
                                          entity_type, local_entity);
    }
 
    // Extract degrees of freedom
    auto x_c = md::submdspan(xdofs, c, md::full_extent);
    for (std::int32_t entity_dof : closure_dofs)
      entity_xdofs.push_back(x_c[entity_dof]);
  }
 
  return {std::move(entity_xdofs), eshape};
}

std::vector<std::int32_t>
compute_incident_entities(const Topology& topology,
                          std::span<const std::int32_t> entities, int d0,
                          int d1);
 
namespace impl
{
template <typename F>
void try_locally(F&& fn, std::exception_ptr& error)
{
  if (error)
    return;
 
  try
  {
    fn();
  }
  catch (...)
  {
    error = std::current_exception();
  }
}

inline void mpi_check(MPI_Comm comm, std::string_view op_name,
                      const std::exception_ptr& error)
{
  int failed = error ? 1 : 0;
  int any_failed = 0;
  MPI_Allreduce(&failed, &any_failed, 1, MPI_INT, MPI_MAX, comm);
  if (error)
    std::rethrow_exception(error);
  else if (any_failed)
  {
    throw std::runtime_error(
        std::format("{} failed on another rank.", op_name));
  }
}

template <std::floating_point T>
std::vector<double> cell_centroids_local(
    std::span<const int> num_vertices_per_cell,
    const std::vector<std::span<const std::int64_t>>& cells,
    const boost::unordered_flat_map<std::int64_t, std::size_t>& node_to_pos,
    std::span<const T> coords, int gdim)
{
  std::size_t num_cells = 0;
  for (std::size_t i = 0; i < cells.size(); ++i)
    num_cells += cells[i].size() / num_vertices_per_cell[i];
  std::vector<double> centroid(gdim * num_cells, 0);
 
  std::size_t c0 = 0;
  for (std::size_t i = 0; i < cells.size(); ++i)
  {
    const int nv = num_vertices_per_cell[i];
    const double w = 1.0 / nv;
    for (std::size_t c = 0; c < cells[i].size() / nv; ++c)
    {
      for (int v = 0; v < nv; ++v)
      {
        auto it = node_to_pos.find(cells[i][nv * c + v]);
        assert(it != node_to_pos.end());
        std::size_t pos = it->second;
        for (int d = 0; d < gdim; ++d)
          centroid[gdim * (c0 + c) + d] += w * coords[gdim * pos + d];
      }
    }
 
    c0 += cells[i].size() / nv;
  }
 
  return centroid;
}

template <std::floating_point T>
std::vector<double>
compute_cell_centroids(MPI_Comm comm,
                       std::span<const int> num_vertices_per_cell,
                       const std::vector<std::span<const std::int64_t>>& cells,
                       MPI_Comm commg, std::span<const T> x, int gdim)
{
  // Vertices of the cells on this rank, sorted and with duplicates
  // removed, and the coordinates for them
  std::vector<std::int64_t> nodes;
  {
    std::size_t size = 0;
    for (std::span<const std::int64_t> c : cells)
      size += c.size();
    nodes.reserve(size);
    for (std::span<const std::int64_t> c : cells)
      nodes.insert(nodes.end(), c.begin(), c.end());
    dolfinx::radix_sort(nodes);
    auto [unique_end, range_end] = std::ranges::unique(nodes);
    nodes.erase(unique_end, range_end);
  }
  const std::vector<T> coords
      = dolfinx::MPI::distribute_data(comm, nodes, commg, x, gdim);
 
  // Hash map from global vertex index to its position in `nodes` (and
  // so its row in `coords`), turning the many repeated cell-vertex
  // lookups below into an O(1) average lookup rather than an
  // O(log(nodes.size())) binary search each time -- most vertices are
  // shared by several cells, so the same key is looked up repeatedly.
  boost::unordered_flat_map<std::int64_t, std::size_t> node_to_pos;
  node_to_pos.reserve(nodes.size());
  for (std::size_t i = 0; i < nodes.size(); ++i)
    node_to_pos.emplace(nodes[i], i);
 
  return cell_centroids_local(num_vertices_per_cell, cells, node_to_pos,
                              std::span<const T>(coords), gdim);
}

template <std::floating_point T>
std::tuple<std::vector<std::vector<std::int64_t>>,
           std::vector<std::vector<std::int64_t>>,
           std::vector<std::vector<int>>>
partition_cells(MPI_Comm comm, MPI_Comm commt,
                const std::vector<std::span<const std::int64_t>>& cells,
                const std::vector<CellType>& celltypes,
                const std::vector<fem::ElementDofLayout>& doflayouts,
                bool p1_geometry, const graph::Partitioner& partitioner,
                bool ghosting,
                std::optional<std::int32_t> max_facet_to_cell_links,
                int num_threads, MPI_Comm commg, std::span<const T> x,
                std::array<std::size_t, 2> xshape)
{
  const std::int32_t num_cell_types = cells.size();
  std::vector<std::vector<std::int64_t>> cells1(num_cell_types);
  std::vector<std::vector<std::int64_t>> original_idx1(num_cell_types);
  std::vector<std::vector<int>> ghost_owners(num_cell_types);
  if (graph::has_partitioner(partitioner.fn))
  {
    logger().info("Using partitioner with cell data ({} cell types)",
                  num_cell_types);
    graph::AdjacencyList<std::int32_t> dest(0);
    std::exception_ptr error;
 
    // Geometric data can be distributed on ranks that do not participate in
    // topology partitioning. Gather cell centroids collectively over `comm`
    // so that every rank in `commg` participates in the coordinate exchange.
    std::vector<double> centroid;
    const bool needs_centroids
        = std::holds_alternative<graph::geom_partition_fn>(partitioner.fn)
          or std::holds_alternative<graph::hybrid_partition_fn>(partitioner.fn);
    std::vector<std::vector<std::int64_t>> topology(num_cell_types);
    std::vector<std::span<const std::int64_t>> topology_view(num_cell_types);
    if (needs_centroids or commt != MPI_COMM_NULL)
    {
      for (std::int32_t i = 0; i < num_cell_types; ++i)
      {
        if (p1_geometry)
          topology_view[i] = cells[i];
        else
        {
          topology[i] = extract_topology(celltypes[i], doflayouts[i], cells[i]);
          topology_view[i] = topology[i];
        }
      }
    }
 
    if (needs_centroids)
    {
      std::vector<int> num_vertices_per_cell;
      num_vertices_per_cell.reserve(celltypes.size());
      std::ranges::transform(celltypes,
                             std::back_inserter(num_vertices_per_cell),
                             [](CellType c) { return num_cell_vertices(c); });
      centroid = compute_cell_centroids(comm, num_vertices_per_cell,
                                        topology_view, commg, x, xshape[1]);
    }
 
    if (std::holds_alternative<graph::geom_partition_fn>(partitioner.fn))
    {
      try_locally(
          [&dest, &comm, &centroid, &partitioner, &xshape]
          {
            int size = dolfinx::MPI::size(comm);
            const auto& p = std::get<graph::geom_partition_fn>(partitioner.fn);
            dest = graph::regular_adjacency_list(
                p(comm, size, std::span<const double>(centroid), xshape[1],
                  partitioner.node_weights),
                1);
          },
          error);
    }
 
    if (commt != MPI_COMM_NULL)
    {
      // A partitioner such as graph::parmetis::geom_partitioner (which
      // requires nparts to equal the number of ranks calling it) can
      // throw only on the ranks with commt != MPI_COMM_NULL, which may
      // be a strict subset of comm (e.g. cells built on rank 0 only).
      // mpi_check below turns that into a comm-wide decision before any
      // rank reaches the graph::build::distribute collective, or a throw
      // here would leave the rest of comm blocked on it forever.
      try_locally(
          [&dest, &comm, &commt, &celltypes, &topology_view,
           &max_facet_to_cell_links, &num_threads, &partitioner, &centroid,
           &ghosting]
          {
            int size = dolfinx::MPI::size(comm);
            // Shared by the graph::partition_fn and
            // graph::hybrid_partition_fn alternatives below: neither has any
            // other way to obtain the mesh dual graph.
            auto dual_graph
                = [&commt, &celltypes, &topology_view, &max_facet_to_cell_links,
                   &num_threads]() -> graph::AdjacencyList<std::int64_t>
            {
              return build_dual_graph(commt, celltypes, topology_view,
                                      max_facet_to_cell_links, num_threads);
            };
 
            // `dest` is already correct for graph::geom_partition_fn (set
            // above); skip the visit's dispatch (and the AdjacencyList
            // copy an unconditional assignment would cost) for that case.
            if (!std::holds_alternative<graph::geom_partition_fn>(
                    partitioner.fn))
            {
              dest = std::visit(
                  [&dual_graph, &commt, size, &centroid, &partitioner,
                   &ghosting](
                      const auto& p) -> graph::AdjacencyList<std::int32_t>
                  {
                    using P = std::decay_t<decltype(p)>;
                    if constexpr (std::is_same_v<P, graph::hybrid_partition_fn>)
                    {
                      return p(commt, size, dual_graph(),
                               std::span<const double>(centroid),
                               partitioner.node_weights, std::nullopt,
                               ghosting);
                    }
                    else if constexpr (std::is_same_v<P, graph::partition_fn>)
                    {
                      return p(commt, size, dual_graph(),
                               partitioner.node_weights, std::nullopt,
                               ghosting);
                    }
                    else
                    {
                      // std::visit still requires this branch to
                      // compile for graph::geom_partition_fn.
                      static_assert(
                          std::is_same_v<P, graph::geom_partition_fn>);
                      throw std::logic_error("Unreachable.");
                    }
                  },
                  partitioner.fn);
            }
          },
          error);
    }
 
    mpi_check(comm, "Cell partitioning", error);
 
    std::int32_t cell_offset = 0;
    for (std::int32_t i = 0; i < num_cell_types; ++i)
    {
      std::size_t num_cell_nodes = doflayouts[i].num_dofs();
      std::size_t num_cells = cells[i].size() / num_cell_nodes;
 
      // Extract destination AdjacencyList for this cell type
      std::vector<std::int32_t> offsets_i(
          std::next(dest.offsets().begin(), cell_offset),
          std::next(dest.offsets().begin(), cell_offset + num_cells + 1));
      std::vector<std::int32_t> data_i(
          std::next(dest.array().begin(), offsets_i.front()),
          std::next(dest.array().begin(), offsets_i.back()));
      const std::int32_t offset_0 = offsets_i.front();
      std::ranges::transform(offsets_i, offsets_i.begin(),
                             [offset_0](std::int32_t j)
                             { return j - offset_0; });
      graph::AdjacencyList<std::int32_t> dest_i(data_i, offsets_i);
      cell_offset += num_cells;
 
      // Distribute cells (topology, includes higher-order 'nodes') to
      // destination rank
      std::vector<int> src_ranks;
      std::tie(cells1[i], src_ranks, original_idx1[i], ghost_owners[i])
          = graph::build::distribute(comm, cells[i],
                                     {num_cells, num_cell_nodes}, dest_i);
      logger().debug("Got {} cells from distribution", cells1[i].size());
    }
  }
  else // No partitioning: keep cells on their current rank
  {
    // Count cells of each type on this rank. Each cell still needs a
    // globally unique index (assigned below), even though it is not
    // being redistributed, and the counts are needed first to size
    // `original_idx1` and to determine this rank's share via the
    // exclusive scan that follows.
    std::int64_t num_owned = 0;
    for (std::int32_t i = 0; i < num_cell_types; ++i)
    {
      cells1[i] = std::vector<std::int64_t>(cells[i].begin(), cells[i].end());
      std::int32_t num_cell_nodes = doflayouts[i].num_dofs();
      original_idx1[i].resize(cells1[i].size() / num_cell_nodes);
      num_owned += original_idx1[i].size();
    }
 
    // Assign a globally unique index to each cell. `global_offset`
    // starts as the number of cells owned by lower-ranked processes
    // (from the exclusive scan), and is advanced by each cell type's
    // count in turn so that the numbering is contiguous across cell
    // types too.
    std::int64_t global_offset = 0;
    MPI_Exscan(&num_owned, &global_offset, 1, MPI_INT64_T, MPI_SUM, comm);
    for (std::int32_t i = 0; i < num_cell_types; ++i)
    {
      std::iota(original_idx1[i].begin(), original_idx1[i].end(),
                global_offset);
      global_offset += original_idx1[i].size();
    }
  }
 
  return {std::move(cells1), std::move(original_idx1), std::move(ghost_owners)};
}
} // namespace impl

template <typename U>
Mesh<typename std::remove_reference_t<typename U::value_type>> create_mesh(
    MPI_Comm comm, MPI_Comm commt,
    std::vector<std::span<const std::int64_t>> cells,
    const std::vector<fem::CoordinateElement<
        typename std::remove_reference_t<typename U::value_type>>>& elements,
    MPI_Comm commg, const U& x, std::array<std::size_t, 2> xshape,
    const graph::Partitioner& partitioner, GhostMode ghost_mode,
    std::optional<std::int32_t> max_facet_to_cell_links, int num_threads,
    const graph::Reorder& reorder_fn = graph::Reorder{})
{
  using T = typename std::remove_reference_t<typename U::value_type>;
 
  if (cells.size() != elements.size())
  {
    throw std::invalid_argument(
        "Number of cell arrays and elements must match.");
  }
  std::vector<CellType> celltypes;
  std::ranges::transform(elements, std::back_inserter(celltypes),
                         [](auto& e) { return e.cell_shape(); });
  std::vector<fem::ElementDofLayout> doflayouts;
  std::ranges::transform(elements, std::back_inserter(doflayouts),
                         [](auto& e) { return e.create_dof_layout(); });
  for (std::size_t i = 0; i < cells.size(); ++i)
  {
    if (cells[i].size() % doflayouts[i].num_dofs() != 0)
    {
      throw std::invalid_argument("Cell array size is not a multiple of the "
                                  "number of nodes per cell.");
    }
  }
 
  // Note: `extract_topology` extracts topology data, i.e. just the
  // vertices. For other elements the filtered lists may have 'gaps',
  // i.e. the indices might not be contiguous.
  //
  // For 'P1 geometry' the extraction is the identity operator, and cell
  // node data is used directly as cell topology. This avoids copies of
  // the (large) cell array, and lets the geometry node indices be taken
  // from the topology vertices rather than re-derived by sorting the
  // cell array (see below).
  const bool p1_geometry = std::ranges::all_of(
      std::views::iota(std::size_t(0), elements.size()),
      [&celltypes, &doflayouts](std::size_t i)
      { return is_vertex_dof_layout(celltypes[i], doflayouts[i]); });
 
  const std::int32_t num_cell_types = cells.size();
 
  // Partition cells across ranks of `comm` (or, if `partitioner` is not
  // callable, keep them on their current rank and just assign each a
  // globally unique index)
  const bool ghosting = (ghost_mode != GhostMode::none);
  auto [cells1, original_idx1, ghost_owners] = impl::partition_cells(
      comm, commt, cells, celltypes, doflayouts, p1_geometry, partitioner,
      ghosting, max_facet_to_cell_links, num_threads, commg,
      std::span<const T>(x), xshape);
 
  // Extract cell 'topology', i.e. extract the vertices for each cell
  // and discard any 'higher-order' nodes. `cells1_v_storage` is empty
  // for 'P1 geometry', where `cells1_v` views `cells1` directly.
  std::vector<std::vector<std::int64_t>> cells1_v_storage(num_cell_types);
  std::vector<std::span<std::int64_t>> cells1_v(num_cell_types);
  for (std::int32_t i = 0; i < num_cell_types; ++i)
  {
    if (p1_geometry)
      cells1_v[i] = cells1[i];
    else
    {
      cells1_v_storage[i]
          = extract_topology(celltypes[i], doflayouts[i], cells1[i]);
      cells1_v[i] = cells1_v_storage[i];
    }
 
    logger().info("Extract basic topology: {}->{}", cells1[i].size(),
                  cells1_v[i].size());
  }
 
  // Re-order cells and get boundary vertices. The re-ordering is done
  // on the cell topology, i.e. the vertex indices, and the higher-order
  // nodes are re-ordered accordingly.
  // Centroids are required only by a callable graph::reorder_geom_fn.
  // As elsewhere, `reorder_fn` is assumed to hold the same alternative
  // on every rank, so an empty function skips the collective exchange
  // everywhere; impl::reorder_cells then throws, and the failure is
  // made comm-wide by try_locally/mpi_check below.
  std::vector<double> cell_centroids;
  if (const graph::reorder_geom_fn* fn
      = std::get_if<graph::reorder_geom_fn>(&reorder_fn);
      fn and *fn)
  {
    std::vector<int> num_cell_vertices;
    num_cell_vertices.reserve(celltypes.size());
    std::ranges::transform(celltypes, std::back_inserter(num_cell_vertices),
                           [](CellType c)
                           { return mesh::num_cell_vertices(c); });
 
    std::vector<std::span<const std::int64_t>> cells1_v_owned;
    cells1_v_owned.reserve(cells1_v.size());
    for (std::size_t i = 0; i < cells1_v.size(); ++i)
    {
      const int num_cell_vertices_i = num_cell_vertices[i];
      const std::size_t num_owned_cells
          = cells1_v[i].size() / num_cell_vertices_i - ghost_owners[i].size();
      cells1_v_owned.emplace_back(cells1_v[i].data(),
                                  num_owned_cells * num_cell_vertices_i);
    }
 
    cell_centroids
        = impl::compute_cell_centroids(comm, num_cell_vertices, cells1_v_owned,
                                       commg, std::span<const T>(x), xshape[1]);
  }
 
  std::vector<std::int64_t> boundary_v;
  std::exception_ptr error;
  impl::try_locally(
      [&boundary_v, &reorder_fn, &cell_centroids, &xshape,
       &max_facet_to_cell_links, &celltypes, &doflayouts, &ghost_owners,
       &cells1, &cells1_v, &original_idx1, &num_threads]
      {
        boundary_v = impl::reorder_cells(reorder_fn, cell_centroids, xshape[1],
                                         max_facet_to_cell_links, celltypes,
                                         doflayouts, ghost_owners, cells1,
                                         cells1_v, original_idx1, num_threads);
      },
      error);
  impl::mpi_check(comm, "Cell reordering", error);
 
  logger().debug("Got {} boundary vertices", boundary_v.size());
 
  // Create Topology
  std::vector<std::span<const std::int64_t>> cells1_v_span(cells1_v.begin(),
                                                           cells1_v.end());
  std::vector<std::span<const std::int64_t>> original_idx1_span;
  std::ranges::transform(original_idx1, std::back_inserter(original_idx1_span),
                         [](auto& c) { return std::span(c); });
  std::vector<std::span<const int>> ghost_owners_span;
  std::ranges::transform(ghost_owners, std::back_inserter(ghost_owners_span),
                         [](auto& c) { return std::span(c); });
 
  // Note: `vertex_index` holds the sorted input global indices of the
  // topology vertices, which for 'P1 geometry' are exactly the geometry
  // node indices required below.
  auto [topology, vertex_index] = mesh::impl::create_topology(
      comm, celltypes, cells1_v_span, original_idx1_span, ghost_owners_span,
      boundary_v, num_threads);
 
  // Create connectivities required higher-order geometries for creating
  // a Geometry object
  for (int i = 0; i < num_cell_types; ++i)
  {
    const auto& entity_dofs = doflayouts[i].entity_dofs_all();
    for (int dim = 1; dim < topology.dim(); ++dim)
    {
      // Accumulate count of all dofs on this dimension
      int dim_sum
          = std::accumulate(entity_dofs[dim].begin(), entity_dofs[dim].end(), 0,
                            [](int c, auto v) { return c + v.size(); });
 
      logger().debug("Counting entity dofs, dim={}: {}", dim, dim_sum);
      if (dim_sum > 0)
        topology.create_entities(dim, num_threads);
    }
 
    if (elements[i].needs_dof_permutations())
      topology.create_cell_permutations(num_threads);
  }
 
  // Cell 'node' indices (global), as a single flat array. This is
  // `cells1` for a single cell type, and concatenated otherwise.
  std::vector<std::int64_t> nodes2_storage;
  std::span<const std::int64_t> nodes2;
  if (num_cell_types == 1)
    nodes2 = cells1.front();
  else
  {
    std::size_t size = 0;
    for (const std::vector<std::int64_t>& c : cells1)
      size += c.size();
    nodes2_storage.reserve(size);
    for (const std::vector<std::int64_t>& c : cells1)
      nodes2_storage.insert(nodes2_storage.end(), c.begin(), c.end());
    nodes2 = nodes2_storage;
  }
 
  // Sorted list of unique (global) node indices. For 'P1 geometry' the
  // nodes are the vertices, which `create_topology` has already sorted
  // and made unique, so re-deriving them from the (much larger) cell
  // array is avoided.
  std::vector<std::int64_t> nodes1;
  if (p1_geometry)
    nodes1 = std::move(vertex_index);
  else
  {
    nodes1.assign(nodes2.begin(), nodes2.end());
    dolfinx::radix_sort(nodes1);
    auto [unique_end, range_end] = std::ranges::unique(nodes1);
    nodes1.erase(unique_end, range_end);
  }
 
  std::vector coords
      = dolfinx::MPI::distribute_data(comm, nodes1, commg, x, xshape[1]);
 
  // Create geometry object
  Geometry geometry
      = create_geometry(topology, elements, nodes1, nodes2, coords, xshape[1]);
 
#ifndef NDEBUG
  // Nodes in `x` that no cell references are dropped silently.
  {
    std::int64_t num_nodes_local = xshape[0];
    std::int64_t num_nodes = 0;
    int err = MPI_Allreduce(&num_nodes_local, &num_nodes, 1,
                            dolfinx::MPI::mpi_t<std::int64_t>, MPI_SUM, comm);
    dolfinx::MPI::check_error(comm, err);
    if (std::int64_t num_used = geometry.index_map()->size_global();
        num_used != num_nodes and dolfinx::MPI::rank(comm) == 0)
    {
      logger().warn("{} of {} input geometry nodes are not referenced by any "
                    "cell and have been dropped.",
                    num_nodes - num_used, num_nodes);
    }
  }
#endif
 
  return Mesh(comm, std::make_shared<Topology>(std::move(topology)),
              std::move(geometry));
}

template <typename U>
Mesh<typename std::remove_reference_t<typename U::value_type>> create_mesh(
    MPI_Comm comm, MPI_Comm commt, std::span<const std::int64_t> cells,
    const fem::CoordinateElement<
        typename std::remove_reference_t<typename U::value_type>>& element,
    MPI_Comm commg, const U& x, std::array<std::size_t, 2> xshape,
    const graph::Partitioner& partitioner, GhostMode ghost_mode,
    std::optional<std::int32_t> max_facet_to_cell_links, int num_threads,
    const graph::Reorder& reorder_fn = graph::Reorder{})
{
  return create_mesh(comm, commt, std::vector{cells}, std::vector{element},
                     commg, x, xshape, partitioner, ghost_mode,
                     max_facet_to_cell_links, num_threads, reorder_fn);
}

template <typename U>
Mesh<typename std::remove_reference_t<typename U::value_type>>
create_mesh(MPI_Comm comm, std::span<const std::int64_t> cells,
            const fem::CoordinateElement<
                std::remove_reference_t<typename U::value_type>>& element,
            const U& x, std::array<std::size_t, 2> xshape, GhostMode ghost_mode,
            std::optional<std::int32_t> max_facet_to_cell_links = 2)
{
  // A single rank has nothing to partition, so skip the default
  // partitioner and just assign global indices.
  graph::Partitioner partitioner
      = dolfinx::MPI::size(comm) == 1
            ? graph::Partitioner{.fn = graph::partition_fn(nullptr)}
            : graph::Partitioner{};
  return create_mesh(comm, comm, std::vector{cells}, std::vector{element}, comm,
                     x, xshape, partitioner, ghost_mode,
                     max_facet_to_cell_links, 1);
}

template <std::floating_point T>
std::pair<Geometry<T>, std::vector<int32_t>>
create_subgeometry(const Mesh<T>& mesh, int dim,
                   std::span<const std::int32_t> subentity_to_entity)
{
  const Geometry<T>& geometry = mesh.geometry();
 
  // Get the geometry dofs in the sub-geometry based on the entities in
  // sub-geometry
  const std::vector<std::int32_t> x_indices
      = entities_to_geometry(mesh, dim, subentity_to_entity, true).first;
 
  std::vector<std::int32_t> sub_x_dofs = x_indices;
  std::ranges::sort(sub_x_dofs);
  auto [unique_end, range_end] = std::ranges::unique(sub_x_dofs);
  sub_x_dofs.erase(unique_end, range_end);
 
  // Get the sub-geometry dofs owned by this process
  auto x_index_map = geometry.index_map();
  assert(x_index_map);
 
  std::shared_ptr<common::IndexMap> sub_x_dof_index_map;
  std::vector<std::int32_t> subx_to_x_dofmap;
  {
    // An owner change is permitted here: a coordinate dof may be needed
    // by a sub-mesh cell on a ghosting rank but not on its owner.
    auto [map, new_to_old, owners_changed] = common::create_sub_index_map(
        *x_index_map, sub_x_dofs, common::IndexMapOrder::any);
    sub_x_dof_index_map = std::make_shared<common::IndexMap>(std::move(map));
    subx_to_x_dofmap = std::move(new_to_old);
  }
 
  // Create sub-geometry coordinates
  std::span<const T> x = geometry.x();
  std::int32_t sub_num_x_dofs = subx_to_x_dofmap.size();
  std::vector<T> sub_x(3 * sub_num_x_dofs);
  for (std::int32_t i = 0; i < sub_num_x_dofs; ++i)
  {
    std::copy_n(std::next(x.begin(), 3 * subx_to_x_dofmap[i]), 3,
                std::next(sub_x.begin(), 3 * i));
  }
 
  // Create geometry to sub-geometry  map
  std::vector<std::int32_t> x_to_subx_dof_map(
      x_index_map->size_local() + x_index_map->num_ghosts(), -1);
  for (std::size_t i = 0; i < subx_to_x_dofmap.size(); ++i)
    x_to_subx_dof_map[subx_to_x_dofmap[i]] = i;
 
  // Create sub-geometry dofmap
  std::vector<std::int32_t> sub_x_dofmap;
  sub_x_dofmap.reserve(x_indices.size());
  std::ranges::transform(x_indices, std::back_inserter(sub_x_dofmap),
                         [&x_to_subx_dof_map](auto x_dof)
                         {
                           assert(x_to_subx_dof_map[x_dof] != -1);
                           return x_to_subx_dof_map[x_dof];
                         });
 
  // Sub-geometry coordinate element
  CellType sub_xcell
      = cell_entity_type(geometry.cmaps().front().cell_shape(), dim, 0);
 
  // Special handling of point meshes, as they only support constant
  // basis functions
  int degree
      = (sub_xcell == CellType::point) ? 0 : geometry.cmaps().front().degree();
  fem::CoordinateElement<T> sub_cmap(sub_xcell, degree,
                                     geometry.cmaps().front().variant());
 
  // Sub-geometry input_global_indices
  const std::vector<std::int64_t>& igi = geometry.input_global_indices();
  std::vector<std::int64_t> sub_igi;
  sub_igi.reserve(subx_to_x_dofmap.size());
  std::ranges::transform(subx_to_x_dofmap, std::back_inserter(sub_igi),
                         [&igi](auto sub_x_dof) { return igi[sub_x_dof]; });
 
  // Create geometry
  return {Geometry(
              sub_x_dof_index_map,
              std::vector<std::vector<std::int32_t>>{std::move(sub_x_dofmap)},
              {sub_cmap}, std::move(sub_x), geometry.dim(), std::move(sub_igi)),
          std::move(subx_to_x_dofmap)};
}

template <std::floating_point T>
std::tuple<Mesh<T>, EntityMap, EntityMap, std::vector<std::int32_t>>
create_submesh(const Mesh<T>& mesh, int dim,
               std::span<const std::int32_t> entities)
{
  if (dim < 0 or dim > mesh.topology()->dim())
    throw std::invalid_argument("dim out of range for mesh topology.");
 
  // Create sub-topology
  mesh.topology_mutable()->create_connectivity(dim, 0);
  auto [topology, subentity_to_entity, subvertex_to_vertex]
      = mesh::create_subtopology(*mesh.topology(), dim, entities);
 
  // Create sub-geometry
  const int tdim = mesh.topology()->dim();
  mesh.topology_mutable()->create_entities(dim);
  mesh.topology_mutable()->create_connectivity(dim, tdim);
  mesh.topology_mutable()->create_connectivity(tdim, dim);
  mesh.topology_mutable()->create_cell_permutations();
  auto [geometry, subx_to_x_dofmap]
      = mesh::create_subgeometry(mesh, dim, subentity_to_entity);
 
  Mesh<T> submesh
      = Mesh(mesh.comm(), std::make_shared<Topology>(std::move(topology)),
             std::move(geometry));
  EntityMap entity_map(mesh.topology(), submesh.topology(), dim,
                       subentity_to_entity);
  EntityMap vertex_map(mesh.topology(), submesh.topology(), 0,
                       subvertex_to_vertex);
  return {std::move(submesh), std::move(entity_map), std::move(vertex_map),
          std::move(subx_to_x_dofmap)};
}

template <typename T>
MeshTags<T> transfer_meshtags_to_submesh(
    const MeshTags<T>& tags,
    std::shared_ptr<const dolfinx::mesh::Topology> submesh_topology,
    const EntityMap& cell_map, const EntityMap& vertex_map)
{
  int tag_dim = tags.dim();
  int submesh_tdim = submesh_topology->dim();
  auto topology = tags.topology();
  if (tag_dim > submesh_tdim)
  {
    throw std::invalid_argument("Tag dimension must be less than or equal to "
                                "submesh dimension");
  }
 
  // Validate that cell_map/vertex_map relate `topology` (the tags'
  // parent topology) to `submesh_topology`, and have the dimension
  // this function assumes.
  if (cell_map.dim() != submesh_tdim)
  {
    throw std::invalid_argument(
        "cell_map dimension must equal the submesh topology dimension.");
  }
  if (cell_map.topology() != topology)
  {
    throw std::invalid_argument(
        "cell_map topology must match tags.topology().");
  }
  if (cell_map.sub_topology() != submesh_topology)
  {
    throw std::invalid_argument(
        "cell_map sub_topology must match submesh_topology.");
  }
  if (vertex_map.dim() != 0)
    throw std::invalid_argument("vertex_map dimension must be 0.");
  if (vertex_map.topology() != topology)
  {
    throw std::invalid_argument(
        "vertex_map topology must match tags.topology().");
  }
  if (vertex_map.sub_topology() != submesh_topology)
  {
    throw std::invalid_argument(
        "vertex_map sub_topology must match submesh_topology.");
  }
  std::shared_ptr<const dolfinx::common::IndexMap> sub_cell_imap
      = submesh_topology->index_map(submesh_tdim);
  if (!sub_cell_imap)
  {
    throw std::runtime_error(
        std::format("Entities of dimension {} does not exist in mesh topology.",
                    submesh_tdim));
  }
 
  // Create a map from parent entity to submesh cell
  std::int32_t submesh_num_cells
      = sub_cell_imap->size_local() + sub_cell_imap->num_ghosts();
  auto sub_cells = std::ranges::views::iota(0, submesh_num_cells);
  std::vector<std::int32_t> sub_cell_to_parent_entity
      = cell_map.sub_topology_to_topology(sub_cells, false);
 
  // Create a full lookup for all cells on the parent mesh, as the tag can have
  // entities that are not in the submesh
  auto parent_entity_imap = topology->index_map(submesh_tdim);
  if (!parent_entity_imap)
  {
    throw std::runtime_error(std::format(
        "Entities of dimension {} does not exist in parent mesh topology.",
        submesh_tdim));
  }
  std::size_t num_parent_entities
      = parent_entity_imap->size_local() + parent_entity_imap->num_ghosts();
  std::vector<std::int32_t> parent_entity_to_sub_cell(num_parent_entities, -1);
  for (std::size_t i = 0; i < sub_cell_to_parent_entity.size(); ++i)
  {
    parent_entity_to_sub_cell[sub_cell_to_parent_entity[i]]
        = static_cast<std::int32_t>(i);
  }
 
  // Get map from submesh vertex to parent vertex
  std::vector<std::int32_t> sub_to_parent_vertex;
  {
    auto sub_vertex_map = submesh_topology->index_map(0);
    std::int32_t num_sub_vertices
        = sub_vertex_map->size_local() + sub_vertex_map->num_ghosts();
    auto sub_vertices = std::ranges::views::iota(0, num_sub_vertices);
 
    sub_to_parent_vertex
        = vertex_map.sub_topology_to_topology(sub_vertices, false);
  }
  // Access various connectivity maps
  auto sub_e_to_v = submesh_topology->connectivity(tag_dim, 0);
  auto sub_c_to_e = submesh_topology->connectivity(submesh_tdim, tag_dim);
  auto sub_entity_imap = submesh_topology->index_map(tag_dim);
  auto e_to_v = topology->connectivity(tag_dim, 0);
  std::shared_ptr<const dolfinx::graph::AdjacencyList<std::int32_t>>
      e_to_sub_cell = nullptr;
  if (tag_dim != submesh_tdim)
  {
    e_to_sub_cell = topology->connectivity(tag_dim, submesh_tdim);
    if (!e_to_sub_cell)
    {
      throw std::runtime_error(
          std::format("Missing connectivity between {} and {} in parent mesh",
                      tag_dim, submesh_tdim));
    }
  }
 
  if (!sub_e_to_v)
  {
    throw std::runtime_error(std::format(
        "Missing connectivity between {} and {} in submesh", tag_dim, 0));
  }
  if (!sub_c_to_e)
  {
    throw std::runtime_error(
        std::format("Missing connectivity between {} and {} in submesh",
                    submesh_tdim, tag_dim));
  }
  if (!sub_entity_imap)
  {
    throw std::runtime_error(std::format(
        "Entities of dimension {} does not exist in submesh topology.",
        tag_dim));
  }
  if (!e_to_v)
  {
    throw std::runtime_error(
        std::format("Missing connectivity between {} and 0", tag_dim));
  }
 
  // Prepare sub entity to parent map
  std::size_t num_sub_entities
      = sub_entity_imap->size_local() + sub_entity_imap->num_ghosts();
  std::vector<T> submesh_values(num_sub_entities);
  std::vector<std::int8_t> submesh_value_found(num_sub_entities, 0);
  std::vector<std::int32_t> submesh_indices(num_sub_entities);
  std::iota(submesh_indices.begin(), submesh_indices.end(), 0);
 
  std::span<const std::int32_t> tagged_entities = tags.indices();
  std::span<const T> tagged_values = tags.values();
  // For each entity in the tag, find all cells of the submesh connected to this
  // entity
  for (std::size_t i = 0; i < tagged_entities.size(); ++i)
  {
    auto find_and_map_sub_entity
        = [tag_dim, submesh_tdim, &e_to_v, &parent_entity_to_sub_cell,
           &sub_to_parent_vertex, &sub_e_to_v, &sub_c_to_e,
           &e_to_sub_cell](std::int32_t entity)
    {
      // Fast exit if the tag dimension is the same as the submesh dimension,
      // as we can directly map the parent entity to the submesh cell
      if (tag_dim == submesh_tdim)
        return parent_entity_to_sub_cell[entity];
 
      // Given an entity in the parent meshtag, find all submesh-cells that are
      // entities in parent mesh that contain this entity.
      auto entity_vertices = e_to_v->links(entity);
      auto parent_sub_cells = e_to_sub_cell->links(entity);
      auto submesh_cells
          = parent_sub_cells
            | std::views::transform([&parent_entity_to_sub_cell](auto c)
                                    { return parent_entity_to_sub_cell[c]; })
            | std::views::filter([](auto sub_cell) { return sub_cell != -1; });
      for (auto sub_cell : submesh_cells)
      {
        for (auto sub_entity : sub_c_to_e->links(sub_cell))
        {
          // Convert submesh entity vertices to parent vertices
          auto parent_vertices
              = sub_e_to_v->links(sub_entity)
                | std::views::transform([&sub_to_parent_vertex](auto v)
                                        { return sub_to_parent_vertex[v]; });
 
          // Check if all parent vertices of the submesh entity are in the
          // parent entity
          bool entity_matches = std::ranges::all_of(
              parent_vertices,
              [&entity_vertices](auto p_v)
              {
                // With C++23 this can use std::ranges::contains
                return std::ranges::find(entity_vertices, p_v)
                       != std::ranges::end(entity_vertices);
              });
 
          // If a match is found, apply values and exit the lambda immediately
          if (entity_matches)
            return sub_entity;
        }
      }
      return -1;
    };
 
    // Execute the search for the current entity
    std::int32_t sub_entity = find_and_map_sub_entity(tagged_entities[i]);
    if (sub_entity != -1)
    {
      submesh_values[sub_entity] = tagged_values[i];
      submesh_value_found[sub_entity] = 1;
    }
  }
 
  // Filter out the entities that were never mapped. Tracked with a
  // separate flag rather than a sentinel value, since any T value
  // (including numeric_limits<T>::max()) is a legitimate tag value.
  std::vector<std::int32_t> filtered_indices;
  std::vector<T> filtered_values;
  filtered_indices.reserve(num_sub_entities);
  filtered_values.reserve(num_sub_entities);
  for (std::size_t i = 0; i < submesh_values.size(); ++i)
  {
    if (submesh_value_found[i])
    {
      filtered_indices.push_back(submesh_indices[i]);
      filtered_values.push_back(submesh_values[i]);
    }
  }
  filtered_indices.shrink_to_fit();
  filtered_values.shrink_to_fit();
  MeshTags<T> new_meshtag(submesh_topology, tag_dim, filtered_indices,
                          filtered_values, tags.name());
  return new_meshtag;
}
 
} // namespace dolfinx::mesh