generation.h

// Copyright (C) 2005-2024 Anders Logg and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Mesh.h"
#include "cell_types.h"
#include "types.h"
#include "utils.h"
#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <dolfinx/common/Timer.h>
#include <dolfinx/common/local_range.h>
#include <dolfinx/fem/CoordinateElement.h>
#include <dolfinx/graph/ordering.h>
#include <dolfinx/graph/partition.h>
#include <limits>
#include <mpi.h>
#include <optional>
#include <span>
#include <stdexcept>
#include <tuple>
#include <utility>
#include <vector>
 
namespace dolfinx::mesh
{
enum class DiagonalType
{
  left,
  right,
  crossed,
  shared_facet,
  left_right,
  right_left
};
 
namespace impl
{
template <std::floating_point T>
std::tuple<std::vector<T>, std::vector<std::int64_t>>
create_interval_cells(std::array<T, 2> p, std::int64_t n);
 
template <std::floating_point T>
Mesh<T> build_tri(MPI_Comm comm, std::array<std::array<T, 2>, 2> p,
                  std::array<std::int64_t, 2> n,
                  const graph::AnyPartitionFunction& partitioner,
                  DiagonalType diagonal, GhostMode ghost_mode,
                  const graph::Reorder& reorder_fn, int gdim);
 
template <std::floating_point T>
Mesh<T> build_quad(MPI_Comm comm, std::array<std::array<T, 2>, 2> p,
                   std::array<std::int64_t, 2> n,
                   const graph::AnyPartitionFunction& partitioner,
                   GhostMode ghost_mode, const graph::Reorder& reorder_fn,
                   int gdim);
 
template <std::floating_point T>
std::vector<T> create_geom(MPI_Comm comm, std::array<std::array<T, 3>, 2> p,
                           std::array<std::int64_t, 3> n);
 
template <std::floating_point T>
Mesh<T> build_tet(MPI_Comm comm, MPI_Comm subcomm,
                  std::array<std::array<T, 3>, 2> p,
                  std::array<std::int64_t, 3> n,
                  const graph::AnyPartitionFunction& partitioner,
                  GhostMode ghost_mode, const graph::Reorder& reorder_fn);
 
template <std::floating_point T>
Mesh<T> build_hex(MPI_Comm comm, MPI_Comm subcomm,
                  std::array<std::array<T, 3>, 2> p,
                  std::array<std::int64_t, 3> n,
                  const graph::AnyPartitionFunction& partitioner,
                  GhostMode ghost_mode, const graph::Reorder& reorder_fn);
 
template <std::floating_point T>
Mesh<T> build_prism(MPI_Comm comm, MPI_Comm subcomm,
                    std::array<std::array<T, 3>, 2> p,
                    std::array<std::int64_t, 3> n,
                    const graph::AnyPartitionFunction& partitioner,
                    GhostMode ghost_mode, const graph::Reorder& reorder_fn);

template <typename U>
Mesh<typename std::remove_reference_t<typename U::value_type>> finalize_mesh(
    MPI_Comm comm, MPI_Comm commx, std::span<const std::int64_t> cells,
    const fem::CoordinateElement<
        typename std::remove_reference_t<typename U::value_type>>& element,
    const U& x, std::array<std::size_t, 2> xshape,
    const graph::AnyPartitionFunction& partitioner, GhostMode ghost_mode,
    const graph::Reorder& reorder_fn)
{
  return create_mesh(comm, commx, cells, element, commx, x, xshape,
                     graph::Partitioner{.fn = partitioner}, ghost_mode, 2, 1,
                     reorder_fn);
}
} // namespace impl

template <std::floating_point T = double>
Mesh<T> create_box(MPI_Comm comm, MPI_Comm subcomm,
                   std::array<std::array<T, 3>, 2> p,
                   std::array<std::int64_t, 3> n, CellType celltype,
                   graph::AnyPartitionFunction partitioner = {},
                   GhostMode ghost_mode = GhostMode::none,
                   const graph::Reorder& reorder_fn = graph::Reorder{})
{
  if (std::ranges::any_of(n, [](auto e) { return e < 1; }))
    throw std::invalid_argument("At least one cell is required.");
 
  for (int32_t i = 0; i < 3; i++)
  {
    if (p[0][i] >= p[1][i])
      throw std::invalid_argument("It must hold p[0] < p[1].");
  }
 
  for (int32_t i = 0; i < 3; i++)
  {
    if (std::abs(p[1][i] - p[0][i]) / static_cast<T>(n[i])
        < 2.0 * std::numeric_limits<T>::epsilon())
    {
      throw std::invalid_argument(
          "Box seems to have zero width, height or depth. Check dimensions");
    }
  }
 
  if (!graph::has_partitioner(partitioner) and dolfinx::MPI::size(comm) > 1)
    partitioner = graph::partition_graph;
 
  switch (celltype)
  {
  case CellType::tetrahedron:
    return impl::build_tet<T>(comm, subcomm, p, n, partitioner, ghost_mode,
                              reorder_fn);
  case CellType::hexahedron:
    return impl::build_hex<T>(comm, subcomm, p, n, partitioner, ghost_mode,
                              reorder_fn);
  case CellType::prism:
    return impl::build_prism<T>(comm, subcomm, p, n, partitioner, ghost_mode,
                                reorder_fn);
  default:
    throw std::invalid_argument("Generate box mesh. Wrong cell type");
  }
}

template <std::floating_point T = double>
Mesh<T> create_box(MPI_Comm comm, std::array<std::array<T, 3>, 2> p,
                   std::array<std::int64_t, 3> n, CellType celltype,
                   const graph::AnyPartitionFunction& partitioner = {},
                   GhostMode ghost_mode = GhostMode::none,
                   const graph::Reorder& reorder_fn = graph::Reorder{})
{
  return create_box<T>(comm, comm, p, n, celltype, partitioner, ghost_mode,
                       reorder_fn);
}

template <std::floating_point T = double>
Mesh<T> create_rectangle(MPI_Comm comm, std::array<std::array<T, 2>, 2> p,
                         std::array<std::int64_t, 2> n, CellType celltype,
                         graph::AnyPartitionFunction partitioner,
                         DiagonalType diagonal = DiagonalType::right,
                         int gdim = 2, GhostMode ghost_mode = GhostMode::none,
                         const graph::Reorder& reorder_fn = graph::Reorder{})
{
  if (gdim < 2 || gdim > 3)
    throw std::invalid_argument("2 <= gdim <= 3 for rectangle mesh.");
  if (std::ranges::any_of(n, [](auto e) { return e < 1; }))
    throw std::invalid_argument("At least one cell per dimension is required.");
 
  for (int32_t i = 0; i < 2; i++)
  {
    if (p[0][i] >= p[1][i])
      throw std::invalid_argument("It must hold p[0] < p[1].");
  }
 
  if (std::abs(p[1][0] - p[0][0]) < std::numeric_limits<T>::epsilon()
      or std::abs(p[1][1] - p[0][1]) < std::numeric_limits<T>::epsilon())
  {
    throw std::invalid_argument(
        "Rectangle seems to have zero width, height or depth. Check "
        "dimensions");
  }
 
  if (!graph::has_partitioner(partitioner) and dolfinx::MPI::size(comm) > 1)
    partitioner = graph::partition_graph;
 
  switch (celltype)
  {
  case CellType::triangle:
    return impl::build_tri<T>(comm, p, n, partitioner, diagonal, ghost_mode,
                              reorder_fn, gdim);
  case CellType::quadrilateral:
    return impl::build_quad<T>(comm, p, n, partitioner, ghost_mode, reorder_fn,
                               gdim);
  default:
    throw std::invalid_argument("Generate rectangle mesh. Wrong cell type.");
  }
}

template <std::floating_point T = double>
Mesh<T> create_rectangle(MPI_Comm comm, std::array<std::array<T, 2>, 2> p,
                         std::array<std::int64_t, 2> n, CellType celltype,
                         DiagonalType diagonal = DiagonalType::right,
                         int gdim = 2)
{
  return create_rectangle<T>(comm, p, n, celltype,
                             graph::AnyPartitionFunction(), diagonal, gdim);
}

template <std::floating_point T = double>
Mesh<T> create_interval(MPI_Comm comm, std::int64_t n, std::array<T, 2> p,
                        mesh::GhostMode ghost_mode = mesh::GhostMode::none,
                        graph::AnyPartitionFunction partitioner = {},
                        int gdim = 1,
                        const graph::Reorder& reorder_fn = graph::Reorder{})
{
  if (gdim < 1 || gdim > 3)
    throw std::invalid_argument("1 <= gdim <= 3 for interval mesh.");
  if (n < 1)
    throw std::invalid_argument("At least one cell per dimension is required.");
 
  const auto [a, b] = p;
  if (a >= b)
    throw std::invalid_argument("It must hold p[0] < p[1].");
  if (std::abs(a - b) < std::numeric_limits<T>::epsilon())
  {
    throw std::invalid_argument(
        "Length of interval is zero. Check your dimensions.");
  }
 
  if (!graph::has_partitioner(partitioner) and dolfinx::MPI::size(comm) > 1)
    partitioner = graph::partition_graph;
 
  fem::CoordinateElement<T> element(CellType::interval, 1);
  if (dolfinx::MPI::rank(comm) == 0)
  {
    auto [x1d, cells] = impl::create_interval_cells<T>(p, n);
    std::size_t npts = x1d.size();
    if (gdim == 1)
    {
      return impl::finalize_mesh(comm, MPI_COMM_SELF, cells, element, x1d,
                                 {npts, 1}, partitioner, ghost_mode,
                                 reorder_fn);
    }
    std::vector<T> x(npts * gdim, T(0));
    for (std::size_t i = 0; i < npts; i++)
      x[i * gdim] = x1d[i];
    return impl::finalize_mesh(comm, MPI_COMM_SELF, cells, element, x,
                               {npts, static_cast<std::size_t>(gdim)},
                               partitioner, ghost_mode, reorder_fn);
  }
  else
  {
    return impl::finalize_mesh(comm, MPI_COMM_NULL, {}, element,
                               std::vector<T>{},
                               {0, static_cast<std::size_t>(gdim)}, partitioner,
                               ghost_mode, reorder_fn);
  }
}
 
namespace impl
{
 
template <std::floating_point T>
std::tuple<std::vector<T>, std::vector<std::int64_t>>
create_interval_cells(std::array<T, 2> p, std::int64_t n)
{
  const auto [a, b] = p;
 
  const T h = (b - a) / static_cast<T>(n);
 
  // Create vertices
  std::vector<T> x(n + 1);
  std::ranges::generate(x, [i = std::int64_t(0), a, h]() mutable
                        { return a + h * static_cast<T>(i++); });
 
  // Create intervals -> cells=[0, 1, 1, ..., n-1, n-1, n]
  std::vector<std::int64_t> cells(2 * n);
  for (std::size_t ix = 0; ix < cells.size() / 2; ++ix)
  {
    cells[2 * ix] = ix;
    cells[2 * ix + 1] = ix + 1;
  }
 
  return {std::move(x), std::move(cells)};
}
 
template <std::floating_point T>
std::vector<T> create_geom(MPI_Comm comm, std::array<std::array<T, 3>, 2> p,
                           std::array<std::int64_t, 3> n)
{
  // Extract data
  auto [p0, p1] = p;
  const auto [nx, ny, nz] = n;
 
  assert(std::ranges::all_of(n, [](auto e) { return e >= 1; }));
  assert(p0 < p1);
 
  // Structured grid cuboid extents
  const std::array<T, 3> extents = {
      (p1[0] - p0[0]) / static_cast<T>(nx),
      (p1[1] - p0[1]) / static_cast<T>(ny),
      (p1[2] - p0[2]) / static_cast<T>(nz),
  };
 
  const std::int64_t n_points = (nx + 1) * (ny + 1) * (nz + 1);
  const auto [range_begin, range_end] = common::local_range(
      dolfinx::MPI::rank(comm), n_points, dolfinx::MPI::size(comm));
 
  std::vector<T> geom;
  geom.reserve((range_end - range_begin) * 3);
  const std::int64_t sqxy = (nx + 1) * (ny + 1);
  for (std::int64_t v = range_begin; v < range_end; ++v)
  {
    // lexiographic index to spatial index
    const std::int64_t xy_idx = v % sqxy;
    std::array<std::int64_t, 3> idx
        = {xy_idx % (nx + 1), xy_idx / (nx + 1), v / sqxy};
 
    // vertex = p0 + idx * extents (elementwise)
    for (std::size_t i = 0; i < idx.size(); i++)
      geom.push_back(p0[i] + static_cast<T>(idx[i]) * extents[i]);
  }
 
  return geom;
}
 
template <std::floating_point T>
Mesh<T> build_tet(MPI_Comm comm, MPI_Comm subcomm,
                  std::array<std::array<T, 3>, 2> p,
                  std::array<std::int64_t, 3> n,
                  const graph::AnyPartitionFunction& partitioner,
                  GhostMode ghost_mode, const graph::Reorder& reorder_fn)
{
  common::Timer timer("Build BoxMesh (tetrahedra)");
  std::vector<T> x;
  std::vector<std::int64_t> cells;
  fem::CoordinateElement<T> element(CellType::tetrahedron, 1);
  if (subcomm != MPI_COMM_NULL)
  {
    x = create_geom<T>(subcomm, p, n);
 
    const auto [nx, ny, nz] = n;
    const std::int64_t n_cells = nx * ny * nz;
 
    std::array range_c = common::local_range(
        dolfinx::MPI::rank(subcomm), n_cells, dolfinx::MPI::size(subcomm));
    cells.reserve(6 * (range_c[1] - range_c[0]) * 4);
 
    // Create tetrahedra
    for (std::int64_t i = range_c[0]; i < range_c[1]; ++i)
    {
      const std::int64_t iz = i / (nx * ny);
      const std::int64_t j = i % (nx * ny);
      const std::int64_t iy = j / nx;
      const std::int64_t ix = j % nx;
      const std::int64_t v0 = iz * (nx + 1) * (ny + 1) + iy * (nx + 1) + ix;
      const std::int64_t v1 = v0 + 1;
      const std::int64_t v2 = v0 + (nx + 1);
      const std::int64_t v3 = v1 + (nx + 1);
      const std::int64_t v4 = v0 + (nx + 1) * (ny + 1);
      const std::int64_t v5 = v1 + (nx + 1) * (ny + 1);
      const std::int64_t v6 = v2 + (nx + 1) * (ny + 1);
      const std::int64_t v7 = v3 + (nx + 1) * (ny + 1);
 
      // Kuhn decomposition of the hexahedron into 6 tetrahedra sharing
      // the v0-v7 diagonal.
      cells.insert(cells.end(),
                   {v0, v1, v3, v7, v0, v1, v7, v5, v0, v5, v7, v4,
                    v0, v3, v2, v7, v0, v6, v4, v7, v0, v2, v6, v7});
    }
  }
 
  return finalize_mesh(comm, subcomm, cells, element, x, {x.size() / 3, 3},
                       partitioner, ghost_mode, reorder_fn);
}
 
template <std::floating_point T>
mesh::Mesh<T> build_hex(MPI_Comm comm, MPI_Comm subcomm,
                        std::array<std::array<T, 3>, 2> p,
                        std::array<std::int64_t, 3> n,
                        const graph::AnyPartitionFunction& partitioner,
                        GhostMode ghost_mode, const graph::Reorder& reorder_fn)
{
  common::Timer timer("Build BoxMesh (hexahedra)");
  std::vector<T> x;
  std::vector<std::int64_t> cells;
  fem::CoordinateElement<T> element(CellType::hexahedron, 1);
  if (subcomm != MPI_COMM_NULL)
  {
    x = create_geom<T>(subcomm, p, n);
 
    // Create cuboids
    const auto [nx, ny, nz] = n;
    const std::int64_t n_cells = nx * ny * nz;
    std::array range_c = common::local_range(
        dolfinx::MPI::rank(subcomm), n_cells, dolfinx::MPI::size(subcomm));
    cells.reserve((range_c[1] - range_c[0]) * 8);
    for (std::int64_t i = range_c[0]; i < range_c[1]; ++i)
    {
      const std::int64_t iz = i / (nx * ny);
      const std::int64_t j = i % (nx * ny);
      const std::int64_t iy = j / nx;
      const std::int64_t ix = j % nx;
 
      const std::int64_t v0 = (iz * (ny + 1) + iy) * (nx + 1) + ix;
      const std::int64_t v1 = v0 + 1;
      const std::int64_t v2 = v0 + (nx + 1);
      const std::int64_t v3 = v1 + (nx + 1);
      const std::int64_t v4 = v0 + (nx + 1) * (ny + 1);
      const std::int64_t v5 = v1 + (nx + 1) * (ny + 1);
      const std::int64_t v6 = v2 + (nx + 1) * (ny + 1);
      const std::int64_t v7 = v3 + (nx + 1) * (ny + 1);
      cells.insert(cells.end(), {v0, v1, v2, v3, v4, v5, v6, v7});
    }
  }
 
  return finalize_mesh(comm, subcomm, cells, element, x, {x.size() / 3, 3},
                       partitioner, ghost_mode, reorder_fn);
}
 
template <std::floating_point T>
Mesh<T> build_prism(MPI_Comm comm, MPI_Comm subcomm,
                    std::array<std::array<T, 3>, 2> p,
                    std::array<std::int64_t, 3> n,
                    const graph::AnyPartitionFunction& partitioner,
                    GhostMode ghost_mode, const graph::Reorder& reorder_fn)
{
  std::vector<T> x;
  std::vector<std::int64_t> cells;
  fem::CoordinateElement<T> element(CellType::prism, 1);
  if (subcomm != MPI_COMM_NULL)
  {
    x = create_geom<T>(subcomm, p, n);
 
    const std::int64_t nx = n[0];
    const std::int64_t ny = n[1];
    const std::int64_t nz = n[2];
    const std::int64_t n_cells = nx * ny * nz;
    std::array range_c = common::local_range(
        dolfinx::MPI::rank(subcomm), n_cells, dolfinx::MPI::size(subcomm));
    const std::int64_t cell_range = range_c[1] - range_c[0];
 
    // Create cuboids
    cells.reserve(2 * cell_range * 6);
    for (std::int64_t i = range_c[0]; i < range_c[1]; ++i)
    {
      const std::int64_t iz = i / (nx * ny);
      const std::int64_t j = i % (nx * ny);
      const std::int64_t iy = j / nx;
      const std::int64_t ix = j % nx;
 
      const std::int64_t v0 = (iz * (ny + 1) + iy) * (nx + 1) + ix;
      const std::int64_t v1 = v0 + 1;
      const std::int64_t v2 = v0 + (nx + 1);
      const std::int64_t v3 = v1 + (nx + 1);
      const std::int64_t v4 = v0 + (nx + 1) * (ny + 1);
      const std::int64_t v5 = v1 + (nx + 1) * (ny + 1);
      const std::int64_t v6 = v2 + (nx + 1) * (ny + 1);
      const std::int64_t v7 = v3 + (nx + 1) * (ny + 1);
      cells.insert(cells.end(), {v0, v1, v2, v4, v5, v6});
      cells.insert(cells.end(), {v1, v2, v3, v5, v6, v7});
    }
  }
 
  return finalize_mesh(comm, subcomm, cells, element, x, {x.size() / 3, 3},
                       partitioner, ghost_mode, reorder_fn);
}
 
template <std::floating_point T>
Mesh<T> build_tri(MPI_Comm comm, std::array<std::array<T, 2>, 2> p,
                  std::array<std::int64_t, 2> n,
                  const graph::AnyPartitionFunction& partitioner,
                  DiagonalType diagonal, GhostMode ghost_mode,
                  const graph::Reorder& reorder_fn, int gdim)
{
  fem::CoordinateElement<T> element(CellType::triangle, 1);
  if (gdim < 2 || gdim > 3)
    throw std::invalid_argument("2 <= gdim <= 3 for tri mesh.");
 
  if (dolfinx::MPI::rank(comm) == 0)
  {
    const auto [p0, p1] = p;
    const auto [nx, ny] = n;
 
    const auto [a, c] = p0;
    const auto [b, d] = p1;
 
    const T ab = (b - a) / static_cast<T>(nx);
    const T cd = (d - c) / static_cast<T>(ny);
 
    // Create vertices and cells
    std::int64_t nv, nc;
    switch (diagonal)
    {
    case DiagonalType::crossed:
      nv = (nx + 1) * (ny + 1) + nx * ny;
      nc = 4 * nx * ny;
      break;
    default:
      nv = (nx + 1) * (ny + 1);
      nc = 2 * nx * ny;
    }
 
    std::vector<T> x;
    x.reserve(nv * 2);
    std::vector<std::int64_t> cells;
    cells.reserve(nc * 3);
 
    // Create main vertices
    for (std::int64_t iy = 0; iy <= ny; iy++)
    {
      T x1 = c + cd * static_cast<T>(iy);
      for (std::int64_t ix = 0; ix <= nx; ix++)
        x.insert(x.end(), {a + ab * static_cast<T>(ix), x1});
    }
 
    // Create midpoint vertices if the mesh type is crossed
    switch (diagonal)
    {
    case DiagonalType::crossed:
      for (std::int64_t iy = 0; iy < ny; iy++)
      {
        T x1 = c + cd * (static_cast<T>(iy) + 0.5);
        for (std::int64_t ix = 0; ix < nx; ix++)
        {
          T x0 = a + ab * (static_cast<T>(ix) + 0.5);
          x.insert(x.end(), {x0, x1});
        }
      }
      break;
    default:
      break;
    }
 
    // Create triangles
    switch (diagonal)
    {
    case DiagonalType::crossed:
    {
      for (std::int64_t iy = 0; iy < ny; iy++)
      {
        for (std::int64_t ix = 0; ix < nx; ix++)
        {
          std::int64_t v0 = iy * (nx + 1) + ix;
          std::int64_t v1 = v0 + 1;
          std::int64_t v2 = v0 + (nx + 1);
          std::int64_t v3 = v1 + (nx + 1);
          std::int64_t vmid = (nx + 1) * (ny + 1) + iy * nx + ix;
 
          // Note that v0 < v1 < v2 < v3 < vmid
          cells.insert(cells.end(), {v0, v1, vmid, v0, v2, vmid, v1, v3, vmid,
                                     v2, v3, vmid});
        }
      }
      break;
    }
    default:
    {
      DiagonalType local_diagonal = diagonal;
      for (std::int64_t iy = 0; iy < ny; iy++)
      {
        // Set up alternating diagonal
        switch (diagonal)
        {
        case DiagonalType::right_left:
          if (iy % 2)
            local_diagonal = DiagonalType::right;
          else
            local_diagonal = DiagonalType::left;
          break;
        case DiagonalType::left_right:
          if (iy % 2)
            local_diagonal = DiagonalType::left;
          else
            local_diagonal = DiagonalType::right;
          break;
        default:
          break;
        }
        for (std::int64_t ix = 0; ix < nx; ix++)
        {
          std::int64_t v0 = iy * (nx + 1) + ix;
          std::int64_t v1 = v0 + 1;
          std::int64_t v2 = v0 + (nx + 1);
          std::int64_t v3 = v1 + (nx + 1);
          switch (local_diagonal)
          {
          case DiagonalType::left:
          {
            cells.insert(cells.end(), {v0, v1, v2, v1, v2, v3});
            if (diagonal == DiagonalType::right_left
                or diagonal == DiagonalType::left_right)
            {
              local_diagonal = DiagonalType::right;
            }
            break;
          }
          default:
          {
            cells.insert(cells.end(), {v0, v1, v3, v0, v2, v3});
            if (diagonal == DiagonalType::right_left
                or diagonal == DiagonalType::left_right)
            {
              local_diagonal = DiagonalType::left;
            }
          }
          }
        }
      }
    }
    }
 
    std::size_t npts = x.size() / 2;
    if (gdim == 2)
    {
      return finalize_mesh(comm, MPI_COMM_SELF, cells, element, x, {npts, 2},
                           partitioner, ghost_mode, reorder_fn);
    }
    std::vector<T> xg(npts * gdim, T(0));
    for (std::size_t i = 0; i < npts; i++)
    {
      xg[i * gdim] = x[2 * i];
      xg[i * gdim + 1] = x[2 * i + 1];
    }
    return finalize_mesh(comm, MPI_COMM_SELF, cells, element, xg,
                         {npts, static_cast<std::size_t>(gdim)}, partitioner,
                         ghost_mode, reorder_fn);
  }
  else
  {
    return finalize_mesh(comm, MPI_COMM_NULL, {}, element, std::vector<T>{},
                         {0, static_cast<std::size_t>(gdim)}, partitioner,
                         ghost_mode, reorder_fn);
  }
}
 
template <std::floating_point T>
Mesh<T> build_quad(MPI_Comm comm, std::array<std::array<T, 2>, 2> p,
                   std::array<std::int64_t, 2> n,
                   const graph::AnyPartitionFunction& partitioner,
                   GhostMode ghost_mode, const graph::Reorder& reorder_fn,
                   int gdim)
{
  if (gdim < 2 || gdim > 3)
    throw std::invalid_argument("2 <= gdim <= 3 for quad mesh.");
 
  fem::CoordinateElement<T> element(CellType::quadrilateral, 1);
  if (dolfinx::MPI::rank(comm) == 0)
  {
    const auto [nx, ny] = n;
    const auto [a, c] = p[0];
    const auto [b, d] = p[1];
 
    const T ab = (b - a) / static_cast<T>(nx);
    const T cd = (d - c) / static_cast<T>(ny);
 
    // Create vertices
    std::vector<T> x;
    x.reserve((nx + 1) * (ny + 1) * 2);
    for (std::int64_t ix = 0; ix <= nx; ix++)
    {
      T x0 = a + ab * static_cast<T>(ix);
      for (std::int64_t iy = 0; iy <= ny; iy++)
        x.insert(x.end(), {x0, c + cd * static_cast<T>(iy)});
    }
 
    // Create rectangles
    std::vector<std::int64_t> cells;
    cells.reserve(nx * ny * 4);
    for (std::int64_t ix = 0; ix < nx; ix++)
    {
      for (std::int64_t iy = 0; iy < ny; iy++)
      {
        std::int64_t i0 = ix * (ny + 1);
        cells.insert(cells.end(), {i0 + iy, i0 + iy + 1, i0 + iy + ny + 1,
                                   i0 + iy + ny + 2});
      }
    }
 
    std::size_t npts = x.size() / 2;
    if (gdim == 2)
    {
      return finalize_mesh(comm, MPI_COMM_SELF, cells, element, x, {npts, 2},
                           partitioner, ghost_mode, reorder_fn);
    }
    std::vector<T> xg(npts * gdim, T(0));
    for (std::size_t i = 0; i < npts; i++)
    {
      xg[i * gdim] = x[2 * i];
      xg[i * gdim + 1] = x[2 * i + 1];
    }
    return finalize_mesh(comm, MPI_COMM_SELF, cells, element, xg,
                         {npts, static_cast<std::size_t>(gdim)}, partitioner,
                         ghost_mode, reorder_fn);
  }
  else
  {
    return finalize_mesh(comm, MPI_COMM_NULL, {}, element, std::vector<T>{},
                         {0, static_cast<std::size_t>(gdim)}, partitioner,
                         ghost_mode, reorder_fn);
  }
}
} // namespace impl
} // namespace dolfinx::mesh