dbscan.h

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/index/auto_range_index.h"
#include "clustering/index/range_query.h"
#include "clustering/math/dsu.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering {

template <class T, class QueryModel = index::AutoRangeIndex<T>>
  requires index::RangeIndex<QueryModel, T>
class DBSCAN {
public:
  static constexpr std::int32_t UNCLASSIFIED = -2; 
  static constexpr std::int32_t NOISY = -1; 

  explicit DBSCAN(T eps, std::size_t minPts, std::size_t nJobs = 0)
      : m_eps(eps), m_minPts(minPts), m_nJobs(math::clampedJobCount(nJobs)), m_labels({0}) {
    CLUSTERING_ALWAYS_ASSERT(minPts >= 1);
    // 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.
  }
 
  DBSCAN(const DBSCAN &) = delete;
  DBSCAN &operator=(const DBSCAN &) = delete;
  DBSCAN(DBSCAN &&) = delete;
  DBSCAN &operator=(DBSCAN &&) = delete;
  ~DBSCAN() = default;

  void run(const NDArray<T, 2> &X) {
    const std::size_t n = X.dim(0);
    ensureLabelsShape(n);
    m_clusterId = 0;
 
    if (n == 0) {
      return;
    }
 
    // One adjacency query per point is the unit of work the range-index backends fan out on;
    // shouldSpawnPool with minOpsPerWorker=16 matches the KDTree adjacency sweep's own
    // `shouldParallelize(n, 4, 2)` gate (@c n / 4 >= 2 * workerCount => `n >= 8` * workerCount)
    // so the pool spawn and the backend fan-out fire at the same shape. @c n * d would
    // under-estimate DBSCAN work -- the backend does per-query tree walks that are much heavier
    // than @c d ops -- and caused @c n_jobs=16 at low @c d to fall back to serial here while
    // every worker-side kernel would have cleared its own gate.
    const std::size_t poolJobs = effectiveWorkerCount(n, X.dim(1), m_nJobs);
    const math::Pool pool{math::shouldSpawnPool(n, poolJobs, /*minOpsPerWorker=*/16)
                              ? &math::sharedPool(poolJobs)
                              : nullptr};
 
    // Pool-aware query models parallelize their own construction (the KDTree build forks
    // subtrees); models without that constructor keep the plain shape contract.
    QueryModel queryModel = [&] {
      if constexpr (std::is_constructible_v<QueryModel, const NDArray<T, 2> &, math::Pool>) {
        return QueryModel(X, pool);
      } else {
        return QueryModel(X);
      }
    }();
    // The backend derives the core flags from full degrees; core rows may carry only their
    // upper-half neighbours per the @ref clustering::index::CoreAdjacency contract, which is
    // exactly the half the component build below reads.
    const auto [adj, isCore, extraEdges] = queryModel.query(m_eps, m_minPts, pool);
 
    // Connected components over core-core edges: density-reachability is the transitive closure of
    // "core within eps of core", which a disjoint-set union builds directly.
    const std::vector<std::uint32_t> componentRoots =
        buildComponentRoots(adj, isCore, extraEdges, pool);
 
    // Dense cluster ids in first-core-index order so the lowest-index core of each component names
    // its cluster. Writing each core's id into the label buffer now freezes it so the border pass
    // reads it without another root lookup.
    std::int32_t *labels = m_labels.data();
    std::vector<std::int32_t> rootCluster(n, UNCLASSIFIED);
    for (std::size_t i = 0; i < n; ++i) {
      if (isCore[i] == 0) {
        continue;
      }
      const auto root = componentRoots[i];
      std::int32_t &slot = rootCluster[root];
      if (slot == UNCLASSIFIED) {
        slot = static_cast<std::int32_t>(m_clusterId);
        ++m_clusterId;
      }
      labels[i] = slot;
    }
 
    assignBorderAndNoise(adj, isCore);
  }

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

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

  void reset() {
    m_labels = NDArray<std::int32_t, 1>({0});
    m_clusterId = 0;
  }
 
private:
  void ensureLabelsShape(std::size_t n) {
    if (m_labels.dim(0) != n) {
      m_labels = NDArray<std::int32_t, 1>({n});
    }
  }
 
  static constexpr std::size_t kSmall2dMaxParallelN = 25'000;
  static constexpr std::size_t kSmall2dWorkerCap = 8;
 
  [[nodiscard]] static std::size_t effectiveWorkerCount(std::size_t n, std::size_t d,
                                                        std::size_t requested) noexcept {
    if (d <= 2 && n <= kSmall2dMaxParallelN && requested > kSmall2dWorkerCap) {
      return kSmall2dWorkerCap;
    }
    return requested;
  }

  static std::vector<std::uint32_t> buildComponentRoots(
      const std::vector<std::vector<std::int32_t>> &adj, const std::vector<std::uint8_t> &isCore,
      const std::vector<std::pair<std::int32_t, std::int32_t>> &extraEdges, math::Pool pool) {
    const std::size_t n = adj.size();
    std::vector<std::uint32_t> roots(n);
 
    const auto uniteRange = [&](auto &dsu, std::size_t lo, std::size_t hi) {
      for (std::size_t i = lo; i < hi; ++i) {
        if (isCore[i] == 0) {
          continue;
        }
        const auto iu = static_cast<std::uint32_t>(i);
        // No orientation filter: a backend may carry an edge in either direction only (a
        // clique shortcut can thin the mirror side), and repeated unions are idempotent.
        for (const std::int32_t neighbor : adj[i]) {
          const auto j = static_cast<std::size_t>(neighbor);
          if (j != i && isCore[j] != 0) {
            dsu.unite(iu, static_cast<std::uint32_t>(j));
          }
        }
      }
    };
    // Representative edges the backend carried outside the rows; both endpoints are cores by
    // the CoreAdjacency contract, and unions are idempotent, so no filtering is needed.
    const auto uniteExtraRange = [&](auto &dsu, std::size_t lo, std::size_t hi) {
      for (std::size_t e = lo; e < hi; ++e) {
        dsu.unite(static_cast<std::uint32_t>(extraEdges[e].first),
                  static_cast<std::uint32_t>(extraEdges[e].second));
      }
    };
 
    const std::size_t workers = pool.workerCount();
    if (pool.pool == nullptr || workers <= 1) {
      UnionFind<std::uint32_t> components(n);
      uniteRange(components, 0, n);
      uniteExtraRange(components, 0, extraEdges.size());
      for (std::size_t i = 0; i < n; ++i) {
        roots[i] = components.find(static_cast<std::uint32_t>(i));
      }
      return roots;
    }
 
    // More blocks than workers lets dynamic stealing balance the skewed per-row degree; the
    // shared structure makes block placement irrelevant to the result.
    AtomicUnionFind<std::uint32_t> components(n);
    pool.parallelForBlocks(std::size_t{0}, n, workers * 4,
                           [&](std::size_t lo, std::size_t hi) { uniteRange(components, lo, hi); });
    pool.parallelForBlocks(std::size_t{0}, n, std::size_t{0}, [&](std::size_t lo, std::size_t hi) {
      for (std::size_t i = lo; i < hi; ++i) {
        roots[i] = components.find(static_cast<std::uint32_t>(i));
      }
    });
    return roots;
  }

  void assignBorderAndNoise(const std::vector<std::vector<std::int32_t>> &adj,
                            const std::vector<std::uint8_t> &isCore) {
    const std::size_t n = adj.size();
    std::int32_t *labels = m_labels.data();
    for (std::size_t p = 0; p < n; ++p) {
      if (isCore[p] != 0) {
        continue;
      }
      std::int32_t best = NOISY;
      for (const std::int32_t neighbor : adj[p]) {
        const auto q = static_cast<std::size_t>(neighbor);
        if (isCore[q] == 0) {
          continue;
        }
        const std::int32_t cluster = labels[q];
        if (best == NOISY || cluster < best) {
          best = cluster;
        }
      }
      labels[p] = best;
    }
  }
 
  T m_eps;                     
  std::size_t m_minPts;        
  std::size_t m_nJobs;         
  std::size_t m_clusterId = 0; 
  NDArray<std::int32_t, 1> m_labels;
};
 
} // namespace clustering