feat: initial release
Assisted-by: GLM 5.3 Flash
This commit is contained in:
+509
@@ -0,0 +1,509 @@
|
||||
// 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"
|
||||
)
|
||||
|
||||
// Regularised linear regression by coordinate descent: the lasso and
|
||||
// its elastic net generalisation. The fit runs on the standardised
|
||||
// design (every column centred and scaled to unit population standard
|
||||
// deviation, the response centred), the penalty falls on the slopes
|
||||
// and never on the intercept, and the standardisation is recorded on
|
||||
// the result and inverted before the coefficients leave, so the
|
||||
// numbers a caller receives apply to the design as supplied. The path
|
||||
// over a documented lambda grid is warm-started: every fit after the
|
||||
// first begins at the previous lambda's coefficients, which is why
|
||||
// the path costs a fraction of the cold fits it replaces.
|
||||
|
||||
// ElasticNetResult carries one regularised fit at a single lambda.
|
||||
type ElasticNetResult struct {
|
||||
// Intercept and Coefficients are the fit on the original scale of
|
||||
// the design as supplied: the standardisation the solver worked on
|
||||
// has been inverted. Coefficients holds one slope per design
|
||||
// column, in order.
|
||||
Intercept float64
|
||||
Coefficients []float64
|
||||
// Fitted and Residuals align with the rows of the design.
|
||||
Fitted []float64
|
||||
Residuals []float64
|
||||
// ColumnMeans and ColumnScales record the standardisation the fit
|
||||
// ran under: the solver saw (x − ColumnMeans)/ColumnScales, and
|
||||
// the lambda is measured against the centred response in the
|
||||
// response's own units.
|
||||
ColumnMeans []float64
|
||||
ColumnScales []float64
|
||||
// Lambda and Alpha are the penalty and the L2 mixing actually used.
|
||||
Lambda float64
|
||||
Alpha float64
|
||||
// Iterations counts full coordinate-descent cycles and Converged
|
||||
// reports whether no slope moved more than the tolerance in the
|
||||
// last cycle. An exhausted budget returns the fit found so far
|
||||
// rather than an error: coordinate descent on this convex
|
||||
// objective cannot diverge, so the iterate stays a usable answer
|
||||
// and the flag says how much to trust it.
|
||||
Iterations int
|
||||
Converged bool
|
||||
}
|
||||
|
||||
// LassoPathResult carries a warm-started regularisation path.
|
||||
type LassoPathResult struct {
|
||||
// Alpha is the mixing the path ran with.
|
||||
Alpha float64
|
||||
// Lambdas are the penalty values of the path in descending order.
|
||||
// The documented grid is lassoGridSteps values log-spaced from
|
||||
// lambdaMax down to lambdaMax/1000, where lambdaMax is the
|
||||
// smallest penalty at which the lasso keeps every slope at zero,
|
||||
// max_j |Σ_i x_ij (y_i − ȳ)|/n divided by alpha. For alpha = 0
|
||||
// the pure-lasso lambdaMax (alpha taken as 1) defines the same
|
||||
// grid, so ridge walks exactly the path the lasso would.
|
||||
Lambdas []float64
|
||||
// Intercepts and Coefficients hold one fit per lambda, on the
|
||||
// original scale, in the order of Lambdas.
|
||||
Intercepts []float64
|
||||
Coefficients [][]float64
|
||||
// Iterations counts the coordinate-descent cycles each lambda
|
||||
// needed. The warm start is what keeps them low along the path:
|
||||
// each fit but the first begins at the previous lambda's
|
||||
// coefficients, far closer to the answer than the zero start a
|
||||
// cold fit uses, and the counts fall accordingly.
|
||||
Iterations []int
|
||||
// Converged reports whether every fit on the path converged
|
||||
// within the iteration budget.
|
||||
Converged bool
|
||||
}
|
||||
|
||||
// The documented shape of the lambda grid and the documented stopping
|
||||
// rule of the coordinate descent: lassoGridSteps log-spaced values
|
||||
// covering three decades down from lambdaMax, and a fit is settled
|
||||
// once no slope moves more than lassoTolerance in a full cycle, with
|
||||
// lassoMaxIterations cycles the most any single fit may spend.
|
||||
const (
|
||||
lassoGridSteps = 100
|
||||
lassoGridRatio = 1e-3
|
||||
lassoTolerance = 1e-8
|
||||
lassoMaxIter = 10000
|
||||
)
|
||||
|
||||
// Lasso fits y by linear regression under the least absolute
|
||||
// shrinkage penalty at a single lambda: the pure lasso, alpha = 1 of
|
||||
// ElasticNet. The slopes and only the slopes are penalised; the
|
||||
// intercept is never. See ElasticNet for the contract the two share.
|
||||
func Lasso(x, y *core.Array, lambda float64) (*ElasticNetResult, error) {
|
||||
return ElasticNet(x, y, lambda, 1)
|
||||
}
|
||||
|
||||
// ElasticNet fits y = X·β under the elastic net penalty
|
||||
//
|
||||
// (1/2n)·Σᵢ(yᵢ − β₀ − xᵢ·β)² + λ·(α·Σ|βⱼ| + (1−α)/2·Σβⱼ²),
|
||||
//
|
||||
// the convex combination of the lasso (alpha 1) and ridge (alpha 0)
|
||||
// penalties, by coordinate descent on the standardised design. The
|
||||
// design carries n rows and p columns exactly as LinearRegression's,
|
||||
// but a rank-deficient design is not an error here: the penalty keeps
|
||||
// the problem well posed, and more columns than rows is a legitimate
|
||||
// use. n must be at least 2, the design must carry at least one
|
||||
// non-constant column, alpha must lie in [0, 1] and lambda must not
|
||||
// be negative; complex or non-finite input is an error.
|
||||
//
|
||||
// Two exact anchors hold. On an orthonormal design the lasso answer
|
||||
// is the soft-thresholded projection, coefficient by coefficient, and
|
||||
// the fit reproduces that closed form. At alpha = 0 the answer is the
|
||||
// Tikhonov ridge system (ZᵀZ/n + λI)β = Zᵀ(y − ȳ)/n over the
|
||||
// standardised design, and the fit agrees with the answer the shared
|
||||
// LU solve delivers for that system.
|
||||
func ElasticNet(x, y *core.Array, lambda, alpha float64) (*ElasticNetResult, error) {
|
||||
const name = "ElasticNet"
|
||||
if math.IsNaN(lambda) || lambda < 0 {
|
||||
return nil, base.Errf("%s: lambda must not be negative, got %g", name, lambda)
|
||||
}
|
||||
if math.IsNaN(alpha) || alpha < 0 || alpha > 1 {
|
||||
return nil, base.Errf("%s: alpha must lie in [0, 1], got %g", name, alpha)
|
||||
}
|
||||
n, p, fx, fy, err := lassoInputs(name, x, y)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
std, err := lassoStandardise(name, n, p, x, y, fx, fy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
beta, iterations, converged := elasticNetFit(std, lambda, alpha, nil, &lassoWorkspace{})
|
||||
return lassoResult(name, n, p, x, y, fx, fy, std, beta, lambda, alpha, iterations, converged)
|
||||
}
|
||||
|
||||
// LassoPath fits the whole regularisation path over the documented
|
||||
// lambda grid, warm-started: the first fit, at the largest lambda,
|
||||
// starts from zero coefficients (which are the exact answer there for
|
||||
// the lasso), and every later fit starts at the previous lambda's
|
||||
// answer. alpha is the elastic net mixing as in ElasticNet, and the
|
||||
// alpha = 0 ridge walks the same grid the pure lasso defines. The
|
||||
// result records how many coordinate cycles each lambda spent, which
|
||||
// is the evidence the warm start earns its keep.
|
||||
func LassoPath(x, y *core.Array, alpha float64) (*LassoPathResult, error) {
|
||||
const name = "LassoPath"
|
||||
if math.IsNaN(alpha) || alpha < 0 || alpha > 1 {
|
||||
return nil, base.Errf("%s: alpha must lie in [0, 1], got %g", name, alpha)
|
||||
}
|
||||
n, p, fx, fy, err := lassoInputs(name, x, y)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
std, err := lassoStandardise(name, n, p, x, y, fx, fy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// lambdaMax: the smallest penalty the lasso answers with an
|
||||
// all-zero slope vector, divided by alpha because the elastic net
|
||||
// subgradient carries only the alpha fraction as the L1 term. The
|
||||
// ridge (alpha = 0) has no sparsity threshold to define a largest
|
||||
// lambda, so it borrows the pure-lasso grid and walks the same
|
||||
// path. A response orthogonal to every column gives lambdaMax 0:
|
||||
// no penalty does anything, and the grid value is immaterial, so a
|
||||
// unit grid keeps the documented shape.
|
||||
gridAlpha := alpha
|
||||
if gridAlpha == 0 {
|
||||
gridAlpha = 1
|
||||
}
|
||||
// The projection norms are computed with the same multiply by the
|
||||
// reciprocal the coordinate sweep uses, so at the top of the grid
|
||||
// the soft threshold compares identical bits against identical
|
||||
// bits and the all-zero answer is exact, not nearly zero.
|
||||
inv := 1 / float64(n)
|
||||
lambdaMax := 0.0
|
||||
for j := range p {
|
||||
col := std.xs[j]
|
||||
s := 0.0
|
||||
for i := range n {
|
||||
s += col[i] * std.yc[i]
|
||||
}
|
||||
if a := math.Abs(s*inv) / gridAlpha; a > lambdaMax {
|
||||
lambdaMax = a
|
||||
}
|
||||
}
|
||||
if lambdaMax == 0 {
|
||||
lambdaMax = 1
|
||||
}
|
||||
out := &LassoPathResult{
|
||||
Alpha: alpha,
|
||||
Lambdas: make([]float64, lassoGridSteps),
|
||||
Intercepts: make([]float64, lassoGridSteps),
|
||||
Coefficients: make([][]float64, lassoGridSteps),
|
||||
Iterations: make([]int, lassoGridSteps),
|
||||
Converged: true,
|
||||
}
|
||||
warm := make([]float64, p)
|
||||
ws := &lassoWorkspace{beta: make([]float64, p), r: make([]float64, n)}
|
||||
for k := range lassoGridSteps {
|
||||
lambda := lambdaMax * math.Pow(lassoGridRatio, float64(k)/float64(lassoGridSteps-1))
|
||||
beta, iterations, converged := elasticNetFit(std, lambda, alpha, warm, ws)
|
||||
copy(warm, beta)
|
||||
// The path records the fit itself and has no field for a
|
||||
// residual: lifting the standardised slopes back to the design
|
||||
// as supplied is the whole reading, so the per-lambda pass over
|
||||
// the original design that Fitted and Residuals would need is
|
||||
// not run.
|
||||
intercept, coefficients := lassoCoefficients(std, beta)
|
||||
out.Lambdas[k] = lambda
|
||||
out.Intercepts[k] = intercept
|
||||
out.Coefficients[k] = coefficients
|
||||
out.Iterations[k] = iterations
|
||||
out.Converged = out.Converged && converged
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// lassoStandardisation holds the standardised system the coordinate
|
||||
// descent runs on: standardised columns stored column-wise for cache
|
||||
// friendly sweeps, their squared norms, and the centred response.
|
||||
type lassoStandardisation struct {
|
||||
xs [][]float64 // the standardised columns, each of length n
|
||||
zs []float64 // z_j = Σᵢ x²ᵢⱼ/n, the standardised column energies
|
||||
yc []float64 // the centred response
|
||||
yMean float64 // the response mean the intercept reabsorbs
|
||||
means []float64 // the original column means
|
||||
scales []float64 // the original column population standard deviations
|
||||
}
|
||||
|
||||
// lassoWorkspace holds the coordinate descent's running buffers: the
|
||||
// slope vector and the residual the sweeps keep current. One
|
||||
// workspace serves a whole path, so a hundred warm-started fits
|
||||
// allocate none of them again.
|
||||
type lassoWorkspace struct {
|
||||
beta []float64
|
||||
r []float64
|
||||
}
|
||||
|
||||
// lassoInputs validates a design and response exactly as the rest of
|
||||
// the estimation entries do, and returns the row and column counts
|
||||
// with the dense raw payloads when they exist.
|
||||
func lassoInputs(name string, x, y *core.Array) (n, p int, fx, fy []float64, err error) {
|
||||
if x.NDim() != 2 {
|
||||
return 0, 0, nil, nil, base.Errf("%s: the design must be rank 2, got shape %s", name, base.ShapeText(x.Shape()))
|
||||
}
|
||||
if y.NDim() != 1 {
|
||||
return 0, 0, nil, nil, base.Errf("%s: the response must be rank 1", name)
|
||||
}
|
||||
if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
|
||||
return 0, 0, nil, nil, base.Errf("%s: complex inputs are not supported", name)
|
||||
}
|
||||
n, p = x.Shape()[0], x.Shape()[1]
|
||||
if y.Len() != n {
|
||||
return 0, 0, nil, nil, base.Errf("%s: the design has %d rows but the response %d", name, n, y.Len())
|
||||
}
|
||||
if n < 2 {
|
||||
return 0, 0, nil, nil, base.Errf("%s: at least two observations are needed, got %d", name, n)
|
||||
}
|
||||
if p == 0 {
|
||||
return 0, 0, nil, nil, base.Errf("%s: the design must carry at least one column", name)
|
||||
}
|
||||
if err := checkFinite(name, "the design", x); err != nil {
|
||||
return 0, 0, nil, nil, err
|
||||
}
|
||||
if err := checkFinite(name, "the response", y); err != nil {
|
||||
return 0, 0, nil, nil, err
|
||||
}
|
||||
if rawFloats(x) != nil {
|
||||
fx = rawFloats(x)
|
||||
}
|
||||
if rawFloats(y) != nil {
|
||||
fy = rawFloats(y)
|
||||
}
|
||||
return n, p, fx, fy, nil
|
||||
}
|
||||
|
||||
// lassoStandardise builds the standardised system: every column
|
||||
// centred and divided by its population standard deviation, the
|
||||
// response centred. A constant column has no standard deviation to
|
||||
// divide by and is refused: its slope is unidentified at any penalty.
|
||||
func lassoStandardise(name string, n, p int, x, y *core.Array, fx, fy []float64) (*lassoStandardisation, error) {
|
||||
std := &lassoStandardisation{
|
||||
xs: make([][]float64, p),
|
||||
zs: make([]float64, p),
|
||||
yc: make([]float64, n),
|
||||
means: make([]float64, p),
|
||||
scales: make([]float64, p),
|
||||
}
|
||||
inv := 1 / float64(n)
|
||||
for j := range p {
|
||||
mean := 0.0
|
||||
if fx != nil {
|
||||
for i := range n {
|
||||
mean += fx[i*p+j]
|
||||
}
|
||||
} else {
|
||||
for i := range n {
|
||||
mean += x.FloatAt(i*p + j)
|
||||
}
|
||||
}
|
||||
mean *= inv
|
||||
variance := 0.0
|
||||
if fx != nil {
|
||||
for i := range n {
|
||||
d := fx[i*p+j] - mean
|
||||
variance += d * d
|
||||
}
|
||||
} else {
|
||||
for i := range n {
|
||||
d := x.FloatAt(i*p+j) - mean
|
||||
variance += d * d
|
||||
}
|
||||
}
|
||||
scale := math.Sqrt(variance * inv)
|
||||
if scale == 0 {
|
||||
return nil, base.Errf("%s: design column %d is constant and cannot be standardised", name, j)
|
||||
}
|
||||
col := make([]float64, n)
|
||||
if fx != nil {
|
||||
for i := range n {
|
||||
col[i] = (fx[i*p+j] - mean) / scale
|
||||
}
|
||||
} else {
|
||||
for i := range n {
|
||||
col[i] = (x.FloatAt(i*p+j) - mean) / scale
|
||||
}
|
||||
}
|
||||
std.xs[j] = col
|
||||
std.means[j] = mean
|
||||
std.scales[j] = scale
|
||||
}
|
||||
yMean := 0.0
|
||||
// The walk is bounded by the response's own element count: a rebased
|
||||
// view's payload may run past its visible elements.
|
||||
if fy != nil {
|
||||
for _, v := range fy[:n] {
|
||||
yMean += v
|
||||
}
|
||||
} else {
|
||||
for i := range n {
|
||||
yMean += y.FloatAt(i)
|
||||
}
|
||||
}
|
||||
yMean *= inv
|
||||
std.yMean = yMean
|
||||
for i := range n {
|
||||
var yv float64
|
||||
if fy != nil {
|
||||
yv = fy[i]
|
||||
} else {
|
||||
yv = y.FloatAt(i)
|
||||
}
|
||||
std.yc[i] = yv - yMean
|
||||
}
|
||||
for j := range p {
|
||||
s := 0.0
|
||||
for i := range n {
|
||||
s += std.xs[j][i] * std.xs[j][i]
|
||||
}
|
||||
std.zs[j] = s * inv
|
||||
}
|
||||
return std, nil
|
||||
}
|
||||
|
||||
// elasticNetFit runs the coordinate descent on the standardised
|
||||
// system. start, when given, is the warm start the path hands down;
|
||||
// a nil start begins at zero, the exact answer at the top of the
|
||||
// grid. The running buffers live in ws and are reused across calls:
|
||||
// the path fits one lambda after another into the same workspace, and
|
||||
// a single fit simply owns one it just allocated. Each cycle sweeps
|
||||
// the coordinates in order, moving each slope to the exact minimiser
|
||||
// of the objective along its own coordinate and folding the movement
|
||||
// into the running residual, until a full cycle moves nothing by more
|
||||
// than the tolerance. The per-coordinate minimiser is the
|
||||
// soft-thresholded least-squares coordinate over the denominator the
|
||||
// elastic net puts there: pure lasso thresholding when alpha is 1,
|
||||
// plain ridge division when alpha is 0.
|
||||
func elasticNetFit(std *lassoStandardisation, lambda, alpha float64, start []float64, ws *lassoWorkspace) (beta []float64, iterations int, converged bool) {
|
||||
n := len(std.yc)
|
||||
p := len(std.xs)
|
||||
inv := 1 / float64(n)
|
||||
if start == nil {
|
||||
beta = append(ws.beta[:0], make([]float64, p)...)
|
||||
} else {
|
||||
beta = append(ws.beta[:0], start...)
|
||||
}
|
||||
ws.beta = beta
|
||||
// The running residual r = yc − Xβ, kept current by folding each
|
||||
// coordinate's movement in as it happens: the sweep then reads the
|
||||
// partial residual every coordinate needs without a full product
|
||||
// per coordinate.
|
||||
r := append(ws.r[:0], std.yc...)
|
||||
ws.r = r
|
||||
if start != nil {
|
||||
for j := range p {
|
||||
bj := beta[j]
|
||||
if bj == 0 {
|
||||
continue
|
||||
}
|
||||
col := std.xs[j]
|
||||
for i := range n {
|
||||
r[i] -= col[i] * bj
|
||||
}
|
||||
}
|
||||
}
|
||||
for iter := 1; iter <= lassoMaxIter; iter++ {
|
||||
worst := 0.0
|
||||
for j := range p {
|
||||
col := std.xs[j]
|
||||
zj := std.zs[j]
|
||||
s := 0.0
|
||||
for i := range n {
|
||||
s += col[i] * r[i]
|
||||
}
|
||||
rho := s*inv + zj*beta[j]
|
||||
newBeta := softThreshold(rho, lambda*alpha) / (zj + lambda*(1-alpha))
|
||||
if newBeta != beta[j] {
|
||||
d := newBeta - beta[j]
|
||||
for i := range n {
|
||||
r[i] -= col[i] * d
|
||||
}
|
||||
beta[j] = newBeta
|
||||
if a := math.Abs(d); a > worst {
|
||||
worst = a
|
||||
}
|
||||
}
|
||||
}
|
||||
if worst < lassoTolerance {
|
||||
return beta, iter, true
|
||||
}
|
||||
}
|
||||
return beta, lassoMaxIter, false
|
||||
}
|
||||
|
||||
// softThreshold is the proximal operator of the absolute value: the
|
||||
// identity past the threshold and zero inside it. It is what makes
|
||||
// the lasso exact on an orthonormal design, where the coordinates
|
||||
// decouple and every slope is this function of its own projection.
|
||||
func softThreshold(v, t float64) float64 {
|
||||
if v > t {
|
||||
return v - t
|
||||
}
|
||||
if v < -t {
|
||||
return v + t
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// lassoCoefficients lifts the standardised slopes back to the design as
|
||||
// supplied: the slope of a standardised column scales by the column's
|
||||
// standard deviation, and the intercept reabsorbs the column means, in
|
||||
// the columns' own order, so the accumulation is the one the fit's
|
||||
// reporting pass runs.
|
||||
func lassoCoefficients(std *lassoStandardisation, betaStd []float64) (intercept float64, coefficients []float64) {
|
||||
p := len(betaStd)
|
||||
coefficients = make([]float64, p)
|
||||
intercept = std.yMean
|
||||
for j := range p {
|
||||
coefficients[j] = betaStd[j] / std.scales[j]
|
||||
intercept -= coefficients[j] * std.means[j]
|
||||
}
|
||||
return intercept, coefficients
|
||||
}
|
||||
|
||||
// lassoResult lifts the standardised slopes back to the design as
|
||||
// supplied: the slope of a standardised column scales by the column's
|
||||
// standard deviation, and the intercept reabsorbs the column means.
|
||||
// Fitted and Residuals are then computed on the original design, the
|
||||
// same final sweep WeightedLinearRegression runs.
|
||||
func lassoResult(name string, n, p int, x, y *core.Array, fx, fy []float64, std *lassoStandardisation, betaStd []float64, lambda, alpha float64, iterations int, converged bool) (*ElasticNetResult, error) {
|
||||
intercept, coefficients := lassoCoefficients(std, betaStd)
|
||||
out := &ElasticNetResult{
|
||||
Intercept: intercept,
|
||||
Coefficients: coefficients,
|
||||
ColumnMeans: std.means,
|
||||
ColumnScales: std.scales,
|
||||
Lambda: lambda,
|
||||
Alpha: alpha,
|
||||
Iterations: iterations,
|
||||
Converged: converged,
|
||||
}
|
||||
out.Fitted = make([]float64, n)
|
||||
out.Residuals = make([]float64, n)
|
||||
for i := range n {
|
||||
f := out.Intercept
|
||||
if fx != nil {
|
||||
row := fx[i*p : i*p+p]
|
||||
for j, xj := range row {
|
||||
f += coefficients[j] * xj
|
||||
}
|
||||
} else {
|
||||
for j := range p {
|
||||
f += coefficients[j] * x.FloatAt(i*p+j)
|
||||
}
|
||||
}
|
||||
var yv float64
|
||||
if fy != nil {
|
||||
yv = fy[i]
|
||||
} else {
|
||||
yv = y.FloatAt(i)
|
||||
}
|
||||
out.Fitted[i] = f
|
||||
out.Residuals[i] = yv - f
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
Reference in New Issue
Block a user