feat: initial release

Assisted-by: GLM 5.3 Flash

stats/bench_regression_test.go

@@ -0,0 +1,638 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package stats
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// The randomised and cross-checked tests that guard the pair counting of
+// KendallTau, the covariance diagonals of the regression entries and the
+// pairwise-slope walk of Theil-Sen against the implementations they
+// replace, followed by the benchmarks of those paths.
+
+// randSource is a deterministic linear congruential generator for the
+// randomised checks: a fixed recurrence makes every failing case
+// reproducible from its seed alone.
+type randSource struct{ state uint64 }
+
+// next returns the next draw in [0, 1).
+func (r *randSource) next() float64 {
+	r.state = r.state*6364136223846793005 + 1442695040888963407
+	return float64(r.state>>11) / float64(uint64(1)<<53)
+}
+
+// intn returns the next draw in [0, n).
+func (r *randSource) intn(n int) int { return int(r.next() * float64(n)) }
+
+// bruteForceTally counts the pairs of two paired samples the way the
+// definition reads: every ordered pair of positions, its two differences
+// multiplied, the sign of the product deciding the side, and the ties of
+// a sample counted over the same pair walk. It is the oracle the
+// merge-sort pair counting is checked against.
+func bruteForceTally(xs, ys []float64) pairTally {
+	var tally pairTally
+	for i := 1; i < len(xs); i++ {
+		for j := 0; j < i; j++ {
+			switch prod := (xs[i] - xs[j]) * (ys[i] - ys[j]); {
+			case prod > 0:
+				tally.concordant++
+			case prod < 0:
+				tally.discordant++
+			}
+			if xs[i] == xs[j] {
+				tally.tiedX++
+			}
+			if ys[i] == ys[j] {
+				tally.tiedY++
+			}
+		}
+	}
+	return tally
+}
+
+// bruteForceTau is the O(n²) τ-b of the definition, the oracle the
+// returned correlation is checked against.
+func bruteForceTau(xs, ys []float64) float64 {
+	tally := bruteForceTally(xs, ys)
+	n := float64(len(xs))
+	n0 := n * (n - 1) / 2
+	return (float64(tally.concordant) - float64(tally.discordant)) /
+		math.Sqrt((n0-float64(tally.tiedX))*(n0-float64(tally.tiedY)))
+}
+
+// TestKendallTauCountsMatchBruteForce checks the merge-sort pair
+// counting against the pair enumeration it replaces. The values are
+// drawn from a handful of levels so that ties in either sample are the
+// common case, which is where the block scan and the tie accounting
+// carry the work; continuous draws, degenerate ranges and a constant
+// second sample cover the rest. Both sides are exact integers, so the
+// comparison is on the returned bits rather than on a tolerance.
+func TestKendallTauCountsMatchBruteForce(t *testing.T) {
+	src := &randSource{state: 0x9e3779b97f4a7c15}
+	levels := []float64{-1, 0, 1, 2, 3}
+	for trial := range 400 {
+		n := 2 + src.intn(38)
+		xs := make([]float64, n)
+		ys := make([]float64, n)
+		// A third of the draws are continuous, the rest come from the
+		// levels, so both samples are tie-heavy on average.
+		draw := func() float64 {
+			if src.intn(3) == 0 {
+				return src.next()*4 - 2
+			}
+			return levels[src.intn(len(levels))]
+		}
+		for i := range n {
+			xs[i] = draw()
+			ys[i] = draw()
+		}
+		if trial%4 == 0 {
+			// The pairing is perfectly monotone, its τ exactly 1.
+			ys = append([]float64(nil), xs...)
+		}
+		if trial%9 == 0 {
+			// A constant chunk inside otherwise varying samples: the
+			// widest blocks of equal values the scan has to handle.
+			for i := n / 3; i < 2*n/3; i++ {
+				xs[i] = xs[0]
+			}
+		}
+		want := bruteForceTally(xs, ys)
+		got := tallyPairs(append([]float64(nil), xs...), append([]float64(nil), ys...))
+		if got != want {
+			t.Fatalf("trial %d, n %d: pair counts %+v, want %+v\nx = %v\ny = %v", trial, n, got, want, xs, ys)
+		}
+		if want.tiedX == int64(n*(n-1)/2) || want.tiedY == int64(n*(n-1)/2) {
+			continue // a constant sample has no τ-b to compare
+		}
+		gotTau, err := KendallTau(mustFloats(t, xs), mustFloats(t, ys))
+		if err != nil {
+			t.Fatalf("trial %d, n %d: KendallTau: %v", trial, n, err)
+		}
+		wantTau := bruteForceTau(xs, ys)
+		if gotTau != wantTau {
+			t.Fatalf("trial %d, n %d: τ = %v (%#x), want %v (%#x)\nx = %v\ny = %v",
+				trial, n, gotTau, math.Float64bits(gotTau), wantTau, math.Float64bits(wantTau), xs, ys)
+		}
+	}
+}
+
+// invertDiagonal returns the diagonal of the inverse of a dense matrix
+// by Gauss-Jordan elimination with partial pivoting: an independent
+// route to the covariance diagonals the fits read through the shared LU
+// solve. It is accurate to rounding, which is all the cross-checks below
+// need from it.
+func invertDiagonal(m [][]float64) []float64 {
+	p := len(m)
+	aug := make([][]float64, p)
+	for i := range p {
+		aug[i] = make([]float64, 2*p)
+		copy(aug[i], m[i])
+		aug[i][p+i] = 1
+	}
+	for k := range p {
+		piv := k
+		for i := k + 1; i < p; i++ {
+			if math.Abs(aug[i][k]) > math.Abs(aug[piv][k]) {
+				piv = i
+			}
+		}
+		aug[k], aug[piv] = aug[piv], aug[k]
+		d := aug[k][k]
+		for j := range 2 * p {
+			aug[k][j] /= d
+		}
+		for i := range p {
+			if i == k {
+				continue
+			}
+			f := aug[i][k]
+			if f == 0 {
+				continue
+			}
+			for j := range 2 * p {
+				aug[i][j] -= f * aug[k][j]
+			}
+		}
+	}
+	diag := make([]float64, p)
+	for i := range p {
+		diag[i] = aug[i][p+i]
+	}
+	return diag
+}
+
+// copyMatrix copies a p×p system, the shape SolveSystem factors in place.
+func copyMatrix(m [][]float64) [][]float64 {
+	out := make([][]float64, len(m))
+	for i, row := range m {
+		out[i] = append([]float64(nil), row...)
+	}
+	return out
+}
+
+// TestSolveSystemUnitColumnsMatchPerColumnSolves checks the shared
+// factorisation the covariance diagonals read against the
+// per-coefficient solves it replaced: the same normal-equations matrix,
+// once against all p unit columns and once against each column alone, on
+// a positive-definite system built from a deterministic design. The
+// diagonal entries are compared bit for bit.
+func TestSolveSystemUnitColumnsMatchPerColumnSolves(t *testing.T) {
+	const p = 9
+	src := &randSource{state: 0xda3e39cb94b95bdb}
+	m := make([][]float64, p)
+	for i := range p {
+		m[i] = make([]float64, p)
+	}
+	row := make([]float64, p)
+	for range 40 {
+		row[0] = 1
+		for c := 1; c < p; c++ {
+			row[c] = src.next() - 0.5
+		}
+		for i := range p {
+			for j := range p {
+				m[i][j] += row[i] * row[j]
+			}
+		}
+	}
+	units := make([][]float64, p)
+	for j := range p {
+		units[j] = make([]float64, p)
+		units[j][j] = 1
+	}
+	together, err := base.SolveSystem("test", copyMatrix(m), units)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for j := range p {
+		e := make([]float64, p)
+		e[j] = 1
+		alone, err := base.SolveSystem("test", copyMatrix(m), [][]float64{e})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := together[j][j], alone[0][j]; got != want {
+			t.Fatalf("diagonal %d of the batched solve = %v (%#x), want %v (%#x) from the single-column solve",
+				j, got, math.Float64bits(got), want, math.Float64bits(want))
+		}
+	}
+}
+
+// TestLogisticRegressionWaldStandardErrors rebuilds the Fisher
+// information at the reported fit from the reported probabilities and
+// inverts it by an independent elimination: the Wald standard errors are
+// that inverse's diagonal, so the check covers the shared factorisation
+// the entry now reads them through on a design no digest pins.
+func TestLogisticRegressionWaldStandardErrors(t *testing.T) {
+	const n, p = 300, 6
+	src := &randSource{state: 0x2545f4914f6cdd1d}
+	design := make([]float64, n*p)
+	response := make([]float64, n)
+	for r := range n {
+		eta := 0.3
+		for c := range p {
+			v := src.next() - 0.5
+			design[r*p+c] = v
+			eta += 0.8 * v
+		}
+		design[r*p] = 1
+		if src.next() < 1/(1+math.Exp(-eta)) {
+			response[r] = 1
+		}
+	}
+	res, err := LogisticRegression(mustFloats(t, design, n, p), mustFloats(t, response, n))
+	if err != nil {
+		t.Fatal(err)
+	}
+	fisher := make([][]float64, p)
+	for i := range p {
+		fisher[i] = make([]float64, p)
+	}
+	for r := range n {
+		w := res.Fitted[r] * (1 - res.Fitted[r])
+		for i := range p {
+			for j := range p {
+				fisher[i][j] += design[r*p+i] * design[r*p+j] * w
+			}
+		}
+	}
+	diag := invertDiagonal(fisher)
+	for j := range p {
+		want := math.Sqrt(diag[j])
+		if rel := math.Abs(res.StandardErrors[j]-want) / want; rel > 1e-9 {
+			t.Fatalf("standard error %d = %.15g, want %.15g (relative %.3g)",
+				j, res.StandardErrors[j], want, rel)
+		}
+	}
+}
+
+// TestHuberRegressionStandardErrors does the same for the Huber fit:
+// σ²·(XᵀWX)⁻¹ with σ the reported robust scale and W the reported final
+// weights, its diagonal checked against an independent elimination of
+// the weighted normal equations the fit writes its own weights into.
+func TestHuberRegressionStandardErrors(t *testing.T) {
+	const n, p = 300, 6
+	src := &randSource{state: 0x853c49e6748fea9b}
+	design := make([]float64, n*p)
+	response := make([]float64, n)
+	for r := range n {
+		design[r*p] = 1
+		response[r] = 1
+		for c := 1; c < p; c++ {
+			v := src.next() - 0.5
+			design[r*p+c] = v
+			response[r] += float64(c) * 0.5 * v
+		}
+		response[r] += 0.2 * (src.next() - 0.5)
+	}
+	response[7] += 5 // the gross point the robust fit exists for
+	res, err := HuberRegression(mustFloats(t, design, n, p), mustFloats(t, response, n))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !(res.Scale > 0) {
+		t.Fatalf("the robust scale collapsed to %g on a contaminated sample", res.Scale)
+	}
+	wxx := make([][]float64, p)
+	for i := range p {
+		wxx[i] = make([]float64, p)
+	}
+	for r := range n {
+		w := res.Weights[r]
+		for i := range p {
+			for j := range p {
+				wxx[i][j] += design[r*p+i] * design[r*p+j] * w
+			}
+		}
+	}
+	diag := invertDiagonal(wxx)
+	for j := range p {
+		want := math.Sqrt(res.Scale * res.Scale * diag[j])
+		if rel := math.Abs(res.StandardErrors[j]-want) / want; rel > 1e-9 {
+			t.Fatalf("standard error %d = %.15g, want %.15g (relative %.3g)",
+				j, res.StandardErrors[j], want, rel)
+		}
+	}
+}
+
+// TestTheilSenParallelPairs checks the pairwise-slope walk against the
+// serial enumeration it replaces, element for element and bit for bit,
+// on a sample large enough to take the parallel path and with repeated
+// predictors, which are the pairs the walk has to skip. The reference
+// walk reproduces the original append order, so a slope landing in the
+// wrong buffer slot moves the median and fails the test.
+func TestTheilSenParallelPairs(t *testing.T) {
+	const n = 1400
+	src := &randSource{state: 0xc0ffee123456789a}
+	xs := make([]float64, n)
+	ys := make([]float64, n)
+	for i := range n {
+		// A tenth of a unit apart, so a tenth of the pairs share a
+		// predictor and carry no slope.
+		xs[i] = math.Floor(src.next()*300) / 10
+		ys[i] = 2*xs[i] + src.next()
+	}
+	gotIntercept, gotSlope, err := TheilSenRegression(mustFloats(t, xs), mustFloats(t, ys))
+	if err != nil {
+		t.Fatal(err)
+	}
+	slopes := make([]float64, 0, n*(n-1)/2)
+	for i := range n {
+		for j := i + 1; j < n; j++ {
+			if dx := xs[j] - xs[i]; dx != 0 {
+				slopes = append(slopes, (ys[j]-ys[i])/dx)
+			}
+		}
+	}
+	slope := medianSlice(slopes)
+	intercepts := make([]float64, n)
+	for i := range n {
+		intercepts[i] = ys[i] - slope*xs[i]
+	}
+	intercept := medianSlice(intercepts)
+	if gotSlope != slope || gotIntercept != intercept {
+		t.Fatalf("Theil-Sen = (%v, %v), want (%v, %v) from the serial walk",
+			gotIntercept, gotSlope, intercept, slope)
+	}
+}
+
+// regDesign builds a deterministic (n, p) design with an intercept column
+// and covariates in (−0.5, 0.5), plus the row-major values for the
+// cross-checks that rebuild a normal-equations matrix from them.
+func regDesign(tb testing.TB, n, p int, seed uint64) (*core.Array, []float64) {
+	tb.Helper()
+	src := &randSource{state: seed}
+	vals := make([]float64, n*p)
+	for r := range n {
+		vals[r*p] = 1
+		for c := 1; c < p; c++ {
+			vals[r*p+c] = src.next() - 0.5
+		}
+	}
+	a, err := core.FromFloats(vals, n, p)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	return a, vals
+}
+
+// BenchmarkKendallTau measures the rank correlation at a size where the
+// pair enumeration it replaced costs milliseconds: 8.4 million pairs,
+// counted by the merge sort in a fraction of that.
+func BenchmarkKendallTau(b *testing.B) {
+	const n = 4096
+	src := &randSource{state: 0x123456789abcdef}
+	xv := make([]float64, n)
+	yv := make([]float64, n)
+	for i := range n {
+		// Rounded draws, so ties in both samples are part of the walk.
+		xv[i] = math.Floor(src.next()*512) / 512
+		yv[i] = math.Floor(src.next()*512) / 512
+	}
+	x, err := core.FromFloats(xv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := KendallTau(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkLinearRegressionCovariance measures the least-squares fit with
+// its inference on a wide design, where the covariance diagonal is a
+// visible share of the work: p is large enough that the p solves it needs
+// outgrow the normal equations themselves.
+func BenchmarkLinearRegressionCovariance(b *testing.B) {
+	const n, p = 400, 48
+	x, _ := regDesign(b, n, p, 0x5eed0001)
+	src := &randSource{state: 0x5eed0002}
+	yv := make([]float64, n)
+	for i := range n {
+		yv[i] = src.next()*2 - 1
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := LinearRegression(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkPoissonRegressionInference measures the count fit with its
+// Wald inference, the covariance diagonal of the Fisher information
+// included.
+func BenchmarkPoissonRegressionInference(b *testing.B) {
+	const n, p = 600, 30
+	x, vals := regDesign(b, n, p, 0x5eed0003)
+	src := &randSource{state: 0x5eed0004}
+	yv := make([]float64, n)
+	for i := range n {
+		eta := 0.4
+		for c := range p {
+			eta += 0.3 * vals[i*p+c]
+		}
+		mu := math.Exp(eta)
+		yv[i] = math.Round(mu * (0.5 + src.next()))
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := PoissonRegression(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkLogisticRegressionInference measures the binary fit with its
+// Wald inference, the covariance diagonal of the Fisher information
+// included.
+func BenchmarkLogisticRegressionInference(b *testing.B) {
+	const n, p = 1500, 30
+	x, vals := regDesign(b, n, p, 0x5eed0005)
+	src := &randSource{state: 0x5eed0006}
+	yv := make([]float64, n)
+	for i := range n {
+		eta := 0.2 + 0.5*vals[i*p]
+		for c := 1; c < p; c++ {
+			eta += 0.5 * vals[i*p+c]
+		}
+		if src.next() < 1/(1+math.Exp(-eta)) {
+			yv[i] = 1
+		}
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := LogisticRegression(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkHuberRegressionInference measures the robust fit with its
+// standard errors: the reweighting rounds, their medians, and the
+// weighted normal equations read for the covariance diagonal.
+func BenchmarkHuberRegressionInference(b *testing.B) {
+	const n, p = 800, 8
+	x, vals := regDesign(b, n, p, 0x5eed0007)
+	src := &randSource{state: 0x5eed0008}
+	yv := make([]float64, n)
+	for i := range n {
+		yv[i] = 1 + 0.5*vals[i*p+1] - 0.3*vals[i*p+2] + 0.2*(src.next()-0.5)
+	}
+	yv[11] += 6
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := HuberRegression(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkTheilSenRegression measures the median-of-slopes fit: the
+// pairwise walk over the pairs with distinct predictors and the median
+// taken over them.
+func BenchmarkTheilSenRegression(b *testing.B) {
+	const n = 512
+	src := &randSource{state: 0x5eed0009}
+	xv := make([]float64, n)
+	yv := make([]float64, n)
+	for i := range n {
+		xv[i] = src.next()*10 - 5
+		yv[i] = 1 + 2*xv[i] + src.next()
+	}
+	x, err := core.FromFloats(xv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, _, err := TheilSenRegression(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkLassoPath measures the warm-started regularisation path over
+// the documented hundred-lambda grid, on an elastic net mixing that
+// exercises the standardised coordinate descent. The design carries no
+// intercept column: the penalty standardises every column, and a
+// constant one has no slope to report.
+func BenchmarkLassoPath(b *testing.B) {
+	const n, p = 400, 12
+	src := &randSource{state: 0x5eed000a}
+	xv := make([]float64, n*p)
+	yv := make([]float64, n)
+	for i := range n {
+		for c := range p {
+			xv[i*p+c] = src.next() - 0.5
+		}
+		yv[i] = 1 + 0.8*xv[i*p+1] - 0.5*xv[i*p+2] + 0.1*(src.next()-0.5)
+	}
+	x, err := core.FromFloats(xv, n, p)
+	if err != nil {
+		b.Fatal(err)
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := LassoPath(x, y, 1); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkKendallTauScaling reports the merge-sort pair count against
+// the sample size: the linear-ish growth the count replaced the
+// quadratic walk with.
+func BenchmarkKendallTauScaling(b *testing.B) {
+	for _, n := range []int{512, 2048, 8192} {
+		src := &randSource{state: 0x51ed + uint64(n)}
+		xv := make([]float64, n)
+		yv := make([]float64, n)
+		for i := range n {
+			xv[i] = math.Floor(src.next()*512) / 512
+			yv[i] = math.Floor(src.next()*512) / 512
+		}
+		x, err := core.FromFloats(xv, n)
+		if err != nil {
+			b.Fatal(err)
+		}
+		y, err := core.FromFloats(yv, n)
+		if err != nil {
+			b.Fatal(err)
+		}
+		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				if _, err := KendallTau(x, y); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+// BenchmarkTheilSenRegressionCap runs the fit at the observation cap,
+// where the pairwise slope list holds millions of entries and the middle
+// order statistic is taken by selection rather than by a full sort.
+func BenchmarkTheilSenRegressionCap(b *testing.B) {
+	const n = TheilSenMaxObservations
+	src := &randSource{state: 0x5eed1234}
+	xv := make([]float64, n)
+	yv := make([]float64, n)
+	for i := range n {
+		xv[i] = src.next()*10 - 5
+		yv[i] = 1 + 2*xv[i] + src.next()
+	}
+	x, err := core.FromFloats(xv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	y, err := core.FromFloats(yv, n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, _, err := TheilSenRegression(x, y); err != nil {
+			b.Fatal(err)
+		}
+	}
+}