auto_mst_backend.h

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <type_traits>
#include <variant>
 
#include "clustering/hdbscan/mst_output.h"
#include "clustering/hdbscan/policy/boruvka_mst_backend.h"
#include "clustering/hdbscan/policy/nn_descent_mst_backend.h"
#include "clustering/hdbscan/policy/prim_mst_backend.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering::hdbscan {

template <class T> class AutoMstBackend {
  static_assert(std::is_same_v<T, float>,
                "AutoMstBackend<T> supports only float; a double specialization is out of scope.");
 
public:
#ifdef CLUSTERING_HDBSCAN_BORUVKA_LOW_DIM_CEIL
  static constexpr std::size_t boruvkaLowDimCeil = CLUSTERING_HDBSCAN_BORUVKA_LOW_DIM_CEIL;
#else
  static constexpr std::size_t boruvkaLowDimCeil = 16;
#endif
 
#ifdef CLUSTERING_HDBSCAN_BORUVKA_DIM_CEIL
  static constexpr std::size_t boruvkaDimCeil = CLUSTERING_HDBSCAN_BORUVKA_DIM_CEIL;
#else
  static constexpr std::size_t boruvkaDimCeil = 60;
#endif
 
  static_assert(boruvkaLowDimCeil <= boruvkaDimCeil,
                "boruvkaLowDimCeil must not exceed boruvkaDimCeil; the dispatch order assumes a "
                "nested Boruvka regime at low d and a relaxed one at moderate d.");

  static constexpr bool primFitsBudget(std::size_t n) noexcept {
    if (n == 0) {
      return true;
    }
    constexpr std::size_t kNsqBudget = kPrimMrdMatrixByteBudget / sizeof(T);
    return n <= kNsqBudget / n;
  }
 
  AutoMstBackend() = default;

  void run(const NDArray<T, 2> &X, std::size_t minSamples, math::Pool pool, MstOutput<T> &out) {
    ensureShape(X.dim(0), X.dim(1));
    std::visit([&](auto &b) { b.run(X, minSamples, pool, out); }, m_held);
  }

  [[nodiscard]] std::size_t heldIndex() const noexcept { return m_held.index(); }

  std::size_t peekArm(std::size_t n, std::size_t d) {
    ensureShape(n, d);
    return m_held.index();
  }
 
private:
  void ensureShape(std::size_t n, std::size_t d) {
    if (n == m_lastN && d == m_lastD) {
      return;
    }
    if (d <= boruvkaLowDimCeil) {
      if (!std::holds_alternative<BoruvkaMstBackend<T>>(m_held)) {
        m_held.template emplace<BoruvkaMstBackend<T>>();
      }
    } else if (d <= boruvkaDimCeil && primFitsBudget(n)) {
      // Streaming Prim beats both Boruvka and NN-Descent in the `d <= 60` band while @c n
      // stays inside its quadratic compute budget. Above @c boruvkaDimCeil the dense pairwise
      // work is dominated by the @c d-wide distance compute and NN-Descent wins; the cap keeps
      // Prim out of that regime.
      if (!std::holds_alternative<PrimMstBackend<T>>(m_held)) {
        m_held.template emplace<PrimMstBackend<T>>();
      }
    } else if (d <= boruvkaDimCeil) {
      if (!std::holds_alternative<BoruvkaMstBackend<T>>(m_held)) {
        m_held.template emplace<BoruvkaMstBackend<T>>();
      }
    } else {
      if (!std::holds_alternative<NnDescentMstBackend<T>>(m_held)) {
        m_held.template emplace<NnDescentMstBackend<T>>();
      }
    }
    m_lastN = n;
    m_lastD = d;
  }
 
  // Variant order is fixed: Boruvka (0), Prim (1), NN-Descent (2). @ref heldIndex leans on this.
  std::variant<BoruvkaMstBackend<T>, PrimMstBackend<T>, NnDescentMstBackend<T>> m_held{
      std::in_place_type<PrimMstBackend<T>>};
  std::size_t m_lastN = 0;
  std::size_t m_lastD = 0;
};
 
} // namespace clustering::hdbscan