mark.h

// Copyright (C) 2026 Paul T. Kühner and Jack S. Hale
//
// This file is part of DOLFINX (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <algorithm>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/MPI.h>
#include <dolfinx/la/Vector.h>
#include <format>
#include <limits>
#include <mpi.h>
#include <numeric>
#include <span>
#include <spdlog/spdlog.h>
#include <stdexcept>
#include <type_traits>
#include <vector>
 
namespace dolfinx::refinement
{

template <std::floating_point T>
std::vector<std::int32_t> mark_maximum(std::span<const T> values,
                                       const common::IndexMap& index_map,
                                       std::type_identity_t<T> theta)
{
  if ((theta <= 0) or (theta > 1))
  {
    throw std::invalid_argument(
        std::format("theta must satisfy 0 < theta <= 1, got {}.", theta));
  }
 
  std::int32_t n = values.size();
  std::int32_t size = index_map.size_local() + index_map.num_ghosts();
  if (n != size)
  {
    throw std::invalid_argument(
        std::format("values must have size index_map.size_local() + "
                    "index_map.num_ghosts() = {}, got {}.",
                    size, n));
  }
 
  T local_max = index_map.size_local() == 0
                    ? std::numeric_limits<T>::lowest()
                    : std::ranges::max(values.first(index_map.size_local()));
 
  T max = 0;
  MPI_Allreduce(&local_max, &max, 1, dolfinx::MPI::mpi_t<T>, MPI_MAX,
                index_map.comm());
 
  T threshold = theta * max;
 
  auto mark = [threshold](T e) { return e > threshold; };
 
  std::vector<std::int32_t> indices;
  indices.reserve(std::ranges::count_if(values, mark));
  for (std::int32_t i = 0; i < n; ++i)
  {
    if (mark(values[i]))
      indices.push_back(i);
  }
 
  spdlog::info("Marking (maximum): marked {} of {} local entries (owned + "
               "ghost).",
               indices.size(), n);
 
  return indices;
}

template <std::floating_point T>
std::vector<std::int32_t>
mark_equidistribution(std::span<const T> values,
                      const common::IndexMap& index_map,
                      std::type_identity_t<T> theta)
{
  if ((theta <= 0) or (theta > 1))
  {
    throw std::invalid_argument(
        std::format("theta must satisfy 0 < theta <= 1, got {}.", theta));
  }
 
  std::int32_t n = values.size();
  std::int32_t size = index_map.size_local() + index_map.num_ghosts();
  if (n != size)
  {
    throw std::invalid_argument(
        std::format("values must have size index_map.size_local() + "
                    "index_map.num_ghosts() = {}, got {}.",
                    size, n));
  }
 
  auto owned = values.first(index_map.size_local());
  T local_squared_norm = std::accumulate(owned.begin(), owned.end(), T{0});
 
  T squared_norm;
  MPI_Allreduce(&local_squared_norm, &squared_norm, 1, dolfinx::MPI::mpi_t<T>,
                MPI_SUM, index_map.comm());
 
  T threshold
      = theta * theta * squared_norm / static_cast<T>(index_map.size_global());
 
  auto mark = [threshold](T e) { return e > threshold; };
 
  std::vector<std::int32_t> indices;
  indices.reserve(std::ranges::count_if(values, mark));
  for (std::int32_t i = 0; i < n; ++i)
  {
    if (mark(values[i]))
      indices.push_back(i);
  }
 
  spdlog::info(
      "Marking (equidistribution): marked {} of {} local entries (owned + "
      "ghost).",
      indices.size(), n);
 
  return indices;
}
 
} // namespace dolfinx::refinement