local_range.h

 
// Copyright (C) 2007-2023 Magnus Vikstrøm, Garth N. Wells and Paul T. Kühner
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <cstdint>
#include <stdexcept>
 
namespace dolfinx::common
{
constexpr std::array<std::int64_t, 2> local_range(int index, std::int64_t N,
                                                  int size)
{
  if (index < 0)
    throw std::invalid_argument("local_range: index must be non-negative");
  if (N < 0)
    throw std::invalid_argument("local_range: N must be non-negative");
  if (size <= 0)
    throw std::invalid_argument("local_range: size must be positive");
 
  // Compute number of items per rank and remainder
  const std::int64_t n = N / size;
  const std::int64_t r = N % size;
 
  // Compute local range
  if (index < r)
    return {index * (n + 1), index * (n + 1) + n + 1};
  else
    return {index * n + r, index * n + r + n};
}
} // namespace dolfinx::common