nn_descent.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/index/nn_descent/detail/join_step.h"
#include "clustering/index/nn_descent/detail/neighbor_heap.h"
#include "clustering/index/nn_descent/detail/rp_tree_init.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering::index {

template <class T> class NnDescentIndex {
  static_assert(std::is_same_v<T, float>,
                "NnDescentIndex<T> supports only float; add a specialization to extend.");
 
public:
  struct KnnEntry {
    std::int32_t idx;
    T sqDist;
  };

  explicit NnDescentIndex(std::size_t k, std::size_t maxIter = 10, T delta = T{0.001},
                          std::uint64_t seed = 0)
      : m_k(k), m_maxIter(maxIter), m_delta(delta), m_seed(seed) {
    CLUSTERING_ALWAYS_ASSERT(k >= 1);
  }
 
  NnDescentIndex(const NnDescentIndex &) = delete;
  NnDescentIndex &operator=(const NnDescentIndex &) = delete;
  NnDescentIndex(NnDescentIndex &&) = default;
  NnDescentIndex &operator=(NnDescentIndex &&) = default;
  ~NnDescentIndex() = default;

  void build(const NDArray<T, 2> &X, math::Pool pool) {
    const std::size_t n = X.dim(0);
    const std::size_t d = X.dim(1);
    const bool kFitsN = (n == 0) || (m_k < n);
    CLUSTERING_ALWAYS_ASSERT(kFitsN);
 
    const bool sameShape =
        (m_lastN == n) && (m_lastD == d) && (m_lastK == m_k) && m_bank.has_value();
 
    m_lastIterations = 0;
 
    if (!sameShape) {
      // Cold start: discard any prior bank and rebuild from RP-tree init.
      m_bank.emplace(n, m_k);
      if (n == 0 || m_k == 0) {
        captureShape(X, n, d);
        m_neighborsView.assign(n, {});
        return;
      }
      // Leaf limit sized to roughly @c 2k so the leaf-pair enumeration is `O(k^2)` per leaf,
      // matching Dong 2011's recommendation. A minimum of `max(2k, 8)` keeps very small @c k
      // from degenerating into singleton leaves.
      const std::size_t leafLimit = std::max<std::size_t>(2 * m_k, 8);
      nn_descent::detail::RpTreeInit<T>::build(X, leafLimit, m_seed, *m_bank);
    } else {
      // Warm start: reload the bank from the previously published sorted view so the heap
      // invariant is restored after the destructive @c toSortedLists pass, then re-flag every
      // slot as "new" so the next join iteration revisits them.
      reloadBankFromView();
    }
 
    if (n == 0 || m_k == 0) {
      captureShape(X, n, d);
      m_neighborsView.assign(n, {});
      return;
    }
 
    // Join loop. Each iteration counts updated slots; once the fraction drops below @p delta the
    // graph is considered converged.
    const std::size_t total = n * m_k;
    for (std::size_t iter = 0; iter < m_maxIter; ++iter) {
      const std::size_t updates = nn_descent::detail::JoinStep<T>::run(X, *m_bank, pool);
      ++m_lastIterations;
      if (total == 0) {
        break;
      }
      const double ratio = static_cast<double>(updates) / static_cast<double>(total);
      if (ratio < static_cast<double>(m_delta)) {
        break;
      }
    }
 
    captureShape(X, n, d);
 
    // Materialize the public neighbor view, sorted ascending by squared distance. This
    // destroys the bank's heap order; a subsequent warm-start call reheapifies via
    // @c reloadBankFromView below.
    m_neighborsView = m_bank->template toSortedLists<KnnEntry>();
    m_bankOrderValid = false;
  }

  [[nodiscard]] const std::vector<std::vector<KnnEntry>> &neighbors() const noexcept {
    return m_neighborsView;
  }

  [[nodiscard]] bool isConnected() const {
    const std::size_t n = m_neighborsView.size();
    if (n <= 1) {
      return true;
    }
    std::vector<std::vector<std::int32_t>> adj(n);
    for (std::size_t u = 0; u < n; ++u) {
      for (const KnnEntry &e : m_neighborsView[u]) {
        if (e.idx < 0 || std::cmp_greater_equal(e.idx, n)) {
          continue;
        }
        adj[u].push_back(e.idx);
        adj[static_cast<std::size_t>(e.idx)].push_back(static_cast<std::int32_t>(u));
      }
    }
    std::vector<std::uint8_t> visited(n, 0);
    std::vector<std::int32_t> stack;
    stack.reserve(n);
    stack.push_back(0);
    visited[0] = 1;
    std::size_t visitedCount = 1;
    while (!stack.empty()) {
      const std::int32_t u = stack.back();
      stack.pop_back();
      for (const std::int32_t v : adj[static_cast<std::size_t>(u)]) {
        if (visited[static_cast<std::size_t>(v)] == 0U) {
          visited[static_cast<std::size_t>(v)] = 1;
          ++visitedCount;
          stack.push_back(v);
        }
      }
    }
    return visitedCount == n;
  }

  [[nodiscard]] std::size_t k() const noexcept { return m_k; }

  [[nodiscard]] std::size_t lastIterations() const noexcept { return m_lastIterations; }
 
private:
  void reloadBankFromView() {
    if (!m_bank.has_value()) {
      return;
    }
    if (m_bankOrderValid) {
      // Bank is already in heap order; we only need to re-arm the "new" flag so the join loop
      // revisits every slot.
      m_bank->rearmAllAsNew();
      return;
    }
    std::vector<std::pair<T, std::int32_t>> buf;
    buf.reserve(m_k);
    for (std::size_t i = 0; i < m_neighborsView.size(); ++i) {
      buf.clear();
      for (const KnnEntry &e : m_neighborsView[i]) {
        buf.emplace_back(e.sqDist, e.idx);
      }
      m_bank->loadFromSorted(static_cast<std::int32_t>(i), buf);
    }
    m_bankOrderValid = true;
  }
 
  void captureShape(const NDArray<T, 2> & /*X*/, std::size_t n, std::size_t d) noexcept {
    m_lastN = n;
    m_lastD = d;
    m_lastK = m_k;
  }
 
  std::size_t m_k;
  std::size_t m_maxIter;
  T m_delta;
  std::uint64_t m_seed;
 
  std::optional<nn_descent::detail::NeighborHeapBank<T>> m_bank;
  std::vector<std::vector<KnnEntry>> m_neighborsView;
 
  std::size_t m_lastN = 0;
  std::size_t m_lastD = 0;
  std::size_t m_lastK = 0;
  std::size_t m_lastIterations = 0;
  bool m_bankOrderValid = false;
};
 
} // namespace clustering::index