Package com.apple.foundationdb.kmeans

Class KMeans

java.lang.Object
com.apple.foundationdb.kmeans.KMeans
@API(INTERNAL) public final class KMeans extends Object
A Lloyd-style k-means++ implementation that makes several independent runs, intended for LOCAL cluster restructuring (SPFresh-style). Self-contained and stateless: a single static fit entry point runs the whole algorithm and returns a KMeans.Result describing the chosen partitioning.

The algorithm assumes the external metric is ordinary L2, but the k-means core optimizes the standard squared-L2 objective internally:

  • assignment uses squared distance,
  • centroid update uses arithmetic mean,
  • k-means++ initialization weights candidate centroids by squared distance to the nearest already-chosen centroid, and
  • run selection compares the per-run k-means objective, preferring the runs whose cluster sizes pass the caller's size predicate.

Optional soft size balancing biases the assignment steps of the Lloyd iterations by the running projected cluster size (see KMeans.SizePenalty), which steers where the centroids settle; set relativeLambda == 0 or sizePenalty == null to disable.

After the main Lloyd loop of a run exits — whether via convergence (changed == 0) or by hitting maxIterations — fit(java.util.SplittableRandom, com.apple.foundationdb.linear.DistanceEstimator, com.apple.foundationdb.util.Lens<V, com.apple.foundationdb.linear.RealVector>, com.apple.foundationdb.util.Lens<C, com.apple.foundationdb.linear.RealVector>, java.util.List<V>, int, int, int, double, com.apple.foundationdb.kmeans.KMeans.SizePenalty, java.util.function.Predicate<int[]>) runs one final assignment pass against the latest centroids by distance alone, without any size penalty. This guarantees that the returned KMeans.Result.assignment(), KMeans.Result.clusterSizes(), and KMeans.Result.distances() describe the nearest-centroid partitioning of KMeans.Result.clusterCentroids(), regardless of how the loop terminated and of whether balancing was active.

  • Method Details

    • overflowQuadraticPenalty

      @Nonnull public static KMeans.SizePenalty overflowQuadraticPenalty()
      Returns a KMeans.SizePenalty that penalizes only overflow above the target cluster size, quadratically: (max(0, projected - target))^2 / max(1, target). A reasonable default for soft size balancing.
      Returns:
      a non-null overflow-quadratic penalty function

    • fit

      public static <V, C> KMeans.Result<C> fit(@Nonnull SplittableRandom random, @Nonnull DistanceEstimator distanceEstimator, @Nonnull Lens<V,RealVector> vectorLens, @Nonnull Lens<C,RealVector> centroidLens, @Nonnull List<V> vectors, int k, int maxIterations, int numRuns, double relativeLambda, @Nullable KMeans.SizePenalty sizePenalty, @Nullable Predicate<int[]> sizePredicate)
      Fits k clusters using Lloyd-style squared-L2 k-means with k-means++ initialization, making numRuns independent runs and returning one of them.

      Each run makes at most maxIterations Lloyd iterations of (assignment, centroid update). Empty clusters detected during the centroid update are reseeded with the data point currently farthest from any centroid. After the main loop, a final assignment pass assigns every vector to its nearest centroid by distance alone, so the run's assignment, cluster sizes, distances and objective describe the nearest-centroid partitioning of its centroids.

      Run selection is two-tier. If the cluster sizes of at least one run pass sizePredicate, fit returns the passing run with the lowest objective; if no run passes, it returns the run with the lowest objective overall, so that the caller still gets a concrete partitioning to judge and report. Equivalently, the runs are ordered by passing before failing, then by lower objective, and the first one is returned. The predicate is thus a preference, not a guarantee: a caller that needs the sizes to pass checks the returned KMeans.Result.clusterSizes() itself, which is also how it recognizes the fallback. With sizePredicate == null every run passes, and the run with the lowest objective is returned.

      The size-balancing strength is given relative to the natural scale of the base objective of vectors: during the iterations, the penalty is weighted by lambda = relativeLambda * s, where s is the per-vector objective of the closed-form single-cluster fit (their mean squared L2 distance to their mean in Euclidean mode, their mean cosine distance to their normalized mean in cosine mode). A relative strength is thus measured in units of the data's own spread. That makes it exactly invariant under scaling the vectors, which scales both terms of the assignment score alike, and comparable in practice, though not identical, across metrics and numbers of dimensions. When all vectors are identical (in cosine mode: parallel), s == 0, and so is the weight: balancing is off. Computing s consumes no randomness, and is skipped for k == 1, for relativeLambda == 0 and for sizePenalty == null.

      The penalty acts only in the assignment steps of the iterations, where it steers the centroids: the final pass of each run ignores it, and so does the objective the runs are ranked by. With the penalty active (lambda > 0) the assignment step of an iteration is order-dependent (each vector's size-balance penalty is evaluated against the running projected sizes of all clusters assigned before it). To stop a fixed input ordering from biasing the centroids, fit then reshuffles the assignment order via Fisher–Yates at the top of every iteration. Otherwise, and in the final pass of every run, the assignment is order-independent and shuffling is skipped as wasted work. As the returned partitioning is the nearest-centroid one, a cluster can in rare cases come out empty.

      Type Parameters:
      V - caller's input vector representation
      C - caller's centroid representation
      Parameters:
      random - the random source; consumed by the k-means++ initialization of every run and by the per-iteration shuffling (when the penalty is active)
      distanceEstimator - the distance estimator. Its metric must be either EUCLIDEAN_METRIC or COSINE_METRIC; other metrics are rejected with an UnsupportedOperationException
      vectorLens - lens that extracts a RealVector from each input element
      centroidLens - lens that wraps a RealVector as the caller's centroid type
      vectors - the input points to cluster; must contain at least k elements
      k - number of clusters; must be at least 1. k == 1 is handled in closed form (the single centroid is the mean of all vectors), bypassing iteration, the runs and sizePredicate
      maxIterations - maximum Lloyd iterations per run; must be at least 1
      numRuns - total number of runs, each from its own k-means++ initialization; must be at least 1
      relativeLambda - strength of the size-balancing bias added to the assignment score of the iterations, relative to s (see above); must be finite and non-negative. 0 disables balancing
      sizePenalty - penalty function applied to projected cluster sizes during the assignment steps of the iterations; null disables balancing. Ignored when relativeLambda == 0
      sizePredicate - preference over the final cluster sizes of a run (see above), tested on a copy of each run's sizes; null lets every run pass
      Returns:
      the passing run with the lowest objective or, if no run passes sizePredicate, the run with the lowest objective