Files
tensor/stats/cluster.go
T
petrbalvin af4ee19703
Release / gates (push) Successful in 4m38s
Test / test (push) Successful in 5m16s
Release / release (push) Successful in 35s
feat: initial release
Assisted-by: GLM 5.3 Flash
2026-09-03 10:00:00 +02:00

772 lines
25 KiB
Go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT
package stats
import (
"math"
"sourcedock.dev/petrbalvin/tensor/internal/base"
"sourcedock.dev/petrbalvin/tensor/internal/core"
"sourcedock.dev/petrbalvin/tensor/internal/engine"
)
// Clustering: k-means with k-means++ seeding and the Gaussian mixture
// fitted by expectation maximisation over the multivariate normal
// machinery of mvn.go. Both fits are deterministic for a given
// generator state: every random choice runs through the house
// generator, never through a global source.
// kmeansMaxIterations caps the Lloyd sweeps; a fit that has not
// settled by then is reported with Converged false rather than
// pretended to have converged.
const kmeansMaxIterations = 300
// kmeansTolerance is the convergence tolerance of the Lloyd loop: the
// sweep stops once the largest coordinate movement of any centre is at
// most kmeansTolerance scaled by the largest absolute coordinate in
// the data (floored at 1, so a degenerate all-zero cloud still
// converges on an absolute scale). The rule mirrors the fixed points
// of Lloyd's iteration: past the tolerance the objective moves below
// rounding, and the label assignment no longer changes.
const kmeansTolerance = 1e-10
// KMeansResult carries the fit of k-means over a sample.
type KMeansResult struct {
// Centres are the k fitted centroids, one row of d coordinates
// each.
Centres [][]float64
// Labels assigns every sample row its cluster, 0 to k-1.
Labels []int
// Inertia is the within-cluster sum of squared distances to the
// centres, the objective the Lloyd loop minimises.
Inertia float64
// Iterations counts the Lloyd sweeps taken; Converged reports
// whether the centre movement fell under kmeansTolerance before
// the iteration budget ran out.
Iterations int
Converged bool
}
// KMeans partitions the n rows of x (n rows, d columns) into k
// clusters by Lloyd's iteration, seeded with k-means++ over the house
// generator: the first centre is drawn uniformly among the samples and
// every later centre with probability proportional to the squared
// distance to the nearest centre already chosen, so a well-separated
// cluster cannot be left unseeded except against odds its separation
// sets. The fit is deterministic for a given generator state.
//
// Convergence: the sweeps stop once no centre moves by more than
// kmeansTolerance (scaled, as the constant documents) or after
// kmeansMaxIterations sweeps; Converged names which happened.
//
// Empty clusters: a sweep can strand a centre with no members. The
// stranded centre then moves onto the sample farthest from its
// current centre, the sample that the objective can improve on most;
// if several centres strand in one sweep, each takes the next
// farthest unused sample. The rule is deterministic and keeps every
// cluster live.
//
// A nil generator, a non-finite sample, k below 1 or k above n is an
// error.
func KMeans(g *core.Generator, x *core.Array, k int) (*KMeansResult, error) {
const name = "KMeans"
if g == nil {
return nil, base.Errf("%s: the generator is nil", name)
}
data, n, d, err := clusterReadSample(name, x)
if err != nil {
return nil, err
}
if k < 1 {
return nil, base.Errf("%s: k must be at least 1, got %d", name, k)
}
if k > n {
return nil, base.Errf("%s: k = %d exceeds the %d samples", name, k, n)
}
seeds := kmeansSeeds(g, data, n, d, k)
return kmeansLloyd(data, n, d, k, seeds)
}
// kmeansSeeds draws k initial centres by the k-means++ rule: the
// first uniformly, each later one with probability proportional to the
// squared distance to the nearest centre already drawn. The weighted
// draw walks the cumulative weights against one uniform, the same
// inverse-transform the package's discrete draws use.
func kmeansSeeds(g *core.Generator, data []float64, n, d, k int) [][]float64 {
seeds := make([][]float64, k)
first := min(int(g.Unit()*float64(n)), n-1)
seeds[0] = append([]float64(nil), data[first*d:first*d+d]...)
nearest := make([]float64, n)
for i := range n {
nearest[i] = sqDistance(data[i*d:i*d+d], seeds[0])
}
for c := 1; c < k; c++ {
total := 0.0
for i := range n {
total += nearest[i]
}
// A total of zero means every sample sits exactly on a chosen
// centre: nothing left to separate, and any further centre is
// arbitrary. The position repeats, and the Lloyd loop's
// empty-cluster rule resolves the tie.
if total == 0 {
seeds[c] = append([]float64(nil), data[0:d]...)
} else {
target := g.Unit() * total
walk := 0.0
pick := n - 1
for i := range n {
walk += nearest[i]
if walk >= target {
pick = i
break
}
}
seeds[c] = append([]float64(nil), data[pick*d:pick*d+d]...)
}
for i := range n {
if dist := sqDistance(data[i*d:i*d+d], seeds[c]); dist < nearest[i] {
nearest[i] = dist
}
}
}
return seeds
}
// kmeansChosen reports whether the point equals one of the centres
// already drawn, the guard of the degenerate all-coincident fallback.
func kmeansChosen(centres [][]float64, point []float64) bool {
for _, c := range centres {
same := true
for j, v := range c {
if v != point[j] {
same = false
break
}
}
if same {
return true
}
}
return false
}
// sqDistance returns the squared Euclidean distance between two
// points of equal length.
func sqDistance(a, b []float64) float64 {
total := 0.0
for j, v := range a {
diff := v - b[j]
total += diff * diff
}
return total
}
// kmeansLloyd runs the assignment and update sweeps from the given
// centres to the documented tolerance, applying the empty-cluster
// rule on the way. It is the shared engine of the public entry point
// and of the seeding comparisons the tests make.
func kmeansLloyd(data []float64, n, d, k int, centres [][]float64) (*KMeansResult, error) {
// The scale that makes the tolerance relative: the largest
// absolute coordinate in the sample, floored at 1.
spread := 1.0
for _, v := range data {
if s := math.Abs(v); s > spread {
spread = s
}
}
labels := make([]int, n)
converged := false
iterations := kmeansMaxIterations
// The update sweep's accumulators live outside the iteration: a sweep
// clears them instead of allocating, and the fold below then adds the
// same rows into the same slots in the same order.
counts := make([]int, k)
sums := make([][]float64, k)
for c := range k {
sums[c] = make([]float64, d)
}
for iter := 1; iter <= kmeansMaxIterations; iter++ {
for i := range n {
best, bestDist := 0, math.Inf(1)
for c := range k {
if dist := sqDistance(data[i*d:i*d+d], centres[c]); dist < bestDist {
best, bestDist = c, dist
}
}
labels[i] = best
}
// The update sweep, with the empty-cluster rule: a centre
// whose cluster stranded takes the sample farthest from its
// own centre, skipping the samples earlier stranded centres
// in the same sweep already claimed.
clear(counts)
for c := range k {
clear(sums[c])
}
for i := range n {
counts[labels[i]]++
row := sums[labels[i]]
for j := range d {
row[j] += data[i*d+j]
}
}
worst := 0.0
for c := range k {
switch {
case counts[c] > 0:
for j := range d {
moved := sums[c][j] / float64(counts[c])
if m := math.Abs(moved - centres[c][j]); m > worst {
worst = m
}
centres[c][j] = moved
}
default:
// The farthest sample from this centre: the search is
// for a maximum, so the running best starts at -Inf.
farthest, farDist := -1, math.Inf(-1)
for i := range n {
if labels[i] == c || kmeansChosen(centres, data[i*d:i*d+d]) {
continue
}
if dist := sqDistance(data[i*d:i*d+d], centres[c]); dist > farDist {
farthest, farDist = i, dist
}
}
if farthest < 0 {
return nil, base.Errf("KMeans: cluster %d cannot be repopulated: the sample holds fewer distinct points than k", c)
}
copy(centres[c], data[farthest*d:farthest*d+d])
worst = math.Inf(1)
}
}
if worst <= kmeansTolerance*spread {
converged = true
iterations = iter
break
}
}
return &KMeansResult{
Centres: centres,
Labels: labels,
Inertia: kmeansInertia(data, n, d, centres, labels),
Iterations: iterations,
Converged: converged,
}, nil
}
// kmeansInertia sums the squared distances of every sample to its
// assigned centre.
func kmeansInertia(data []float64, n, d int, centres [][]float64, labels []int) float64 {
total := 0.0
for i := range n {
total += sqDistance(data[i*d:i*d+d], centres[labels[i]])
}
return total
}
// gmmMaxIterations caps the EM sweeps, with the same reporting
// convention kmeansMaxIterations sets.
const gmmMaxIterations = 200
// gmmTolerance is the EM convergence tolerance on the relative change
// of the log likelihood: the sweeps stop once it moves by less than
// the tolerance scaled by 1 + |log likelihood|, past which no
// parameter the M step can produce moves the objective materially.
const gmmTolerance = 1e-9
// gmmCovarianceFloor is the regularisation floor of the M step: after
// every update, gmmCovarianceFloor of the data's mean per-dimension
// variance is added to the covariance diagonal. A component that
// collapses onto a single sample computes a singular covariance and a
// Cholesky factor that does not exist; the floor keeps the factor
// alive at a width far below any sample's spread, so it regularises
// the algebra without moving any fit the data actually supports.
const gmmCovarianceFloor = 1e-6
// GaussianMixtureResult carries the fit of a Gaussian mixture.
type GaussianMixtureResult struct {
// Components is the fitted component count.
Components int
// Weights are the mixture weights, in component order, summing to
// 1.
Weights []float64
// Means holds one mean vector per component.
Means [][]float64
// Covariances holds one d-by-d covariance per component,
// row-major, in component order.
Covariances [][]float64
// Responsibilities are the final E-step posteriors, n rows of
// Components entries each, row-major: P(component j | sample i).
Responsibilities []float64
// LogLikelihood is the maximised observed-data log likelihood.
LogLikelihood float64
// BIC is -2·LogLikelihood + p·ln n at the fitted parameters,
// p the count of free parameters.
BIC float64
// Iterations counts the EM sweeps; Converged reports whether the
// log likelihood settled under gmmTolerance before the budget.
Iterations int
Converged bool
// BICGrid is filled by GaussianMixtureBIC only: the BIC of every
// fit on the component grid 1..len(BICGrid), in grid order. The
// plain fit leaves it nil.
BICGrid []float64
}
// GaussianMixture fits a mixture of `components` multivariate normals
// over the rows of x by expectation maximisation, built on the MVN
// machinery of mvn.go: every component density goes through the
// Cholesky factor of its covariance.
//
// The E step runs in log space: each sample's responsibilities come
// from a log-sum-exp normalisation guarded by the row maximum, so a
// sample sitting far outside every component underflows to a clean 0/1
// split rather than to a NaN.
//
// The M step regularises every covariance by the documented
// gmmCovarianceFloor before the next factorisation.
//
// Initialisation is k-means over the same generator: the partitions
// seed the weights, means and covariances, so the fit is deterministic
// for a given seed and every component starts populated. A component
// whose total responsibility collapses below 1e-10 (the floor at work
// on a stray sample) keeps its previous parameters and a floored
// weight rather than dividing by zero; the renormalised weights keep
// the mixture a mixture.
//
// The sample must be rank 2 and finite, the component count at least 1
// and at most the sample size, and the sample at least two rows.
func GaussianMixture(g *core.Generator, x *core.Array, components int) (*GaussianMixtureResult, error) {
const name = "GaussianMixture"
if g == nil {
return nil, base.Errf("%s: the generator is nil", name)
}
data, n, d, err := clusterReadSample(name, x)
if err != nil {
return nil, err
}
if n < 2 {
return nil, base.Errf("%s: needs at least two samples, got %d", name, n)
}
if components < 1 {
return nil, base.Errf("%s: the component count must be at least 1, got %d", name, components)
}
if components > n {
return nil, base.Errf("%s: %d components exceed the %d samples", name, components, n)
}
weights, means, covs, serr := gmmSeed(g, data, n, d, components)
if serr != nil {
return nil, base.Errf("%s: %w", name, serr)
}
result, err := gmmEM(data, n, d, components, weights, means, covs)
if err != nil {
return nil, err
}
result.BIC = gaussianMixtureBIC(result.LogLikelihood, components, n, d)
return result, nil
}
// GaussianMixtureBIC fits mixtures over the component grid 1 to
// maxComponents and keeps the one with the lowest BIC. The grid and
// the winner travel in the result: BICGrid holds every fit's BIC in
// grid order and Components names the selected count. Ties resolve to
// the smaller model, the parsimony the penalty exists to buy.
func GaussianMixtureBIC(g *core.Generator, x *core.Array, maxComponents int) (*GaussianMixtureResult, error) {
const name = "GaussianMixtureBIC"
if g == nil {
return nil, base.Errf("%s: the generator is nil", name)
}
data, n, d, err := clusterReadSample(name, x)
if err != nil {
return nil, err
}
if n < 2 {
return nil, base.Errf("%s: needs at least two samples, got %d", name, n)
}
if maxComponents < 1 {
return nil, base.Errf("%s: the largest component count must be at least 1, got %d", name, maxComponents)
}
if maxComponents > n {
return nil, base.Errf("%s: %d components exceed the %d samples", name, maxComponents, n)
}
grid := make([]float64, maxComponents)
best := -1
bestBIC := math.Inf(1)
var bestResult *GaussianMixtureResult
for k := 1; k <= maxComponents; k++ {
weights, means, covs, serr := gmmSeed(g, data, n, d, k)
if serr != nil {
return nil, serr
}
result, err := gmmEM(data, n, d, k, weights, means, covs)
if err != nil {
return nil, base.Errf("%s: the fit with %d components failed (%w)", name, k, err)
}
grid[k-1] = gaussianMixtureBIC(result.LogLikelihood, k, n, d)
// Strictly better only: a tie keeps the smaller model.
if grid[k-1] < bestBIC {
bestBIC = grid[k-1]
best = k
bestResult = result
}
}
bestResult.BICGrid = grid
bestResult.BIC = grid[best-1]
return bestResult, nil
}
// gaussianMixtureBIC assembles -2·logL + p·ln n, with p the free
// parameters of a k-component mixture: k-1 weights, k·d mean entries
// and k·d(d+1)/2 covariance entries.
func gaussianMixtureBIC(logL float64, k, n, d int) float64 {
p := (k - 1) + k*d + k*d*(d+1)/2
return -2*logL + float64(p)*math.Log(float64(n))
}
// clusterReadSample validates x as a finite real rank-2 sample and
// returns it flattened row-major with its shape.
func clusterReadSample(name string, x *core.Array) ([]float64, int, int, error) {
if x.NDim() != 2 {
return nil, 0, 0, base.Errf("%s: the sample must be rank 2, got shape %s", name, base.ShapeText(x.Shape()))
}
if x.Dtype() == core.Complex {
return nil, 0, 0, base.Errf("%s: complex inputs are not supported", name)
}
if err := checkFinite(name, "the sample", x); err != nil {
return nil, 0, 0, err
}
n, d := x.Shape()[0], x.Shape()[1]
if d < 1 {
return nil, 0, 0, base.Errf("%s: the sample needs at least one column", name)
}
data := make([]float64, n*d)
if fs := rawFloats(x); fs != nil {
copy(data, fs)
} else {
for i := range data {
data[i] = x.FloatAt(i)
}
}
return data, n, d, nil
}
// gmmSeed initialises the mixture from the k-means partitions of the
// same generator: each component takes its cluster's weight, mean and
// covariance, floored exactly as the M step floors, so the first E
// step factors live covariances even for a cluster that collapsed onto
// one sample.
func gmmSeed(g *core.Generator, data []float64, n, d, k int) ([]float64, [][]float64, [][]float64, error) {
seeds := kmeansSeeds(g, data, n, d, k)
fit, err := kmeansLloyd(data, n, d, k, seeds)
if err != nil {
// A sample with fewer distinct points than k drives the
// empty-cluster rule out of answers: the caller turns this
// into the entry point's own named refusal.
return nil, nil, nil, err
}
meanVariance := clusterMeanVariance(data, n, d)
weights := make([]float64, k)
means := make([][]float64, k)
covs := make([][]float64, k)
counts := make([]int, k)
for i := range n {
counts[fit.Labels[i]]++
}
for c := range k {
weights[c] = math.Max(float64(counts[c]), 1e-10)
mean := make([]float64, d)
if counts[c] > 0 {
for i := range n {
if fit.Labels[i] != c {
continue
}
for j := range d {
mean[j] += data[i*d+j]
}
}
for j := range d {
mean[j] /= float64(counts[c])
}
}
cov := make([]float64, d*d)
for i := range n {
if fit.Labels[i] != c {
continue
}
for a := range d {
for b := range d {
cov[a*d+b] += (data[i*d+a] - mean[a]) * (data[i*d+b] - mean[b])
}
}
}
den := math.Max(float64(counts[c])-1, 1)
for a := range d {
for b := range d {
cov[a*d+b] /= den
}
cov[a*d+a] += gmmCovarianceFloor * meanVariance
}
weights[c] /= float64(n)
means[c] = mean
covs[c] = cov
}
// The floored weights renormalised: the mixture must sum to 1 even
// when a component started empty and took the floor.
total := 0.0
for _, w := range weights {
total += w
}
for c := range k {
weights[c] /= total
}
return weights, means, covs, nil
}
// clusterMeanVariance returns the mean per-dimension variance of the
// sample, the scale the covariance floor is measured in.
func clusterMeanVariance(data []float64, n, d int) float64 {
if n < 2 {
return 1
}
total := 0.0
for j := range d {
mean := 0.0
for i := range n {
mean += data[i*d+j]
}
mean /= float64(n)
v := 0.0
for i := range n {
diff := data[i*d+j] - mean
v += diff * diff
}
total += v / float64(n-1)
}
return total / float64(d)
}
// gmmParallelMinPoints is the sample count one worker must carry before
// the expectation sweep splits across goroutines: a point costs the
// component count in density evaluations and another in exponentials,
// so a shorter chunk is cheaper on the calling goroutine than in a
// pool.
const gmmParallelMinPoints = 32
// gmmEM runs the expectation maximisation sweeps from the seeded
// parameters to the documented tolerance.
func gmmEM(data []float64, n, d, k int, weights []float64, means [][]float64, covs [][]float64) (*GaussianMixtureResult, error) {
const name = "GaussianMixture"
resp := make([]float64, n*k)
logL := math.Inf(-1)
converged := false
iterations := gmmMaxIterations
// The sweep's buffers live outside the sweep: one flat covariance
// copy the factor step refills, one factor per component the
// factorisation clears and refills, and one row of per-point
// normalisers the log-likelihood fold reads, so a sweep allocates
// none of them. The per-component mean and covariance slices the M
// step accumulates into are the parameter storage itself: each is
// cleared before its pass, so the accumulation sees the zero state
// a fresh allocation carried and the result carries the last
// sweep's values in the same slices.
covVals := make([]float64, d*d)
factors := make([][][]float64, k)
for c := range k {
rows := make([][]float64, d)
for i := range d {
rows[i] = make([]float64, d)
}
factors[c] = rows
if len(means[c]) != d {
means[c] = make([]float64, d)
}
if len(covs[c]) != d*d {
covs[c] = make([]float64, d*d)
}
}
consts := make([]float64, k)
logWeights := make([]float64, k)
counts := make([]float64, k)
centred := make([]float64, d)
logNorms := make([]float64, n)
// The covariance floor's scale depends on the sample alone, and the
// sample never moves, so it is measured once for the whole fit.
meanVariance := clusterMeanVariance(data, n, d)
for iter := 1; iter <= gmmMaxIterations; iter++ {
// E step: the components are factored once per sweep, and each
// factor contributes its log determinant and the density's
// normalising constant as one value, read once per component
// instead of once per point. The component's log weight is
// likewise constant across the sweep. The sample loop that
// follows then carries only the quadratic form.
for c := range k {
copy(covVals, covs[c])
if err := mvnCholeskyFlatInto(name, covVals, factors[c], d); err != nil {
return nil, base.Errf("%s: component %d failed to factor (%w)", name, c, err)
}
l := factors[c]
logDet := 0.0
for i := range d {
logDet += math.Log(l[i][i])
}
consts[c] = -0.5*float64(d)*math.Log(2*math.Pi) - logDet
logWeights[c] = math.Log(weights[c])
}
// A point owns its responsibility row and its own normaliser
// entry and nothing else, so the sample loop splits across
// goroutines with no lock and no shared scratch: the worker's two
// buffers are overwritten per point.
engine.ParallelMin(n, gmmParallelMinPoints, func(start, end int) {
logps := make([]float64, k)
solve := make([]float64, d)
for i := start; i < end; i++ {
point := data[i*d : i*d+d]
for c := range k {
logps[c] = logWeights[c] + mvnLogDensitySolve(point, means[c], factors[c], d, consts[c], solve)
}
rowMax := math.Inf(-1)
for _, lp := range logps {
if lp > rowMax {
rowMax = lp
}
}
// A floored weight can park a component at -Inf; the
// max-guarded sum survives it and the row still
// normalises. logps is reused for the exponentials once
// the row maximum has served: the posterior is the
// quotient of the stored value and the row total, so
// one exponential per component serves the row and the
// divisor is the total in [1, k] instead of an exponent
// compounded with the logarithm's rounding.
total := 0.0
for c := range k {
e := math.Exp(logps[c] - rowMax)
logps[c] = e
total += e
}
logNorm := rowMax + math.Log(total)
for c := range k {
resp[i*k+c] = logps[c] / total
}
logNorms[i] = logNorm
}
})
// The observed-data log likelihood folds the normalisers in
// ascending point order on the calling goroutine: the chain is
// the one the serial sweep accumulated, so the split above never
// moves a bit of it.
nextLogL := 0.0
for _, logNorm := range logNorms {
nextLogL += logNorm
}
move := nextLogL - logL
logL = nextLogL
if math.Abs(move) <= gmmTolerance*(1+math.Abs(logL)) {
converged = true
iterations = iter
break
}
// M step: responsibilities to weights, means and floored
// covariances. A component whose total responsibility
// collapses under 1e-10 keeps its previous parameters under a
// floored weight: dividing by zero would poison the sweep,
// and the floor lets the next E step give the component
// another chance.
clear(counts)
for i := range n {
for c := range k {
counts[c] += resp[i*k+c]
}
}
totalWeight := 0.0
for c := range k {
if counts[c] < 1e-10 {
weights[c] = 1e-10
} else {
weights[c] = counts[c] / float64(n)
}
totalWeight += weights[c]
}
for c := range k {
weights[c] /= totalWeight
}
for c := range k {
if counts[c] < 1e-10 {
continue
}
mean := means[c]
clear(mean)
for i := range n {
r := resp[i*k+c]
for j := range d {
mean[j] += r * data[i*d+j]
}
}
for j := range d {
mean[j] /= counts[c]
}
cov := covs[c]
clear(cov)
// The centred row is built once per sample and read by every
// entry the covariance accumulates: the subtraction is the one
// the product below performed per entry, and the remaining
// factorisation of the term, (r·ca)·cb added to the entry,
// keeps its own order and shape.
for i := range n {
r := resp[i*k+c]
point := data[i*d : i*d+d]
for a := range d {
centred[a] = point[a] - mean[a]
}
for a := range d {
scaled := r * centred[a]
row := cov[a*d : a*d+d]
for b, cb := range centred {
row[b] += scaled * cb
}
}
}
for a := range d {
for b := range d {
cov[a*d+b] /= counts[c]
}
cov[a*d+a] += gmmCovarianceFloor * meanVariance
}
}
}
return &GaussianMixtureResult{
Components: k,
Weights: weights,
Means: means,
Covariances: covs,
Responsibilities: resp,
LogLikelihood: logL,
Iterations: iterations,
Converged: converged,
}, nil
}
// mvnLogDensitySolve evaluates the log density of one point under a
// component already factored, with the caller's normalising constant
// (the log determinant already folded in) and its own forward-solve
// scratch. Every entry of solve is written before it is read, so the
// scratch carries no state across calls; the constant is the identical
// expression the caller would otherwise recompute from the same factor.
func mvnLogDensitySolve(point, mean []float64, l [][]float64, d int, normConst float64, solve []float64) float64 {
for i := range d {
total := point[i] - mean[i]
for j := range i {
total -= l[i][j] * solve[j]
}
solve[i] = total / l[i][i]
}
quad := 0.0
for i := range d {
quad += solve[i] * solve[i]
}
return normConst - 0.5*quad
}