459 lines
15 KiB
Go
459 lines
15 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"
|
||
)
|
||
|
||
// 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
|
||
}
|