kmeans.h

Go to the documentation of this file.

#pragma once
 
#include <BS_thread_pool.hpp>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <type_traits>
 
#include "clustering/always_assert.h"
#include "clustering/kmeans/policy/auto_seeder.h"
#include "clustering/kmeans/policy/lloyd.h"
#include "clustering/kmeans/policy/lloyd_fused_gemm.h"
#include "clustering/kmeans/policy/seeder.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering {

template <class T, class Algo = kmeans::LloydFusedGemm<T>, class Seeder = kmeans::AutoSeeder<T>>
  requires kmeans::LloydStrategy<Algo, T> && kmeans::SeederStrategy<Seeder, T>
class KMeans {
  static_assert(std::is_same_v<T, float>,
                "KMeans<T> supports only float; add a double specialization to extend.");
 
public:
  explicit KMeans(std::size_t k, std::size_t nJobs = 0)
      : m_k(k), m_nJobs(math::clampedJobCount(nJobs)), m_centroids({0, 0}), m_labels({0}) {
    CLUSTERING_ALWAYS_ASSERT(k >= 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.
  }
 
  KMeans(const KMeans &) = delete;
  KMeans &operator=(const KMeans &) = delete;
  KMeans(KMeans &&) = delete;
  KMeans &operator=(KMeans &&) = delete;
  ~KMeans() = default;

  void run(const NDArray<T, 2> &X, std::size_t maxIter = 300, T tol = T{1e-4},
           std::uint64_t seed = 0) {
    const std::size_t n = X.dim(0);
    const std::size_t d = X.dim(1);
 
    CLUSTERING_ALWAYS_ASSERT(m_k >= 1);
    CLUSTERING_ALWAYS_ASSERT(n >= m_k);
 
    ensureOutputShape(n, d);
 
    if (n == 0 || d == 0) {
      m_nIter = 0;
      m_converged = true;
      m_inertia = 0.0;
      return;
    }
 
    if (!m_pool.has_value() && math::shouldSpawnPool(n * d * m_k, m_nJobs)) {
      m_pool.emplace(m_nJobs);
    }
    const math::Pool pool{m_pool.has_value() ? &*m_pool : nullptr};
    m_seeder.run(X, m_k, seed, pool, m_centroids);
    m_lloyd.run(X, m_centroids, m_k, maxIter, tol, pool, m_labels, m_inertia, m_nIter, m_converged);
  }

  [[nodiscard]] const NDArray<std::int32_t, 1> &labels() const noexcept { return m_labels; }
  [[nodiscard]] const NDArray<T, 2, Layout::Contig> &centroids() const noexcept {
    return m_centroids;
  }

  [[nodiscard]] double inertia() const noexcept { return m_inertia; }
  [[nodiscard]] std::size_t nIter() const noexcept { return m_nIter; }
  [[nodiscard]] bool converged() const noexcept { return m_converged; }

  void reset() {
    m_centroids = NDArray<T, 2, Layout::Contig>({0, 0});
    m_labels = NDArray<std::int32_t, 1>({0});
    m_inertia = 0.0;
    m_nIter = 0;
    m_converged = false;
    m_lloyd = Algo{};
    m_seeder = Seeder{};
  }
 
private:
  void ensureOutputShape(std::size_t n, std::size_t d) {
    if (m_centroids.dim(0) != m_k || m_centroids.dim(1) != d) {
      m_centroids = NDArray<T, 2, Layout::Contig>({m_k, d});
    }
    if (m_labels.dim(0) != n) {
      m_labels = NDArray<std::int32_t, 1>({n});
    }
  }
 
  std::size_t m_k;
  std::size_t m_nJobs;
  std::optional<BS::light_thread_pool> m_pool;
  NDArray<T, 2, Layout::Contig> m_centroids;
  NDArray<std::int32_t, 1> m_labels;
  double m_inertia = 0.0;
  std::size_t m_nIter = 0;
  bool m_converged = false;
 
  Algo m_lloyd{};
  Seeder m_seeder{};
};
 
} // namespace clustering