510 lines
18 KiB
Go
510 lines
18 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"
|
||
)
|
||
|
||
// 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
|
||
}
|