kmeans.h

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <cstdint>
#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_pool(m_nJobs > 1 ? &math::sharedPool(m_nJobs) : nullptr), m_centroids({0, 0}),
        m_labels({0}) {
    CLUSTERING_ALWAYS_ASSERT(k >= 1);
  }

  KMeans(std::size_t k, std::size_t nJobs, math::OwnedPool &externalPool)
      : m_k(k), m_nJobs(math::clampedJobCount(nJobs)), m_pool(&externalPool), m_centroids({0, 0}),
        m_labels({0}) {
    CLUSTERING_ALWAYS_ASSERT(k >= 1);
  }
 
  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 seedFirst = 0, std::size_t nInit = 1) {
    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);
    CLUSTERING_ALWAYS_ASSERT(nInit >= 1);
 
    ensureOutputShape(n, d);
 
    if (n == 0 || d == 0) {
      m_nIter = 0;
      m_converged = true;
      m_inertia = 0.0;
      return;
    }
 
    const math::Pool pool{m_pool};
 
    if (nInit == 1) {
      m_seeder.run(X, m_k, seedFirst, pool, m_centroids);
      m_lloyd.run(X, m_centroids, m_k, maxIter, tol, pool, m_labels, m_inertia, m_nIter,
                  m_converged);
      return;
    }
 
    ensureBestOfScratch(n, d);
    bool anyImprovement = false;
 
    for (std::size_t i = 0; i < nInit; ++i) {
      const std::uint64_t seed = seedFirst + static_cast<std::uint64_t>(i);
      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);
      if (!anyImprovement || m_inertia < m_bestInertia) {
        anyImprovement = true;
        m_bestInertia = m_inertia;
        m_bestNIter = m_nIter;
        m_bestConverged = m_converged;
        std::memcpy(m_bestCentroids.data(), m_centroids.data(), m_k * d * sizeof(T));
        std::memcpy(m_bestLabels.data(), m_labels.data(), n * sizeof(std::int32_t));
      }
    }
 
    // Promote the best-of-restarts result onto the public accessors. Swapping the storage
    // (rather than copying) keeps the work-side buffers the same shape so the next call
    // skips reallocation.
    std::swap(m_centroids, m_bestCentroids);
    std::swap(m_labels, m_bestLabels);
    m_inertia = m_bestInertia;
    m_nIter = m_bestNIter;
    m_converged = m_bestConverged;
  }

  [[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_bestCentroids = NDArray<T, 2, Layout::Contig>({0, 0});
    m_bestLabels = NDArray<std::int32_t, 1>({0});
    m_inertia = 0.0;
    m_nIter = 0;
    m_converged = false;
    m_bestInertia = 0.0;
    m_bestNIter = 0;
    m_bestConverged = 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});
    }
  }
 
  void ensureBestOfScratch(std::size_t n, std::size_t d) {
    if (m_bestCentroids.dim(0) != m_k || m_bestCentroids.dim(1) != d) {
      m_bestCentroids = NDArray<T, 2, Layout::Contig>({m_k, d});
    }
    if (m_bestLabels.dim(0) != n) {
      m_bestLabels = NDArray<std::int32_t, 1>({n});
    }
  }
 
  std::size_t m_k;
  std::size_t m_nJobs;
  math::OwnedPool *m_pool = nullptr;
  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;
 
  // Scratch holding the best-of-restarts result while iterating (@c nInit > 1). On exit
  // these are swapped with the public buffers so accessors return the lowest-inertia run.
  NDArray<T, 2, Layout::Contig> m_bestCentroids{NDArray<T, 2, Layout::Contig>({0, 0})};
  NDArray<std::int32_t, 1> m_bestLabels{NDArray<std::int32_t, 1>({0})};
  double m_bestInertia = 0.0;
  std::size_t m_bestNIter = 0;
  bool m_bestConverged = false;
 
  Algo m_lloyd{};
  Seeder m_seeder{};
};
 
} // namespace clustering