partition.h

// Copyright (C) 2020-2026 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "AdjacencyList.h"
#include <array>
#include <cstdint>
#include <functional>
#include <mpi.h>
#include <optional>
#include <span>
#include <tuple>
#include <variant>
#include <vector>
 
namespace dolfinx::graph
{
using partition_fn = std::function<graph::AdjacencyList<std::int32_t>(
    MPI_Comm, int, const AdjacencyList<std::int64_t>&,
    std::optional<std::span<const std::int32_t>>,
    std::optional<std::span<const std::int32_t>>, bool)>;

using geom_partition_fn = std::function<std::vector<int>(
    MPI_Comm, int, std::span<const double>, int,
    std::optional<std::span<const std::int32_t>>)>;

using hybrid_partition_fn = std::function<graph::AdjacencyList<std::int32_t>(
    MPI_Comm, int, const AdjacencyList<std::int64_t>&, std::span<const double>,
    std::optional<std::span<const std::int32_t>>,
    std::optional<std::span<const std::int32_t>>, bool)>;

using AnyPartitionFunction
    = std::variant<partition_fn, geom_partition_fn, hybrid_partition_fn>;

bool has_partitioner(const AnyPartitionFunction& partitioner);

AdjacencyList<std::int32_t> partition_graph(
    MPI_Comm comm, int nparts, const AdjacencyList<std::int64_t>& local_graph,
    std::optional<std::span<const std::int32_t>> node_weights,
    std::optional<std::span<const std::int32_t>> edge_weights, bool ghosting);

/// for.
struct Partitioner
{
  AnyPartitionFunction fn = partition_fn(partition_graph);

  std::optional<std::span<const std::int32_t>> node_weights = std::nullopt;
};

/// @todo Add a function that sends data to the 'owner'
namespace build
{
std::tuple<graph::AdjacencyList<std::int64_t>, std::vector<int>,
           std::vector<std::int64_t>, std::vector<int>>
distribute(MPI_Comm comm, const graph::AdjacencyList<std::int64_t>& list,
           const graph::AdjacencyList<std::int32_t>& destinations);

std::tuple<std::vector<std::int64_t>, std::vector<int>,
           std::vector<std::int64_t>, std::vector<int>>
distribute(MPI_Comm comm, std::span<const std::int64_t> list,
           std::array<std::size_t, 2> shape,
           const graph::AdjacencyList<std::int32_t>& destinations);

std::vector<std::int64_t>
compute_ghost_indices(MPI_Comm comm,
                      std::span<const std::int64_t> owned_indices,
                      std::span<const std::int64_t> ghost_indices,
                      std::span<const int> ghost_owners, int num_threads);

std::vector<std::int64_t>
compute_local_to_global(std::span<const std::int64_t> global,
                        std::span<const std::int32_t> local);

std::vector<std::int32_t>
compute_local_to_local(std::span<const std::int64_t> local0_to_global,
                       std::span<const std::int64_t> local1_to_global);
} // namespace build
 
} // namespace dolfinx::graph