dbscan.h

Go to the documentation of this file.

#pragma once
 
#include <BS_thread_pool.hpp>
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <vector>
 
#include "clustering/always_assert.h"
#include "clustering/index/auto_range_index.h"
#include "clustering/index/range_query.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);
    // Defer pool construction to @ref run: at small shapes every hot phase gates serial, and
    // spawning nJobs workers in the ctor burns tens of microseconds of thread-create futex
    // traffic that the fit never amortizes. @ref run emplaces on first need and reuses.
  }
 
  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.
    if (!m_pool.has_value() && math::shouldSpawnPool(n, m_nJobs, /*minOpsPerWorker=*/16)) {
      m_pool.emplace(m_nJobs);
    }
    const math::Pool pool{m_pool.has_value() ? &*m_pool : nullptr};
 
    QueryModel queryModel(X);
    const auto adj = queryModel.query(m_eps, pool);
 
    std::vector<std::uint8_t> isCore(n, 0);
    for (std::size_t i = 0; i < n; ++i) {
      isCore[i] = (adj[i].size() >= m_minPts) ? 1U : 0U;
    }
 
    std::int32_t *labelPtr = m_labels.data();
    std::fill(labelPtr, labelPtr + n, UNCLASSIFIED);
 
    for (std::size_t i = 0; i < n; ++i) {
      if (labelPtr[i] == UNCLASSIFIED && isCore[i] != 0) {
        expandCluster(i, isCore, adj, labelPtr);
        ++m_clusterId;
      }
    }
 
    std::replace(labelPtr, labelPtr + n, UNCLASSIFIED, NOISY);
  }

  [[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});
    }
  }

  void expandCluster(std::size_t seed, const std::vector<std::uint8_t> &isCore,
                     const std::vector<std::vector<std::int32_t>> &adj, std::int32_t *labels) {
    const auto clusterLabel = static_cast<std::int32_t>(m_clusterId);
    labels[seed] = clusterLabel;
    std::vector<std::size_t> frontier;
    frontier.reserve(adj[seed].size());
    for (const std::int32_t n : adj[seed]) {
      frontier.push_back(static_cast<std::size_t>(n));
    }
 
    while (!frontier.empty()) {
      std::vector<std::size_t> next;
      next.reserve(frontier.size());
      for (const std::size_t point : frontier) {
        if (labels[point] != UNCLASSIFIED) {
          continue;
        }
        labels[point] = clusterLabel;
        if (isCore[point] == 0) {
          continue;
        }
        for (const std::int32_t n : adj[point]) {
          if (labels[static_cast<std::size_t>(n)] == UNCLASSIFIED) {
            next.push_back(static_cast<std::size_t>(n));
          }
        }
      }
      frontier.swap(next);
    }
  }
 
  T m_eps;              
  std::size_t m_minPts; 
  std::size_t m_nJobs;  
  std::size_t m_clusterId = 0; 
  std::optional<BS::light_thread_pool> m_pool;
  NDArray<std::int32_t, 1> m_labels;
};
 
} // namespace clustering