InPlaceClusterization.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
#ifndef ACTS_INPLACECLUSTERIZATION_H
#define ACTS_INPLACECLUSTERIZATION_H
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <concepts>
#include <limits>
#include <ranges>
#include <span>
#include <type_traits>
 
namespace Acts::InPlaceClusterization {
 
// For cells that should be clusterized the following functions
// need to be implemented:
namespace traits {
// @brief Get the number of dimensions of the grid the cell lives on
template <typename cell_t>
constexpr auto getCellDimension() {
  return std::tuple_size<decltype(cell_t::coordinates)>();
}

template <typename cell_t>
auto getCellCoordinate(const cell_t &a, unsigned int axis_i) {
  return a.coordinates[axis_i];
}

template <typename cell_t, std::unsigned_integral index_t>
void setLabel(cell_t &a, index_t label) {
  a.label = label;
}

template <typename cell_t>
auto getLabel(const cell_t &a) {
  return a.label;
}
}  // namespace traits

template <typename coordinates_t, std::size_t N, std::unsigned_integral index_t>
struct Cell {
  Cell(const std::array<coordinates_t, N> &the_coordinates, index_t src_index)
      : coordinates(the_coordinates), label(index_t{}), srcIndex(src_index) {}
 
  std::array<coordinates_t, N>
      coordinates;   
  index_t label;     
  index_t srcIndex;  
};

template <typename cell_t>
concept CellWithLabel = requires(const cell_t &a, const cell_t &b,
                                 decltype(traits::getLabel(a)) idx,
                                 unsigned int axis_i) {
  { traits::getLabel(a) } -> std::unsigned_integral;
  traits::setLabel(a, idx);
  {
    traits::getCellCoordinate(a, axis_i) < traits::getCellCoordinate(b, axis_i)
  } -> std::convertible_to<bool>;
  { traits::getCellDimension<cell_t>() } -> std::convertible_to<std::size_t>;
};

template <typename container_t>
concept SequenceContainer = requires(container_t &cont, unsigned int index,
                                     typename container_t::value_type) {
  {
    cont[index]
  } -> std::convertible_to<const typename container_t::value_type &>;
  { cont.empty() } -> std::convertible_to<bool>;
  { cont.size() } -> std::convertible_to<std::size_t>;
} && std::permutable<typename container_t::iterator>;

template <typename cell_collection_t>
concept CellCollection = SequenceContainer<cell_collection_t> &&
                         CellWithLabel<typename cell_collection_t::value_type>;

template <typename coordinates_t>
bool isConnectedCommonEdgeOrCorner(const coordinates_t &coordinates_diff) {
  bool connected = true;
  for (const auto &a_coordinate_diff : coordinates_diff) {
    connected &= a_coordinate_diff <= 1;
  }
  return connected;
}

template <typename coordinates_t>
bool isConnectedCommonEdge(const coordinates_t &coordinates_diff) {
  int connections = 0;
  for (const auto &a_coordinate_diff : coordinates_diff) {
    connections += a_coordinate_diff;
  }
  return connections <= 1;
}
 
enum class EConnectionType { CommonEdgeOrCorner, CommonEdge };

template <CellWithLabel cell_t,
          EConnectionType connection_type = EConnectionType::CommonEdgeOrCorner>
struct ConnectionHelper {
  static constexpr std::size_t NDim = traits::getCellDimension<cell_t>();
  using coordinate_t = std::remove_cvref_t<decltype(traits::getCellCoordinate(
      std::declval<cell_t>(), 0u))>;

  static bool isConnected(
      const std::array<coordinate_t, NDim> &coordinates_diff) {
    if constexpr (connection_type == EConnectionType::CommonEdgeOrCorner) {
      return isConnectedCommonEdgeOrCorner(coordinates_diff);
    } else {
      return isConnectedCommonEdge(coordinates_diff);
    }
  }

  static bool canAbortSearch(
      const std::array<coordinate_t, NDim> &coordinates_diff,
      unsigned int sort_axis_i) {
    assert(sort_axis_i < coordinates_diff.size());
    return coordinates_diff[sort_axis_i] > 1;
  }
};
 
template <
    EConnectionType connection_type = EConnectionType::CommonEdgeOrCorner,
    CellCollection cell_container_t = std::span<Cell<int, 2, unsigned int> > >
auto defaultConnectionHelper([[maybe_unused]] const cell_container_t &cells) {
  return ConnectionHelper<typename cell_container_t::value_type,
                          connection_type>{};
}

template <typename coordinate_t>
auto absDifference(coordinate_t a, coordinate_t b) {
  if constexpr (std::is_signed_v<coordinate_t>) {
    // @TODO check overflow with C++26, or change return type to
    //       an unsigned integer of same size.
    return static_cast<coordinate_t>(std::abs(a - b));
  } else {
    return static_cast<coordinate_t>((a > b ? a - b : b - a));
  }
}

template <unsigned int SORT_AXIS, CellCollection cell_collection_t,
          std::unsigned_integral index_t = unsigned int,
          typename connection_helper_t>
void labelSortedCells(cell_collection_t &cells,
                      connection_helper_t &&connection_helper) {
  using cell_t = typename cell_collection_t::value_type;
  static constexpr std::size_t NDim = traits::getCellDimension<cell_t>();
  static_assert(NDim > 0);
  using label_t = std::remove_cvref_t<decltype(traits::getLabel(cells[0]))>;
  using coordinate_t =
      std::remove_cvref_t<decltype(traits::getCellCoordinate(cells[0], 0u))>;
  static_assert(std::numeric_limits<index_t>::max() >=
                std::numeric_limits<label_t>::max());
  // cells are sorted in the first coordinate
  // thus can stop searching for adjacent cells
  // if distance in the sorted coordinate is too large
  for (index_t idx_a = 0; idx_a < cells.size(); ++idx_a) {
    traits::setLabel(cells[idx_a], idx_a);
    for (index_t idx_b = idx_a; idx_b-- > 0;) {
      // Unnecessary default initialization to satisfy clang-tidy:
      std::array<coordinate_t, NDim> diff{};
      for (unsigned int axis_i = 0; axis_i < NDim; ++axis_i) {
        diff[axis_i] =
            absDifference(traits::getCellCoordinate(cells[idx_a], axis_i),
                          traits::getCellCoordinate(cells[idx_b], axis_i));
      }
 
      if (connection_helper.isConnected(diff)) {
        if (traits::getLabel(cells[idx_a]) < idx_a) {
          // Unnecessary default initialization to satisfy clang-tidy:
          label_t min_label{};
          label_t max_label{};
          if (traits::getLabel(cells[idx_a]) < traits::getLabel(cells[idx_b])) {
            min_label = traits::getLabel(cells[idx_a]);
            max_label = traits::getLabel(cells[idx_b]);
            traits::setLabel(cells[idx_b], min_label);
          } else {
            max_label = traits::getLabel(cells[idx_a]);
            min_label = traits::getLabel(cells[idx_b]);
            traits::setLabel(cells[idx_a], min_label);
          }
          // nothing will be done if the min and the max label are identical
          if (min_label != max_label) {
            // can only encounter cells with label max_label down to index
            // max_label
            for (index_t idx = idx_a; idx-- > max_label;) {
              if (traits::getLabel(cells[idx]) == max_label) {
                traits::setLabel(cells[idx], min_label);
              }
            }
          }
        } else {
          traits::setLabel(cells[idx_a], traits::getLabel(cells[idx_b]));
        }
      } else if (connection_helper.canAbortSearch(diff, SORT_AXIS)) {
        // difference too large in sorted coordinate, there won't be any more
        // candidates for merging. Can abort search for cell idx_a
        break;
      }
    }
  }
}

template <unsigned int AXIS, CellCollection cell_collection_t,
          std::size_t NDim = 2, std::unsigned_integral index_t = unsigned int>
  requires(NDim > 0 && AXIS < NDim)
void sortCellsByCoordinate(cell_collection_t &cells) {
  using cell_t = typename cell_collection_t::value_type;
  std::sort(cells.begin(), cells.end(), [](const cell_t &a, const cell_t &b) {
    return traits::getCellCoordinate(a, AXIS) <
           traits::getCellCoordinate(b, AXIS);
  });
}

template <CellCollection cell_collection_t>
void groupCellsByLabel(cell_collection_t &cells) {
  using cell_t = typename cell_collection_t::value_type;
  // stable sort to retain coordinate ordering
  std::stable_sort(cells.begin(), cells.end(),
                   [](const cell_t &a, const cell_t &b) {
                     return traits::getLabel(a) < traits::getLabel(b);
                   });
}

template <
    unsigned int SORT_AXIS, std::unsigned_integral index_t = unsigned int,
    CellCollection cell_collection_t = std::span<Cell<int, 2, unsigned int> >,
    typename connection_helper_t =
        ConnectionHelper<Cell<int, 2, unsigned int> > >
void clusterize(cell_collection_t &cells,
                connection_helper_t &&connection_helper =
                    ConnectionHelper<typename cell_collection_t::value_type,
                                     EConnectionType::CommonEdgeOrCorner>{}) {
  assert(cells.size() <= std::numeric_limits<index_t>::max());
  sortCellsByCoordinate<SORT_AXIS>(cells);
  labelSortedCells<SORT_AXIS, cell_collection_t, index_t>(cells,
                                                          connection_helper);
  groupCellsByLabel(cells);
}

template <CellCollection cell_collection_t,
          std::unsigned_integral index_t = unsigned int>
std::size_t countLabels(const cell_collection_t &cells) {
  assert(cells.size() <= std::numeric_limits<index_t>::max());
  std::size_t nlabels = cells.empty() ? 0 : 1;
  for (index_t idx_a = 1; idx_a < cells.size(); ++idx_a) {
    nlabels +=
        traits::getLabel(cells[idx_a]) != traits::getLabel(cells[idx_a - 1]);
  }
  return nlabels;
}

template <CellCollection cell_collection_t, typename func_t>
void for_each_cluster(cell_collection_t &cells, func_t func) {
  if (cells.empty()) {
    return;
  }
  using index_t = decltype(traits::getLabel(cells[0]));
  index_t idx_begin = 0;
  index_t idx = 1;
  for (; idx < cells.size(); ++idx) {
    if (traits::getLabel(cells[idx]) != traits::getLabel(cells[idx - 1])) {
      func(cells, idx_begin, idx);
      idx_begin = idx;
    }
  }
  if (idx_begin < idx) {
    func(cells, idx_begin, idx);
  }
}

template <CellCollection cell_collection_t, typename range_collection_t>
void addCellRanges(const cell_collection_t &cells, range_collection_t &ranges) {
  for_each_cluster(
      cells, [&ranges]([[maybe_unused]] const cell_collection_t &all_cells,
                       unsigned int idx_begin, unsigned int idx_end) {
        ranges.emplace_back(idx_begin, idx_end);
      });
}

template <unsigned int AXIS, CellCollection cell_collection_t,
          typename range_collection_t>
void addCellRangesAndSort(cell_collection_t &cells,
                          range_collection_t &ranges) {
  for_each_cluster(cells, [&ranges](cell_collection_t &all_cells,
                                    unsigned int idx_begin,
                                    unsigned int idx_end) {
    auto cluster_range =
        std::span(all_cells.begin() + idx_begin, all_cells.begin() + idx_end);
    using cell_t = typename cell_collection_t::value_type;
    std::ranges::sort(cluster_range, [](const cell_t &a, const cell_t &b) {
      return traits::getCellCoordinate(a, AXIS) <
             traits::getCellCoordinate(b, AXIS);
    });
    ranges.emplace_back(idx_begin, idx_end);
  });
}
 
}  // namespace Acts::InPlaceClusterization
#endif