nn_descent_mst_backend.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/hdbscan/mst_output.h"
#include "clustering/index/nn_descent.h"
#include "clustering/math/detail/avx2_helpers.h"
#include "clustering/math/dsu.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering::hdbscan {

struct NnDescentMstConfig {
  std::size_t kExtra = 10;
  std::size_t maxIter = 10;
  float delta = 0.001F;
  std::uint64_t seed = 0;
};

template <class T> class NnDescentMstBackend {
  static_assert(
      std::is_same_v<T, float>,
      "NnDescentMstBackend<T> supports only float; a double specialization is out of scope.");
 
public:
  NnDescentMstBackend() = default;

  explicit NnDescentMstBackend(NnDescentMstConfig config) : m_config(config) {}

  void run(const NDArray<T, 2> &X, std::size_t minSamples, math::Pool pool, MstOutput<T> &out) {
    const std::size_t n = X.dim(0);
    CLUSTERING_ALWAYS_ASSERT(minSamples >= 1);
    CLUSTERING_ALWAYS_ASSERT(minSamples < n);
 
    out.edges.clear();
    out.edges.reserve(n - 1);
    out.coreDistances = NDArray<T, 1>(std::array<std::size_t, 1>{n});
 
    // Phase 1: rebuild the NN-Descent kNN graph per fit. The kNN graph encodes the current
    // input's neighbour structure; caching it across fits would risk silent failure under
    // in-place mutation of the caller's buffer. The index is reconstructed on every call so
    // the graph always reflects the data @c run was invoked with.
    const std::size_t k = minSamples + m_config.kExtra;
    CLUSTERING_ALWAYS_ASSERT(k < n);
    m_index.emplace(k, m_config.maxIter, m_config.delta, m_config.seed);
    m_index->build(X, pool);
 
    const auto &neighbors = m_index->neighbors();
 
    // Phase 2: core distance per point is the minSamples-th nearest squared distance. The
    // NN-Descent graph returns entries sorted ascending by squared distance; index minSamples - 1
    // picks the correct slot. We assume the graph's k is at least minSamples so the slot exists.
    T *coreDistData = out.coreDistances.data();
    for (std::size_t i = 0; i < n; ++i) {
      CLUSTERING_ALWAYS_ASSERT(neighbors[i].size() >= minSamples);
      coreDistData[i] = neighbors[i][minSamples - 1].sqDist;
    }
 
    // Phase 3: MRD-weighted edge list. The kNN graph is directed (i -> j when j is a neighbour of
    // i); to form an undirected edge set we emit (min(i,j), max(i,j), weight) and dedupe by
    // sorting plus a forward sweep. The weight takes the max of core[i], core[j], and sqDist(i,j).
    struct Edge {
      T weight;
      std::int32_t u;
      std::int32_t v;
    };
    std::vector<Edge> edges;
    edges.reserve(n * k);
    for (std::size_t i = 0; i < n; ++i) {
      const T coreI = coreDistData[i];
      for (const auto &e : neighbors[i]) {
        const auto j = static_cast<std::size_t>(e.idx);
        if (j == i) {
          continue;
        }
        const T coreJ = coreDistData[j];
        T w = e.sqDist;
        if (coreI > w) {
          w = coreI;
        }
        if (coreJ > w) {
          w = coreJ;
        }
        const auto u32 = static_cast<std::int32_t>(std::min(i, j));
        const auto v32 = static_cast<std::int32_t>(std::max(i, j));
        edges.push_back(Edge{w, u32, v32});
      }
    }
 
    // Sort ascending by weight then by endpoints so Kruskal consumes a deterministic order and
    // duplicate (u, v) entries from the directed kNN view become adjacent. We do not explicitly
    // dedupe: Kruskal's union-find check rejects a second edge between the same pair implicitly.
    std::sort(edges.begin(), edges.end(), [](const Edge &a, const Edge &b) {
      if (a.weight != b.weight) {
        return a.weight < b.weight;
      }
      if (a.u != b.u) {
        return a.u < b.u;
      }
      return a.v < b.v;
    });
 
    // Phase 4: Kruskal on the sorted edge list. The UnionFind is keyed on the signed int32 MST
    // index type; cast via uint32 for the union-find that requires an unsigned index.
    UnionFind<std::uint32_t> uf(n);
    for (const Edge &e : edges) {
      if (uf.unite(static_cast<std::uint32_t>(e.u), static_cast<std::uint32_t>(e.v))) {
        out.edges.push_back(MstEdge<T>{e.u, e.v, e.weight});
        if (out.edges.size() + 1 == n) {
          break;
        }
      }
    }
 
    // Phase 5: connectivity fallback. When the kNN-graph MRD-MST leaves more than one component,
    // enumerate minimum-weight bridges between every pair of disjoint components and Kruskal them
    // in until a single spanning tree remains. Each bridge candidate is the (i, j) MRD edge of
    // minimum weight where i and j live in distinct components; weight uses the true squared
    // Euclidean distance computed row-by-row, then lifted by max with both core distances.
    if (uf.countComponents() > 1) {
      resolveDisconnectedComponents(X, coreDistData, pool, uf, out);
    }
  }
 
private:
  void resolveDisconnectedComponents(const NDArray<T, 2> &X, const T *coreDistData, math::Pool pool,
                                     UnionFind<std::uint32_t> &uf, MstOutput<T> &out) {
    const std::size_t n = X.dim(0);
    const std::size_t d = X.dim(1);
 
    // Materialize the current component membership as a root -> member list so subsequent
    // all-pairs scans iterate only over the members of each component rather than the full
    // point set per pair.
    std::vector<std::vector<std::uint32_t>> members;
    std::vector<std::uint32_t> rootToSlot(n, std::numeric_limits<std::uint32_t>::max());
    for (std::uint32_t i = 0; i < static_cast<std::uint32_t>(n); ++i) {
      const std::uint32_t r = uf.find(i);
      std::uint32_t slot = rootToSlot[r];
      if (slot == std::numeric_limits<std::uint32_t>::max()) {
        slot = static_cast<std::uint32_t>(members.size());
        rootToSlot[r] = slot;
        members.emplace_back();
      }
      members[slot].push_back(i);
    }
 
    while (uf.countComponents() > 1) {
      // For each pair of surviving components (a, b) find the minimum-MRD bridge. Compare every
      // member of a to every member of b; track the minimum-weight edge. When all pairs are
      // scanned, Kruskal the resulting bridges into the MST. The outer loop runs c - 1 times in
      // the worst case but typically terminates after one round because inserting a single bridge
      // collapses components transitively.
      struct Bridge {
        T weight;
        std::int32_t u;
        std::int32_t v;
      };
 
      // Materialise the live (a, b) work list before fan-out so the parallel body indexes a flat
      // array. Skipping empty slots and slots that already share a root keeps the worker count
      // honest after the first Kruskal round collapses components transitively.
      struct PairIdx {
        std::uint32_t a;
        std::uint32_t b;
      };
      std::vector<PairIdx> pairs;
      pairs.reserve(members.size() * (members.size() - 1) / 2);
      for (std::uint32_t a = 0; a < static_cast<std::uint32_t>(members.size()); ++a) {
        if (members[a].empty()) {
          continue;
        }
        const std::uint32_t rootA = uf.find(members[a][0]);
        for (std::uint32_t b = a + 1; b < static_cast<std::uint32_t>(members.size()); ++b) {
          if (members[b].empty()) {
            continue;
          }
          if (uf.find(members[b][0]) == rootA) {
            continue;
          }
          pairs.push_back(PairIdx{a, b});
        }
      }
 
      std::vector<Bridge> bridges(pairs.size(), Bridge{std::numeric_limits<T>::infinity(), 0, 0});
 
      auto computePair = [&](std::size_t pairIdx) noexcept {
        const auto &p = pairs[pairIdx];
        const auto &memA = members[p.a];
        const auto &memB = members[p.b];
        T bestW = std::numeric_limits<T>::infinity();
        std::int32_t bestU = 0;
        std::int32_t bestV = 0;
        for (const std::uint32_t ia : memA) {
          const T coreI = coreDistData[ia];
          const T *rowI = X.data() + (static_cast<std::size_t>(ia) * d);
          for (const std::uint32_t jb : memB) {
            const T coreJ = coreDistData[jb];
            const T *rowJ = X.data() + (static_cast<std::size_t>(jb) * d);
            const T sq = math::detail::sqEuclideanRowPtr(rowI, rowJ, d);
            T w = sq;
            if (coreI > w) {
              w = coreI;
            }
            if (coreJ > w) {
              w = coreJ;
            }
            if (w < bestW) {
              bestW = w;
              bestU = static_cast<std::int32_t>(ia);
              bestV = static_cast<std::int32_t>(jb);
            }
          }
        }
        bridges[pairIdx] = Bridge{bestW, bestU, bestV};
      };
 
      // Cost per pair scales with `|memA| * |memB| * d`; the gate uses the largest pair as a
      // proxy because per-pair work is wildly heterogeneous and underestimating any single big
      // pair would leave one worker stuck while others finish.
      std::size_t maxPairOps = 0;
      for (const auto &p : pairs) {
        const std::size_t ops = members[p.a].size() * members[p.b].size() * d;
        maxPairOps = std::max(maxPairOps, ops);
      }
      const std::size_t totalPairOps = maxPairOps * pairs.size();
      if (pool.pool != nullptr && pool.shouldParallelizeWork(totalPairOps)) {
        pool.pool
            ->submit_blocks(std::size_t{0}, pairs.size(),
                            [&](std::size_t lo, std::size_t hi) {
                              for (std::size_t i = lo; i < hi; ++i) {
                                computePair(i);
                              }
                            })
            .wait();
      } else {
        for (std::size_t i = 0; i < pairs.size(); ++i) {
          computePair(i);
        }
      }
 
      // Drop the infinity sentinel before the sort so Kruskal never sees an unset slot. In
      // practice every pair admits at least one finite bridge (cross-component points exist).
      bridges.erase(std::remove_if(bridges.begin(), bridges.end(),
                                   [](const Bridge &br) {
                                     return br.weight == std::numeric_limits<T>::infinity();
                                   }),
                    bridges.end());
 
      // Sort bridges ascending by weight and Kruskal them into the MST. A bridge whose endpoints
      // are already merged (via a prior bridge in the same round) is skipped by union-find.
      std::sort(bridges.begin(), bridges.end(),
                [](const Bridge &a, const Bridge &b) { return a.weight < b.weight; });
      bool progress = false;
      for (const Bridge &br : bridges) {
        if (uf.unite(static_cast<std::uint32_t>(br.u), static_cast<std::uint32_t>(br.v))) {
          out.edges.push_back(MstEdge<T>{br.u, br.v, br.weight});
          progress = true;
          if (out.edges.size() + 1 == n) {
            break;
          }
        }
      }
 
      // Guard against an infinite loop: if no bridge was accepted the graph cannot reach
      // connectivity. In practice the precondition (cross-component points exist) guarantees a
      // finite bridge and this branch is never taken.
      CLUSTERING_ALWAYS_ASSERT(progress);
 
      // Rebuild the component membership for the next iteration: a merge may collapse multiple
      // slots into one root, so the next round's all-pairs scan should only visit surviving
      // components. Slots whose root changed get their members absorbed into the winning slot.
      if (uf.countComponents() > 1) {
        std::vector<std::vector<std::uint32_t>> nextMembers;
        std::fill(rootToSlot.begin(), rootToSlot.end(), std::numeric_limits<std::uint32_t>::max());
        for (const auto &comp : members) {
          for (const std::uint32_t i : comp) {
            const std::uint32_t r = uf.find(i);
            std::uint32_t slot = rootToSlot[r];
            if (slot == std::numeric_limits<std::uint32_t>::max()) {
              slot = static_cast<std::uint32_t>(nextMembers.size());
              rootToSlot[r] = slot;
              nextMembers.emplace_back();
            }
            nextMembers[slot].push_back(i);
          }
        }
        members = std::move(nextMembers);
      }
    }
  }
 
  NnDescentMstConfig m_config{};
  std::optional<index::NnDescentIndex<T>> m_index;
};
 
} // namespace clustering::hdbscan