tensor/stats/regression_edge_pin_test.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package stats

import (
	"math"
	"strings"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Regression-edge pins for the linear model: the statistics of a
// design without an intercept column, the tail of the coefficient
// p-values, and the exact-fit report.

// fact returns n!, small n only.
func fact(n int) float64 {
	r := 1.0
	for i := 2; i <= n; i++ {
		r *= float64(i)
	}
	return r
}

// binom returns the binomial coefficient, small n only.
func binom(n, k int) float64 { return fact(n) / (fact(k) * fact(n-k)) }

// tTailExactEven returns P(|T| > t) for even df in closed form. There,
// the incomplete beta has a = df/2, an integer, and b = 1/2, so the
// integral is elementary: nothing is approximated and nothing from the
// library is used.
func tTailExactEven(tv float64, df int) float64 {
	m := df / 2
	z := float64(df) / (float64(df) + tv*tv)
	// ∫_0^z u^(m−1)(1−u)^(−1/2) du, with u = 1 − s², is
	// 2∫_{√(1−z)}^{1} (1−s²)^(m−1) ds.
	s0 := math.Sqrt(1 - z)
	antiderivative := func(s float64) float64 {
		sum := 0.0
		for k := 0; k <= m-1; k++ {
			sum += binom(m-1, k) * math.Pow(-1, float64(k)) * math.Pow(s, float64(2*k+1)) / float64(2*k+1)
		}
		return sum
	}
	num := 2 * (antiderivative(1) - antiderivative(s0))
	lg, _ := math.Lgamma(float64(m))
	lgb, _ := math.Lgamma(0.5)
	lgs, _ := math.Lgamma(float64(m) + 0.5)
	beta := math.Exp(lg + lgb - lgs) // B(m, 1/2)
	return num / beta
}

// TestTwoSidedTAccuracy pins the coefficient tail against exact closed
// forms and against the far tail, where the previous 2·(1 − T_cdf)
// form lost every digit and returned an exact zero.
func TestTwoSidedTAccuracy(t *testing.T) {
	// Exact references: Cauchy (df = 1) and the df = 2 closed form.
	for _, tc := range []struct {
		t    float64
		want float64
	}{{0.5, 2 * math.Atan(1/0.5) / math.Pi}, {2, 2 * math.Atan(1.0/2) / math.Pi}} {
		got, err := twoSidedT(tc.t, 1)
		if err != nil {
			t.Fatal(err)
		}
		if math.Abs(got-tc.want) > 1e-13*tc.want {
			t.Fatalf("twoSidedT(%v, 1) = %.17g, want %.17g", tc.t, got, tc.want)
		}
	}
	for _, tc := range []struct {
		t    float64
		want float64
	}{{0.5, 1 - 0.5/math.Sqrt(2+0.25)}, {2, 1 - 2/math.Sqrt(6)}} {
		got, err := twoSidedT(tc.t, 2)
		if err != nil {
			t.Fatal(err)
		}
		if math.Abs(got-tc.want) > 1e-13*tc.want {
			t.Fatalf("twoSidedT(%v, 2) = %.17g, want %.17g", tc.t, got, tc.want)
		}
	}
	// Even degrees of freedom: the elementary closed form.
	for _, tc := range []struct {
		t  float64
		df int
	}{{2, 6}, {8, 6}, {0.5, 6}, {2, 10}, {8, 10}, {3, 20}} {
		got, err := twoSidedT(tc.t, tc.df)
		if err != nil {
			t.Fatalf("twoSidedT(%v, %d): %v", tc.t, tc.df, err)
		}
		want := tTailExactEven(tc.t, tc.df)
		if math.Abs(got-want) > 1e-9*want {
			t.Fatalf("twoSidedT(%v, %d) = %.17g, closed form says %.17g", tc.t, tc.df, got, want)
		}
	}
	// Large df, against a reference computed outside the library to
	// 60 digits: the tail of t(8 | df = 1e6) is
	// 1.2455063433202503e-15. The cancelling form returned 1.33227e-15
	// here, 7 % high, so this pins the accuracy rather than the order.
	const independentTail = 1.2455063433202503e-15
	got, err := twoSidedT(8, 1000000)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(got-independentTail) > 1e-6*independentTail {
		t.Fatalf("twoSidedT(8, 1e6) = %.17g, the independent reference is %.17g", got, independentTail)
	}
	// The far tail must stay positive: the cancelling form returned 0
	// for t = 30, df = 100, where the true tail is 8.4e-52.
	far, err := twoSidedT(30, 100)
	if err != nil {
		t.Fatal(err)
	}
	if far <= 0 {
		t.Fatalf("twoSidedT(30, 100) = %v, want a positive tail", far)
	}
	if far > 1e-40 {
		t.Fatalf("twoSidedT(30, 100) = %v, want a tail near 8.4e-52", far)
	}
	// Exact value at the centre.
	if p, err := twoSidedT(0, 7); err != nil || p != 1 {
		t.Fatalf("twoSidedT(0, 7) = %v (err %v), want exactly 1", p, err)
	}
}

// TestLinearRegressionWithoutIntercept pins the statistics of a design
// with no constant column: the uncentred total sum of squares is the
// null model, and the model degrees of freedom are the column count.
func TestLinearRegressionWithoutIntercept(t *testing.T) {
	t.Run("single column", func(t *testing.T) {
		x := mustMatrix(t, []float64{1, 0, -1}, 3, 1)
		y := mustFloats(t, []float64{3, 1, 2}, 3)
		res, err := LinearRegression(x, y)
		if err != nil {
			t.Fatalf("LinearRegression: %v", err)
		}
		// beta = Σxy/Σx² = 1/2, rss = 13.5, Σy² = 14.
		if math.Abs(res.Coefficients[0]-0.5) > 1e-15 {
			t.Fatalf("slope = %v, want 0.5", res.Coefficients[0])
		}
		wantR2 := 1 - 13.5/14
		if math.Abs(res.RSquared-wantR2) > 1e-12 {
			t.Fatalf("R² = %v, want %v (uncentred)", res.RSquared, wantR2)
		}
		if res.DModel != 1 {
			t.Fatalf("DModel = %d, want 1", res.DModel)
		}
		if !(res.FStatistic > 0) || res.FPValue <= 0 || res.FPValue > 1 {
			t.Fatalf("F = %v, p = %v, want a positive statistic and a probability", res.FStatistic, res.FPValue)
		}
	})
	t.Run("two columns", func(t *testing.T) {
		x := mustMatrix(t, []float64{1, 0, 0, 1, -1, 1}, 3, 2)
		y := mustFloats(t, []float64{3, 1, 2}, 3)
		res, err := LinearRegression(x, y)
		if err != nil {
			t.Fatalf("LinearRegression: %v", err)
		}
		// beta = [5/3, 7/3], rss = 16/3, Σy² = 14.
		wantR2 := 1 - (16.0/3)/14
		if math.Abs(res.RSquared-wantR2) > 1e-12 {
			t.Fatalf("R² = %v, want %v", res.RSquared, wantR2)
		}
		if res.DModel != 2 {
			t.Fatalf("DModel = %d, want 2", res.DModel)
		}
	})
	t.Run("intercept unchanged", func(t *testing.T) {
		// The centred form still applies when a constant column is
		// present: an exact line fits perfectly.
		x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2)
		y := mustFloats(t, []float64{3, 5, 7}, 3)
		res, err := LinearRegression(x, y)
		if err != nil {
			t.Fatalf("LinearRegression: %v", err)
		}
		if math.Abs(res.RSquared-1) > 1e-12 {
			t.Fatalf("R² = %v, want 1", res.RSquared)
		}
		if res.DModel != 1 {
			t.Fatalf("DModel = %d, want p−1 = 1", res.DModel)
		}
	})
}

// TestLinearRegressionExactFitReport pins the exact-fit report: zero
// standard errors mean an infinite statistic, not a zero one, and the
// p-value is zero rather than absent.
func TestLinearRegressionExactFitReport(t *testing.T) {
	x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3, 1, 4}, 4, 2)
	y := mustFloats(t, []float64{3, 5, 7, 9}, 4)
	res, err := LinearRegression(x, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if res.RSquared != 1 {
		t.Fatalf("R² = %v, want exactly 1", res.RSquared)
	}
	for j := range 2 {
		if res.StandardErrors[j] != 0 {
			t.Fatalf("se[%d] = %v, want 0", j, res.StandardErrors[j])
		}
		if !math.IsInf(res.TStatistics[j], 0) {
			t.Fatalf("t[%d] = %v, want ±Inf", j, res.TStatistics[j])
		}
		if res.PValues[j] != 0 {
			t.Fatalf("p[%d] = %v, want 0", j, res.PValues[j])
		}
	}
}

// TestLinearRegressionRefusesNonFinite pins the input contract: the
// other tests in the package refuse non-finite data and this one must
// not answer with a silent column of NaN.
func TestLinearRegressionRefusesNonFinite(t *testing.T) {
	x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2)
	for _, bad := range []float64{math.NaN(), math.Inf(1), math.Inf(-1)} {
		y := mustFloats(t, []float64{3, bad, 7}, 3)
		if _, err := LinearRegression(x, y); err == nil {
			t.Fatalf("expected an error for the response holding %v", bad)
		} else if !strings.Contains(err.Error(), "non-finite") {
			t.Fatalf("error = %v, want a non-finite refusal", err)
		}
	}
	xb := mustMatrix(t, []float64{1, 1, 1, 2, 1, math.NaN()}, 3, 2)
	if _, err := LinearRegression(xb, mustFloats(t, []float64{3, 5, 7}, 3)); err == nil {
		t.Fatal("expected an error for a design holding NaN")
	}
}

// mustMatrix builds an (r, c) float64 array.
func mustMatrix(t *testing.T, vals []float64, r, c int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, r, c)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}