772 lines
25 KiB
Go
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
|
||
|
|
}
|