Files

459 lines
15 KiB
Go
Raw Permalink Normal View History

2026-09-03 10:00:00 +02:00
// 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"
)
// Quantile regression: the linear fit that asks not for the mean of
// the response but for its tau-th conditional quantile, by minimising
// the check loss
//
// ρ_τ(r) = r·τ when r ≥ 0, r·(τ − 1) when r < 0,
//
// the asymmetric absolute loss that rewards a fitted line for putting
// the right fraction of the data beneath it. The implementation is the
// Frisch-Newton interior-point method on the dual, the form
// Portnoy and Koenker put the method in: the check loss problem is
// the linear program
//
// min Σᵢ τ·uᵢ + (1−τ)·vᵢ subject to u − v = y − Xβ, u, v ≥ 0,
//
// and its dual asks for
//
// max wᵀy subject to Xᵀw = (1−τ)·Xᵀ1, w ∈ [0, 1]ⁿ.
//
// At the optimum an observation with a positive residual carries w = 1,
// one with a negative residual w = 0, and the observations the fit
// reproduces exactly carry w strictly inside the box. A logarithmic
// barrier is laid on the box, every iteration solves the barrier's
// Newton system exactly in the p×p form XᵀD⁻¹X through the shared LU
// solve, and the barrier parameter falls geometrically. The start is
// exactly feasible, w = 1−τ for every observation, and the primal fit
// is read back off the interior set at every iteration, so the run
// records a monotone descent of the true check loss.
// The documented schedule of the interior-point loop: the barrier
// parameter starts at the mean absolute residual of the ordinary
// least squares start, shrinks by this factor every iteration, and the
// run is settled when no coefficient of the dual moves by more than
// the tolerance, or when the parameter has fallen sixteen orders of
// magnitude, whichever comes first.
const (
quantileMaxIterations = 100
quantileMuShrink = 0.25
quantileMuFloorRatio = 1e-16
quantileStepFraction = 0.995
quantileTolerance = 1e-14
)
// QuantileRegressionResult carries a quantile regression fit.
type QuantileRegressionResult struct {
// Coefficients are the quantile estimates β̂, one per design
// column, in the design's own order. The intercept, supplied by
// the caller as a constant column, is estimated like any other
// coefficient: the check loss pulls it to the response's tau-th
// quantile at x = 0.
Coefficients []float64
// Fitted and Residuals align with the rows of the design.
Fitted []float64
Residuals []float64
// Tau is the quantile the fit minimises the check loss for.
Tau float64
// CheckLoss is the minimised check loss Σᵢ ρ_τ(rᵢ) at the fit.
CheckLoss float64
// Objective records the best check loss seen after every
// interior-point iteration, starting from the ordinary least
// squares start. It is the instrument that shows the optimisation
// descending: monotone non-increasing by construction, because an
// iterate only enters the record by improving on every iterate
// before it.
Objective []float64
// Iterations counts the interior-point iterations taken;
// Converged reports whether the run settled by its own stopping
// rules.
Iterations int
Converged bool
}
// QuantileRegression fits y = X·β for the tau-th conditional quantile
// by the Frisch-Newton interior-point method on the dual of the check
// loss program. The design carries n rows and p columns exactly as
// LinearRegression's, the intercept included by the caller as a
// constant column when wanted, and the same validations apply: n > p,
// a full-rank design, real finite input. tau must lie strictly inside
// (0, 1): the closed ends have no regression answer, tau = 0 and
// tau = 1 being the envelope of the data rather than a fit.
//
// The run starts from the ordinary least squares fit and follows the
// barrier path in the dual, where the equality constraint is satisfied
// exactly from the first iterate to the last. The primal fit is read
// back off the dual's interior set, the observations the optimal fit
// reproduces exactly, by least squares over that set; a degenerate
// optimum, whose interior set is underdetermined, is resolved to the
// smallest-norm member of the optimal face by a ridge of the order of
// the solve's own rounding.
func QuantileRegression(x, y *core.Array, tau float64) (*QuantileRegressionResult, error) {
const name = "QuantileRegression"
if x.NDim() != 2 {
return nil, base.Errf("%s: the design must be rank 2, got shape %s", name, base.ShapeText(x.Shape()))
}
if y.NDim() != 1 {
return nil, base.Errf("%s: the response must be rank 1", name)
}
if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
return nil, base.Errf("%s: complex inputs are not supported", name)
}
n, p := x.Shape()[0], x.Shape()[1]
if y.Len() != n {
return nil, base.Errf("%s: the design has %d rows but the response %d", name, n, y.Len())
}
if n <= p {
return nil, base.Errf("%s: need n > p, got %d observations and %d columns", name, n, p)
}
if p == 0 {
return nil, base.Errf("%s: the design must carry at least one column", name)
}
// NaN compares false against both bounds, so this refuses it too.
if !(tau > 0 && tau < 1) {
return nil, base.Errf("%s: tau must lie strictly inside (0, 1), got %g", name, tau)
}
if err := checkFinite(name, "the design", x); err != nil {
return nil, err
}
if err := checkFinite(name, "the response", y); err != nil {
return nil, err
}
fx := rawFloats(x)
fy := rawFloats(y)
if fy == nil {
// The response is a view or a narrower dtype: widen it once so
// the barrier loops below sweep a plain slice, the same bits
// the widening accessor returns.
fy = make([]float64, n)
for i := range n {
fy[i] = y.FloatAt(i)
}
}
// The start is the ordinary least squares answer. It is not the
// optimum of any asymmetric loss, but it is inside the basin, its
// residual scale sets the barrier's starting parameter, and a
// design the shared LU cannot solve is refused here, rank
// deficiency and all.
beta, err := leastSquaresSolve(name, n, p, x, y, fx, fy, nil)
if err != nil {
return nil, err
}
residuals := make([]float64, n)
regressionResiduals(n, p, x, y, fx, fy, beta, residuals)
meanAbs := 0.0
for _, r := range residuals {
meanAbs += math.Abs(r)
}
meanAbs /= float64(n)
lossAt := func(b []float64) float64 {
total := 0.0
for i := range n {
r := fy[i]
for j := range p {
var xj float64
if fx != nil {
xj = fx[i*p+j]
} else {
xj = x.FloatAt(i*p + j)
}
r -= b[j] * xj
}
if r >= 0 {
total += r * tau
} else {
// r and tau − 1 are both negative, so the product is
// the positive |r|·(1 − tau) the check loss asks for.
total += r * (tau - 1)
}
}
return total
}
loss := lossAt(beta)
out := &QuantileRegressionResult{
Coefficients: beta,
Tau: tau,
CheckLoss: loss,
Objective: []float64{loss},
}
if meanAbs == 0 {
// The ordinary least squares fit reproduces the response
// exactly: the check loss is zero and no asymmetric loss can
// beat zero. The run is settled before it begins.
out.Fitted = make([]float64, n)
out.Residuals = make([]float64, n)
copy(out.Residuals, residuals)
for i := range n {
out.Fitted[i] = fy[i]
}
out.Converged = true
return out, nil
}
// The dual barrier state. w = 1−τ for every observation satisfies
// the equality Xᵀw = (1−τ)·Xᵀ1 identically and sits exactly in the
// middle of the box, the interior point the method asks for.
mu := meanAbs
muFloor := mu * quantileMuFloorRatio
w := make([]float64, n)
dw := make([]float64, n)
gradient := make([]float64, n)
curvature := make([]float64, n)
cols := make([]float64, n)
colSums := make([]float64, p)
// The backtracking buffer holds only complete candidate vectors: every
// iteration rewrites all n slots before the objective reads any, so one
// buffer serves the whole run.
trial := make([]float64, n)
for i := range n {
w[i] = 1 - tau
}
for i := range n {
for j := range p {
var xj float64
if fx != nil {
xj = fx[i*p+j]
} else {
xj = x.FloatAt(i*p + j)
}
colSums[j] += xj
}
}
bestBeta := append([]float64(nil), beta...)
bestLoss := loss
converged := false
iterations := quantileMaxIterations
// The barrier system's storage rides the whole run: the upper
// triangle is refilled by accumulation from an explicit zero and
// its mirror copies the lower one, and the right-hand wrapper is
// fixed because the solve writes through rhs in place.
normal := make([][]float64, p)
for j := range p {
normal[j] = make([]float64, p)
}
rhs := make([]float64, p)
solveRHS := [][]float64{rhs}
// The barrier objective Φ = wᵀy + μΣ(log w + log(1−w)) rises
// along the Newton direction of this concave problem by
// ∇ΦᵀΔw = ΔwᵀDΔw ≥ 0 identically, so a backtracking half-step
// finds an ascent all the way to the barrier maximum.
objective := func(wv []float64, muv float64) float64 {
total := 0.0
for i := range n {
total += wv[i]*fy[i] + muv*(math.Log(wv[i])+math.Log(1-wv[i]))
}
return total
}
for iter := 1; iter <= quantileMaxIterations; iter++ {
// The Newton system of the barrier problem, eliminated to the
// p×p form XᵀD⁻¹X Δλ = XᵀD⁻¹z: z the barrier gradient
// y + μ(1/w − 1/(1−w)), D⁻¹ the reciprocal of the barrier's
// curvature μ((1−w)² + w²)/(w²(1−w)²). The solve is exact, the
// equality direction XᵀΔw = 0 comes out of it, and the step
// keeps the equality exactly because it started exact.
for j := range p {
nj := normal[j]
for k := j; k < p; k++ {
nj[k] = 0
}
}
clear(rhs)
for i := range n {
if fx != nil {
for j := range p {
cols[j] = fx[i*p+j]
}
} else {
for j := range p {
cols[j] = x.FloatAt(i*p + j)
}
}
oneMinus := 1 - w[i]
gradient[i] = fy[i] + mu*(1/w[i]-1/oneMinus)
curvature[i] = w[i] * w[i] * oneMinus * oneMinus /
(mu * (oneMinus*oneMinus + w[i]*w[i]))
scale := gradient[i] * curvature[i]
for j := range p {
rhs[j] += cols[j] * scale
for k := j; k < p; k++ {
normal[j][k] += cols[j] * cols[k] * curvature[i]
}
}
}
for j := range p {
for k := range j {
normal[j][k] = normal[k][j]
}
}
solved, err := base.SolveSystem(name, normal, solveRHS)
if err != nil {
return nil, base.Errf("%s: the barrier system is singular (%w)", name, err)
}
deltaLambda := solved[0]
mostMove := 0.0
for i := range n {
pred := 0.0
if fx != nil {
for j := range p {
pred += fx[i*p+j] * deltaLambda[j]
}
} else {
for j := range p {
pred += x.FloatAt(i*p+j) * deltaLambda[j]
}
}
dw[i] = curvature[i] * (gradient[i] - pred)
if m := math.Abs(dw[i]); m > mostMove {
mostMove = m
}
}
// Fraction to the boundary: the largest step that keeps every
// w strictly inside the box, taken at a documented fraction of
// it and never more than the full Newton step.
t := 1.0
for i := range n {
if dw[i] < 0 {
t = min(t, quantileStepFraction*(-w[i]/dw[i]))
}
if dw[i] > 0 {
t = min(t, quantileStepFraction*((1-w[i])/dw[i]))
}
}
// The barrier objective rises along the Newton direction of
// this concave problem (the note above the closure), so a
// backtracking half-step finds an ascent all the way to the
// barrier maximum.
current := objective(w, mu)
for {
for i := range n {
trial[i] = w[i] + t*dw[i]
}
if objective(trial, mu) >= current || t < quantileTolerance {
break
}
t /= 2
}
if t < quantileTolerance || mostMove < quantileTolerance {
// The barrier problem is stationary at this μ: further
// pressure buys nothing until μ falls, and μ is about to.
// The run settles when the barrier itself is exhausted.
converged = true
iterations = iter
break
}
copy(w, trial)
// The primal fit read off this iterate's interior set, and the
// monotone record it feeds.
candidate, err := quantileRecover(name, n, p, x, y, fx, fy, w)
if err != nil {
return nil, err
}
if l := lossAt(candidate); l < bestLoss {
bestLoss = l
copy(bestBeta, candidate)
}
out.Objective = append(out.Objective, bestLoss)
mu *= quantileMuShrink
if mu < muFloor {
converged = true
iterations = iter
break
}
}
if !converged {
return nil, base.Errf("%s: %d iterations did not converge", name, quantileMaxIterations)
}
out.Iterations = iterations
out.Converged = true
out.CheckLoss = bestLoss
copy(out.Coefficients, bestBeta)
out.Fitted = make([]float64, n)
out.Residuals = make([]float64, n)
regressionResiduals(n, p, x, y, fx, fy, out.Coefficients, out.Residuals)
for i := range n {
out.Fitted[i] = fy[i] - out.Residuals[i]
}
return out, nil
}
// quantileRecover reads a primal fit off a dual barrier point. The
// observations whose w sits strictly inside the box are the ones the
// optimal fit reproduces exactly, their residual being zero, so the
// least squares over that interior set is the fit; the degenerate
// case, an interior set that underdetermines the coefficients, is
// resolved to the smallest-norm member of the optimal face by a ridge
// a trillionth of the normal equations' own scale, far below the
// solve's meaningful digits. An empty interior set falls back to the
// plain least squares fit, the same reading the first iterate's
// all-interior box gives.
func quantileRecover(name string, n, p int, x, y *core.Array, fx, fy []float64, w []float64) ([]float64, error) {
const band = 1e-5
normal := make([][]float64, p)
for j := range p {
normal[j] = make([]float64, p)
}
rhs := make([]float64, p)
trace := 0.0
count := 0
for i := range n {
if !(w[i] > band && w[i] < 1-band) {
continue
}
count++
var yv float64
if fy != nil {
yv = fy[i]
} else {
yv = y.FloatAt(i)
}
for a := range p {
var xa float64
if fx != nil {
xa = fx[i*p+a]
} else {
xa = x.FloatAt(i*p + a)
}
rhs[a] += xa * yv
for b := range p {
var xb float64
if fx != nil {
xb = fx[i*p+b]
} else {
xb = x.FloatAt(i*p + b)
}
normal[a][b] += xa * xb
}
}
}
if count < p {
// An interior set too small to identify the coefficients: the
// whole box was effectively at its bounds, and the plain least
// squares fit is the honest reading. The caller's
// best-by-check-loss record keeps this from ever worsening the
// answer.
return leastSquaresSolve(name, n, p, x, y, fx, fy, nil)
}
for a := range p {
trace += normal[a][a]
}
for a := range p {
normal[a][a] += 1e-12 * trace / float64(p)
}
solved, err := base.SolveSystem(name, normal, [][]float64{rhs})
if err != nil {
return nil, base.Errf("%s: the interior set is degenerate (%w)", name, err)
}
return solved[0], nil
}