IndexMap.h

// Copyright (C) 2015-2024 Chris Richardson, Garth N. Wells and Igor Baratta
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "MPI.h"
#include <cstdint>
#include <memory>
#include <span>
#include <tuple>
#include <utility>
#include <vector>
 
namespace dolfinx::common
{
// Forward declaration
class IndexMap;

enum class IndexMapOrder : bool
{
  preserve = true, 
  any = false      
};

std::vector<int32_t>
compute_owned_indices(std::span<const std::int32_t> indices,
                      const IndexMap& map);

std::tuple<std::int64_t, std::vector<std::int32_t>,
           std::vector<std::vector<std::int64_t>>,
           std::vector<std::vector<int>>>
stack_index_maps(
    const std::vector<std::pair<std::reference_wrapper<const IndexMap>, int>>&
        maps);

std::pair<IndexMap, std::vector<std::int32_t>> create_sub_index_map(
    const IndexMap& imap, std::span<const std::int32_t> indices,
    IndexMapOrder order = IndexMapOrder::any, bool allow_owner_change = false);

class IndexMap
{
public:
  IndexMap(MPI_Comm comm, std::int32_t local_size);

  IndexMap(MPI_Comm comm, std::int32_t local_size,
           std::span<const std::int64_t> ghosts, std::span<const int> owners,
           int tag = static_cast<int>(dolfinx::MPI::tag::consensus_nbx));

  IndexMap(MPI_Comm comm, std::int32_t local_size,
           const std::array<std::vector<int>, 2>& src_dest,
           std::span<const std::int64_t> ghosts, std::span<const int> owners);
 
  // Copy constructor
  IndexMap(const IndexMap& map) = delete;

  IndexMap(IndexMap&& map) = default;

  ~IndexMap() = default;

  IndexMap& operator=(IndexMap&& map) = default;
 
  // Copy assignment
  IndexMap& operator=(const IndexMap& map) = delete;

  std::array<std::int64_t, 2> local_range() const noexcept;

  std::int32_t num_ghosts() const noexcept;

  std::int32_t size_local() const noexcept;

  std::int64_t size_global() const noexcept;

  std::span<const std::int64_t> ghosts() const noexcept;

  MPI_Comm comm() const;

  void local_to_global(std::span<const std::int32_t> local,
                       std::span<std::int64_t> global) const;

  void global_to_local(std::span<const std::int64_t> global,
                       std::span<std::int32_t> local) const;

  std::vector<std::int64_t> global_indices() const;

  std::span<const int> owners() const { return _owners; }

  std::pair<std::vector<int>, std::vector<std::int32_t>> index_to_dest_ranks(
      int tag = static_cast<int>(dolfinx::MPI::tag::consensus_nbx)) const;

  std::vector<std::int32_t> shared_indices() const;

  std::span<const int> src() const noexcept;

  std::span<const int> dest() const noexcept;

  std::vector<std::int32_t> weights_src() const;

  std::vector<std::int32_t> weights_dest() const;

  std::array<std::vector<int>, 2> rank_type(int split_type) const;
 
private:
  // Range of indices (global) owned by this process
  std::array<std::int64_t, 2> _local_range;
 
  // Number indices across communicator
  std::int64_t _size_global;
 
  // MPI communicator that map is defined on
  dolfinx::MPI::Comm _comm;
 
  // Local-to-global map for ghost indices
  std::vector<std::int64_t> _ghosts;
 
  // Owning rank on _comm for the ith ghost index
  std::vector<int> _owners;
 
  // Set of ranks that own ghosts
  std::vector<int> _src;
 
  // Set of ranks ghost owned indices
  std::vector<int> _dest;
};
 
} // namespace dolfinx::common