feat: initial release
Assisted-by: GLM 5.3 Flash
This commit is contained in:
@@ -0,0 +1,771 @@
|
||||
// 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
|
||||
}
|
||||
Reference in New Issue
Block a user