hdbscan.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/hdbscan/detail/condensed_tree.h"
#include "clustering/hdbscan/detail/eom_extract.h"
#include "clustering/hdbscan/detail/glosh.h"
#include "clustering/hdbscan/detail/leaf_extract.h"
#include "clustering/hdbscan/detail/single_linkage.h"
#include "clustering/hdbscan/mst_backend.h"
#include "clustering/hdbscan/mst_output.h"
#include "clustering/hdbscan/policy/auto_mst_backend.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering::hdbscan {

enum class ClusterSelectionMethod : std::uint8_t {
  kEom,  
  kLeaf, 
};

enum class MinSamplesConvention : std::uint8_t {
  kSklearn,
  kCampello,
};
 
} // namespace clustering::hdbscan
 
namespace clustering {

template <class T, class MstBackend = hdbscan::AutoMstBackend<T>>
  requires hdbscan::MstBackendStrategy<MstBackend, T>
class HDBSCAN {
  static_assert(std::is_same_v<T, float>,
                "HDBSCAN<T> supports only float; a double specialization is out of scope.");
 
public:
  struct CondensedTreeView {
    std::span<const std::int32_t> parent;
    std::span<const std::int32_t> child;
    std::span<const T> lambda;
    std::span<const std::int32_t> childSize;

    [[nodiscard]] bool empty() const noexcept { return parent.empty(); }
    [[nodiscard]] std::size_t size() const noexcept { return parent.size(); }
  };

  explicit HDBSCAN(
      std::size_t minClusterSize, std::size_t minSamples = 0,
      hdbscan::ClusterSelectionMethod method = hdbscan::ClusterSelectionMethod::kEom,
      std::size_t nJobs = 0,
      hdbscan::MinSamplesConvention convention = hdbscan::MinSamplesConvention::kSklearn)
      : m_minClusterSize(minClusterSize), m_minSamples(minSamples), m_method(method),
        m_nJobs(math::clampedJobCount(nJobs)), m_convention(convention), m_labels({0}),
        m_outlierScores({0}) {
    CLUSTERING_ALWAYS_ASSERT(minClusterSize >= 2);
    // Pool acquisition is deferred to @ref run: at small shapes every hot phase gates serial, so
    // a fully-serial fit should not force the shared registry to materialize a worker pool. The
    // shape-gated borrow from @ref math::sharedPool happens inside @ref run.
  }
 
  HDBSCAN(const HDBSCAN &) = delete;
  HDBSCAN &operator=(const HDBSCAN &) = delete;
  HDBSCAN(HDBSCAN &&) = delete;
  HDBSCAN &operator=(HDBSCAN &&) = delete;
  ~HDBSCAN() = default;

  void run(const NDArray<T, 2> &X) {
    const std::size_t n = X.dim(0);
 
    CLUSTERING_ALWAYS_ASSERT(m_minClusterSize >= 2);
 
    // Translate the user-facing @c minSamples to the non-self neighbour count the MST backends
    // consume. Under the scikit-learn convention the query point counts as one of the neighbours
    // and we feed the backend one less; under Campello we pass it through.
    const std::size_t requestedMinSamples = (m_minSamples == 0) ? m_minClusterSize : m_minSamples;
    if (m_convention == hdbscan::MinSamplesConvention::kSklearn) {
      CLUSTERING_ALWAYS_ASSERT(requestedMinSamples >= 2);
    } else {
      CLUSTERING_ALWAYS_ASSERT(requestedMinSamples >= 1);
    }
    const std::size_t effectiveMinSamples =
        (m_convention == hdbscan::MinSamplesConvention::kSklearn) ? requestedMinSamples - 1
                                                                  : requestedMinSamples;
    CLUSTERING_ALWAYS_ASSERT(effectiveMinSamples >= 1);
    CLUSTERING_ALWAYS_ASSERT(effectiveMinSamples < n);
    CLUSTERING_ALWAYS_ASSERT(n >= m_minClusterSize);
    CLUSTERING_ALWAYS_ASSERT(n <=
                             static_cast<std::size_t>(std::numeric_limits<std::int32_t>::max()));
 
    // Shape-gated pool borrow: the MST backend is the only phase that can parallelise at all;
    // post-MST is serial by contract. Size the work gate by the backend's dominant shape (N * N)
    // so small inputs stay fully serial and never touch the shared registry.
    const math::Pool pool{math::shouldSpawnPool(n * n, m_nJobs, /*minOpsPerWorker=*/1U << 15)
                              ? &math::sharedPool(m_nJobs)
                              : nullptr};
 
    // Phase 1: MST via the pinned backend. The backend writes edges and core distances into
    // `m_mstOutput`.
    m_backend.run(X, effectiveMinSamples, pool, m_mstOutput);
 
    // Convert squared distances to linear distances before the post-MST pipeline consumes them.
    // The backends store squared Euclidean internally (avoids an @c sqrt per pair-distance); the
    // MST structure is invariant under @c d -> `sqrt(d)` (monotone) but the condensed-tree
    // stability DP compares absolute lambda values whose outcome is not invariant under that
    // transform. Linearising here aligns the lambda scale with the reference implementation and
    // with outlier-score bounds users expect.
    {
      const std::size_t nCore = m_mstOutput.coreDistances.dim(0);
      T *coreData = m_mstOutput.coreDistances.data();
      for (std::size_t i = 0; i < nCore; ++i) {
        coreData[i] = std::sqrt(coreData[i]);
      }
      for (auto &edge : m_mstOutput.edges) {
        edge.weight = std::sqrt(edge.weight);
      }
    }
 
    // Phase 2: build the single-linkage dendrogram from the MST edges.
    hdbscan::detail::SingleLinkageTree<T> slt;
    hdbscan::detail::buildSingleLinkageTree(m_mstOutput, n, slt);
 
    // Phase 3: condense the dendrogram under `minClusterSize`.
    hdbscan::detail::CondensedTree<T> condensed;
    hdbscan::detail::condenseTree(slt, n, m_minClusterSize, condensed);
 
    // Phase 4: cluster extraction (EOM or leaf).
    std::vector<std::int32_t> labels;
    if (m_method == hdbscan::ClusterSelectionMethod::kEom) {
      hdbscan::detail::extractEom(condensed, n, labels);
    } else {
      hdbscan::detail::extractLeaf(condensed, n, labels);
    }
 
    // Phase 5: GLOSH outlier scores. The per-point score loop fans out over the same pool the MST
    // borrowed; its internal subtree-max precompute stays serial.
    std::vector<T> scores;
    hdbscan::detail::computeGlosh(condensed, n, labels, scores, pool);
 
    // Finalise result accessors. The label array lands in the public NDArray buffer; ditto the
    // outlier-score array. The condensed tree is retained in its parallel-array form so the
    // public view can borrow from it without an additional copy.
    m_labels = NDArray<std::int32_t, 1>(std::array<std::size_t, 1>{n});
    std::int32_t maxLabel = -1;
    for (std::size_t i = 0; i < n; ++i) {
      m_labels(i) = labels[i];
      maxLabel = std::max(maxLabel, labels[i]);
    }
    m_outlierScores = NDArray<T, 1>(std::array<std::size_t, 1>{n});
    for (std::size_t i = 0; i < n; ++i) {
      m_outlierScores(i) = scores[i];
    }
    m_nClusters = (maxLabel < 0) ? std::size_t{0} : static_cast<std::size_t>(maxLabel) + 1;
 
    m_ctParent = std::move(condensed.parent);
    m_ctChild = std::move(condensed.child);
    m_ctLambda = std::move(condensed.lambdaVal);
    m_ctChildSize = std::move(condensed.childSize);
  }

  [[nodiscard]] const NDArray<std::int32_t, 1> &labels() const noexcept { return m_labels; }

  [[nodiscard]] const NDArray<T, 1> &outlierScores() const noexcept { return m_outlierScores; }

  [[nodiscard]] std::size_t nClusters() const noexcept { return m_nClusters; }

  [[nodiscard]] CondensedTreeView condensedTree() const noexcept {
    return CondensedTreeView{
        .parent = std::span<const std::int32_t>(m_ctParent.data(), m_ctParent.size()),
        .child = std::span<const std::int32_t>(m_ctChild.data(), m_ctChild.size()),
        .lambda = std::span<const T>(m_ctLambda.data(), m_ctLambda.size()),
        .childSize = std::span<const std::int32_t>(m_ctChildSize.data(), m_ctChildSize.size()),
    };
  }

  void reset() {
    m_labels = NDArray<std::int32_t, 1>({0});
    m_outlierScores = NDArray<T, 1>({0});
    m_nClusters = 0;
    m_ctParent = std::vector<std::int32_t>{};
    m_ctChild = std::vector<std::int32_t>{};
    m_ctLambda = std::vector<T>{};
    m_ctChildSize = std::vector<std::int32_t>{};
    m_mstOutput = hdbscan::MstOutput<T>{};
    m_backend = MstBackend{};
  }
 
private:
  std::size_t m_minClusterSize;
  std::size_t m_minSamples;
  hdbscan::ClusterSelectionMethod m_method;
  std::size_t m_nJobs;
  hdbscan::MinSamplesConvention m_convention;
  NDArray<std::int32_t, 1> m_labels;
  NDArray<T, 1> m_outlierScores;
  std::size_t m_nClusters = 0;
 
  // Condensed-tree parallel arrays, filled by the post-MST pipeline and surfaced through
  // @ref condensedTree as a read-only view.
  std::vector<std::int32_t> m_ctParent;
  std::vector<std::int32_t> m_ctChild;
  std::vector<T> m_ctLambda;
  std::vector<std::int32_t> m_ctChildSize;
 
  MstBackend m_backend{};
  hdbscan::MstOutput<T> m_mstOutput{};
};
 
} // namespace clustering