feat: initial release

Assisted-by: GLM 5.3 Flash

stats/gp.go

@@ -0,0 +1,398 @@
+// 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"
+)
+
+// Gaussian-process regression: a prior over functions fixed by a
+// covariance kernel, conditioned exactly on the observations through
+// the Cholesky factor of the Gram matrix, the same house route the
+// multivariate normal takes. The marginal log likelihood is exposed as
+// a plain function of the hyperparameters: the dependency graph gives
+// stats no edge to optim, so the fitting loop stays with the caller,
+// and the house minimiser (optim.Minimise) drives
+// MarginalLogLikelihood from outside the package over the two or three
+// hyperparameters a kernel carries.
+
+// Kernel is the covariance function of a Gaussian process: it returns
+// the prior covariance k(x, y) of the process at one pair of input
+// points, the points of equal dimension. Every house kernel carries a
+// unit amplitude: k(x, x) = 1, and the constructors below enforce the
+// parameter contract, a finite positive scale (and period where one
+// belongs), so a Kernel value is always safe to evaluate.
+type Kernel interface {
+	// Covariance returns k(x, y) for two points of equal length.
+	Covariance(x, y []float64) float64
+}
+
+// gpKernelScale validates one finite positive kernel parameter, the
+// contract every constructor states.
+func gpKernelScale(name, label string, v float64) error {
+	if !(v > 0) || math.IsInf(v, 0) {
+		return base.Errf("%s: %s must be finite and positive, got %g", name, label, v)
+	}
+	return nil
+}
+
+// squaredExponential is the smooth limit kernel
+// k(r) = exp(-r²/(2·lengthScale²)), infinitely differentiable, the
+// default prior for functions believed smooth.
+type squaredExponential struct {
+	lengthScale float64
+}
+
+// SquaredExponentialKernel returns the squared-exponential (RBF)
+// kernel of unit amplitude over the Euclidean distance of the input
+// points, with the given length scale.
+func SquaredExponentialKernel(lengthScale float64) (Kernel, error) {
+	const name = "SquaredExponentialKernel"
+	if err := gpKernelScale(name, "the length scale", lengthScale); err != nil {
+		return nil, err
+	}
+	return squaredExponential{lengthScale: lengthScale}, nil
+}
+
+// Covariance implements Kernel.
+func (k squaredExponential) Covariance(x, y []float64) float64 {
+	r2 := sqDistance(x, y)
+	return math.Exp(-r2 / (2 * k.lengthScale * k.lengthScale))
+}
+
+// matern32 is the once differentiable Matérn with ν = 3/2,
+// k(r) = (1 + √3·r/lengthScale)·exp(-√3·r/lengthScale), the spectral
+// density of which is S(ω) = 4a³/(a² + ω²)² for a = √3/lengthScale:
+// the closed form above and that spectral form are a Fourier pair,
+// k(r) = (1/π)∫₀^∞ S(ω)·cos(ω·r) dω, the defining equation the tests
+// hold the closed form against.
+type matern32 struct {
+	lengthScale float64
+}
+
+// Matern32Kernel returns the Matérn kernel with ν = 3/2 and the given
+// length scale, of unit amplitude over the Euclidean distance.
+func Matern32Kernel(lengthScale float64) (Kernel, error) {
+	const name = "Matern32Kernel"
+	if err := gpKernelScale(name, "the length scale", lengthScale); err != nil {
+		return nil, err
+	}
+	return matern32{lengthScale: lengthScale}, nil
+}
+
+// Covariance implements Kernel.
+func (k matern32) Covariance(x, y []float64) float64 {
+	a := math.Sqrt(3) * math.Sqrt(sqDistance(x, y)) / k.lengthScale
+	return (1 + a) * math.Exp(-a)
+}
+
+// matern52 is the twice differentiable Matérn with ν = 5/2,
+// k(r) = (1 + √5·r/lengthScale + 5r²/(3·lengthScale²))·exp(-√5·r/lengthScale),
+// the roughness the incidence of real fields usually lands between the
+// two lower Matérns and the smooth exponential square.
+type matern52 struct {
+	lengthScale float64
+}
+
+// Matern52Kernel returns the Matérn kernel with ν = 5/2 and the given
+// length scale, of unit amplitude over the Euclidean distance.
+func Matern52Kernel(lengthScale float64) (Kernel, error) {
+	const name = "Matern52Kernel"
+	if err := gpKernelScale(name, "the length scale", lengthScale); err != nil {
+		return nil, err
+	}
+	return matern52{lengthScale: lengthScale}, nil
+}
+
+// Covariance implements Kernel.
+func (k matern52) Covariance(x, y []float64) float64 {
+	r := math.Sqrt(sqDistance(x, y))
+	a := math.Sqrt(5) * r / k.lengthScale
+	return (1 + a + a*a/3) * math.Exp(-a)
+}
+
+// periodic is the periodic kernel
+// k(r) = exp(-2·sin²(π·r/period)/lengthScale²), the prior over
+// functions that repeat exactly with the given period, the length
+// scale setting how sharply neighbouring periods decorrelate.
+type periodic struct {
+	lengthScale float64
+	period      float64
+}
+
+// PeriodicKernel returns the periodic kernel with the given length
+// scale and period, of unit amplitude over the Euclidean distance.
+func PeriodicKernel(lengthScale, period float64) (Kernel, error) {
+	const name = "PeriodicKernel"
+	if err := gpKernelScale(name, "the length scale", lengthScale); err != nil {
+		return nil, err
+	}
+	if err := gpKernelScale(name, "the period", period); err != nil {
+		return nil, err
+	}
+	return periodic{lengthScale: lengthScale, period: period}, nil
+}
+
+// Covariance implements Kernel.
+func (k periodic) Covariance(x, y []float64) float64 {
+	r := math.Sqrt(sqDistance(x, y))
+	s := math.Sin(math.Pi * r / k.period)
+	return math.Exp(-2 * s * s / (k.lengthScale * k.lengthScale))
+}
+
+// GaussianProcessResult carries the posterior of a Gaussian process at
+// the test points.
+type GaussianProcessResult struct {
+	// Mean is the posterior mean function evaluated at the test
+	// points.
+	Mean []float64
+	// Covariance is the posterior covariance between the test
+	// points, m-by-m row-major in the order the test points were
+	// given, the full uncertainty the posterior carries.
+	Covariance []float64
+	// Variance is the diagonal of Covariance, the marginal posterior
+	// variance per test point, clamped at zero: rounding can push a
+	// training-point diagonal a hair below zero where the truth is
+	// zero, and a negative variance reports nothing.
+	Variance []float64
+	// LogLikelihood is the log marginal likelihood of the training
+	// observations under the prior, the objective a hyperparameter
+	// fit maximises (see MarginalLogLikelihood).
+	LogLikelihood float64
+}
+
+// GaussianProcessRegression conditions the prior defined by the
+// kernel on the n noisy observations y over the n training rows of
+// trainX (n rows, d columns), and evaluates the posterior mean and
+// covariance at the m rows of testX. The noise variance sits on the
+// Gram diagonal, K = k(X, X) + noiseVariance·I: zero is a legitimate
+// noiseless fit, and with it the posterior interpolates the training
+// data exactly and has zero variance there, while duplicated training
+// rows name the singular row in an error. The posterior algebra is the
+// Cholesky route: the Gram matrix factors through the multivariate
+// normal machinery of mvn.go, the weights alpha solve the triangular
+// systems against it, and the predictive covariance subtracts the
+// forward-solved cross covariances from the prior.
+//
+// The inputs must be real and finite, the training and test designs
+// of equal width, y of the training length, and the noise variance
+// finite and non-negative.
+func GaussianProcessRegression(kernel Kernel, trainX, trainY *core.Array, noiseVariance float64, testX *core.Array) (*GaussianProcessResult, error) {
+	const name = "GaussianProcessRegression"
+	l, alpha, train, n, d, err := gpFit(name, kernel, trainX, trainY, noiseVariance)
+	if err != nil {
+		return nil, err
+	}
+	test, m, testD, err := gpReadDesign(name, "the test design", testX)
+	if err != nil {
+		return nil, err
+	}
+	if testD != d {
+		return nil, base.Errf("%s: the training design is %d wide but the test design %d", name, d, testD)
+	}
+	out := &GaussianProcessResult{
+		Mean:       make([]float64, m),
+		Covariance: make([]float64, m*m),
+		Variance:   make([]float64, m),
+	}
+	// The cross covariances A: A[i][j] = k(test i, train j), and the
+	// forward-solved V = L⁻¹Aᵀ whose rows pair off against each other
+	// in the predictive covariance.
+	cross := make([][]float64, m)
+	solved := make([][]float64, m)
+	for i := range m {
+		cross[i] = make([]float64, n)
+		for j := range n {
+			cross[i][j] = kernel.Covariance(test[i*d:i*d+d], train[j*d:j*d+d])
+		}
+		solved[i] = gpForwardSolve(l, cross[i])
+	}
+	for i := range m {
+		total := 0.0
+		for j := range n {
+			total += cross[i][j] * alpha[j]
+		}
+		out.Mean[i] = total
+		for j := i; j < m; j++ {
+			// The prior term is the kernel's own diagonal entry,
+			// k(x*, x*), which the unit-amplitude house kernels fix at
+			// 1 but a custom kernel may set otherwise.
+			prior := kernel.Covariance(test[j*d:j*d+d], test[j*d:j*d+d])
+			if j > i {
+				prior = kernel.Covariance(test[i*d:i*d+d], test[j*d:j*d+d])
+			}
+			cov := prior
+			for r := range n {
+				cov -= solved[i][r] * solved[j][r]
+			}
+			out.Covariance[i*m+j] = cov
+			out.Covariance[j*m+i] = cov
+		}
+		diag := out.Covariance[i*m+i]
+		if diag < 0 {
+			diag = 0
+		}
+		out.Variance[i] = diag
+	}
+	out.LogLikelihood = gpLogLikelihood(l, alpha, trainY, n)
+	return out, nil
+}
+
+// MarginalLogLikelihood returns the log marginal likelihood of the
+// observations y under the prior the kernel defines over the training
+// design: the evidence of the hyperparameters, integrating the
+// training responses against their multivariate normal prior. It is
+// the objective a hyperparameter fit maximises; the dependency graph
+// gives stats no edge to optim, so the house minimiser drives this
+// function from outside the package, over the kernel's scale (and
+// period) and the noise variance.
+func MarginalLogLikelihood(kernel Kernel, trainX, trainY *core.Array, noiseVariance float64) (float64, error) {
+	const name = "MarginalLogLikelihood"
+	l, alpha, _, n, _, err := gpFit(name, kernel, trainX, trainY, noiseVariance)
+	if err != nil {
+		return 0, err
+	}
+	return gpLogLikelihood(l, alpha, trainY, n), nil
+}
+
+// gpFit validates the training inputs and returns the Cholesky factor
+// of the regularised Gram matrix, the solved weights alpha and the
+// flattened training design.
+func gpFit(name string, kernel Kernel, trainX, trainY *core.Array, noiseVariance float64) ([][]float64, []float64, []float64, int, int, error) {
+	if kernel == nil {
+		return nil, nil, nil, 0, 0, base.Errf("%s: the kernel is nil", name)
+	}
+	if math.IsNaN(noiseVariance) || math.IsInf(noiseVariance, 0) || noiseVariance < 0 {
+		return nil, nil, nil, 0, 0, base.Errf("%s: the noise variance must be finite and non-negative, got %g", name, noiseVariance)
+	}
+	train, n, d, err := gpReadDesign(name, "the training design", trainX)
+	if err != nil {
+		return nil, nil, nil, 0, 0, err
+	}
+	if err := gpReadResponse(name, trainY, n); err != nil {
+		return nil, nil, nil, 0, 0, err
+	}
+	// The Gram matrix, computed on the lower triangle and mirrored
+	// exactly, so the factorisation's symmetry check reads two
+	// bit-identical halves. It is written straight into one flat
+	// row-major slice, the layout the factorisation reads, instead of
+	// through a row-of-rows form another pass would flatten.
+	gram := make([]float64, n*n)
+	for i := range n {
+		for j := range i + 1 {
+			v := kernel.Covariance(train[i*d:i*d+d], train[j*d:j*d+d])
+			if i == j {
+				v += noiseVariance
+			}
+			gram[i*n+j] = v
+			gram[j*n+i] = v
+		}
+	}
+	l, err := mvnCholeskyFlat(name, gram, n)
+	if err != nil {
+		return nil, nil, nil, 0, 0, base.Errf("%s: %w", name, err)
+	}
+	y := make([]float64, n)
+	fy := rawFloats(trainY)
+	for i := range n {
+		if fy != nil {
+			y[i] = fy[i]
+		} else {
+			y[i] = trainY.FloatAt(i)
+		}
+	}
+	alpha := gpBackSolve(l, gpForwardSolve(l, y))
+	return l, alpha, train, n, d, nil
+}
+
+// gpLogLikelihood assembles the log marginal likelihood from the
+// factor and the solved weights:
+// -½·yᵀ·alpha - Σ ln Lᵢᵢ - (n/2)·ln 2π.
+func gpLogLikelihood(l [][]float64, alpha []float64, trainY *core.Array, n int) float64 {
+	fy := rawFloats(trainY)
+	total := 0.0
+	for i := range n {
+		var yv float64
+		if fy != nil {
+			yv = fy[i]
+		} else {
+			yv = trainY.FloatAt(i)
+		}
+		total -= 0.5 * yv * alpha[i]
+		total -= math.Log(l[i][i])
+	}
+	return total - 0.5*float64(n)*math.Log(2*math.Pi)
+}
+
+// gpReadDesign validates a finite real rank-2 design and returns it
+// flattened row-major.
+func gpReadDesign(name, label string, x *core.Array) ([]float64, int, int, error) {
+	if x.NDim() != 2 {
+		return nil, 0, 0, base.Errf("%s: %s must be rank 2, got shape %s", name, label, 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, label, 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: %s needs at least one column", name, label)
+	}
+	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
+}
+
+// gpReadResponse validates the response vector against the training
+// length.
+func gpReadResponse(name string, y *core.Array, n int) error {
+	if y.NDim() != 1 {
+		return base.Errf("%s: the response must be rank 1", name)
+	}
+	if y.Dtype() == core.Complex {
+		return base.Errf("%s: complex inputs are not supported", name)
+	}
+	if y.Len() != n {
+		return base.Errf("%s: the design has %d rows but the response %d", name, n, y.Len())
+	}
+	return checkFinite(name, "the response", y)
+}
+
+// gpForwardSolve solves L·v = b for the lower triangular L.
+func gpForwardSolve(l [][]float64, b []float64) []float64 {
+	v := make([]float64, len(b))
+	for i := range b {
+		total := b[i]
+		for j := range i {
+			total -= l[i][j] * v[j]
+		}
+		v[i] = total / l[i][i]
+	}
+	return v
+}
+
+// gpBackSolve solves Lᵀ·u = b for the lower triangular L.
+func gpBackSolve(l [][]float64, b []float64) []float64 {
+	n := len(b)
+	u := make([]float64, n)
+	for i := n - 1; i >= 0; i-- {
+		total := b[i]
+		for j := i + 1; j < n; j++ {
+			total -= l[j][i] * u[j]
+		}
+		u[i] = total / l[i][i]
+	}
+	return u
+}