tensor/stats/smallops_test.go

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

package stats

import (
	"math"
	"slices"
	"testing"

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

func smallVector(t *testing.T, vals []float64) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, len(vals))
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestRollingFamily checks the window reductions on a known series:
// alignment, the shrinking length and each reduction's value.
func TestRollingFamily(t *testing.T) {
	a := smallVector(t, []float64{1, 4, 2, 8, 5, 7})
	mean, err := RollingMean(a, 3)
	if err != nil {
		t.Fatalf("RollingMean: %v", err)
	}
	if mean.Len() != 4 {
		t.Fatalf("length %d, want 4", mean.Len())
	}
	wantMean := []float64{7.0 / 3, 14.0 / 3, 5, 20.0 / 3}
	for i := range wantMean {
		if math.Abs(mean.FloatAt(i)-wantMean[i]) > 1e-12 {
			t.Fatalf("mean[%d] = %.12g, want %.12g", i, mean.FloatAt(i), wantMean[i])
		}
	}
	sum, _ := RollingSum(a, 2)
	if sum.FloatAt(0) != 5 || sum.FloatAt(4) != 12 {
		t.Fatalf("sum ends = %g, %g, want 5, 12", sum.FloatAt(0), sum.FloatAt(4))
	}
	max, _ := RollingMax(a, 3)
	if max.FloatAt(0) != 4 || max.FloatAt(3) != 8 {
		t.Fatalf("max ends = %g, %g, want 4, 8", max.FloatAt(0), max.FloatAt(3))
	}
	min, _ := RollingMin(a, 2)
	if min.FloatAt(2) != 2 || min.FloatAt(1) != 2 {
		t.Fatalf("min = %g, %g, want 2, 2", min.FloatAt(1), min.FloatAt(2))
	}
	if _, err := RollingMean(a, 0); err == nil {
		t.Fatal("zero window accepted")
	}
	if _, err := RollingMean(a, 7); err == nil {
		t.Fatal("window past the length accepted")
	}
}

// TestMedianAbsoluteDeviation checks the definition on a sample with
// one wild point: the breakdown-proof scale stays at 1, the middle
// deviation, where the standard deviation would explode.
func TestMedianAbsoluteDeviation(t *testing.T) {
	a := smallVector(t, []float64{1, 2, 3, 4, 100})
	mad, err := MedianAbsoluteDeviation(a)
	if err != nil {
		t.Fatalf("MedianAbsoluteDeviation: %v", err)
	}
	if mad != 1 {
		t.Fatalf("MAD = %g, want 1", mad)
	}
}

// TestTrimmedMean checks the tail fraction: dropping one sample a
// side (ten percent) of the contaminated sample discards both tails,
// leaving the clean mean of 2 to 9.
func TestTrimmedMean(t *testing.T) {
	a := smallVector(t, []float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 100})
	mean, err := TrimmedMean(a, 0.1)
	if err != nil {
		t.Fatalf("TrimmedMean: %v", err)
	}
	if math.Abs(mean-5.5) > 1e-12 {
		t.Fatalf("trimmed mean = %.12g, want 5.5", mean)
	}
	if _, err := TrimmedMean(a, 0.5); err == nil {
		t.Fatal("half fraction accepted")
	}
	if _, err := TrimmedMean(a, -0.1); err == nil {
		t.Fatal("negative fraction accepted")
	}
}

// TestHistogram2D checks the count matrix on a constructed pair of
// samples, including the last-bin absorption of the maximum.
func TestHistogram2D(t *testing.T) {
	x := smallVector(t, []float64{0.5, 1.5, 2.5, 3.5, 2.9})
	y := smallVector(t, []float64{10, 20, 30, 40, 39})
	counts, xEdges, yEdges, err := Histogram2D(x, y, 4, 2)
	if err != nil {
		t.Fatalf("Histogram2D: %v", err)
	}
	if counts.Shape()[0] != 4 || counts.Shape()[1] != 2 {
		t.Fatalf("shape %v, want [4 2]", counts.Shape())
	}
	if len(xEdges) != 5 || len(yEdges) != 3 {
		t.Fatalf("edges %d and %d, want 5 and 3", len(xEdges), len(yEdges))
	}
	// x range [0.5, 3.5] over 4 bins of width 0.75; y range [10, 40]
	// over 2 bins of width 15.
	want := map[[2]int]int{
		{0, 0}: 1, // (0.5, 10)
		{1, 0}: 1, // (1.5, 20)
		{2, 1}: 1, // (2.5, 30)
		{3, 1}: 2, // (2.9, 39) and (3.5, 40) in the last bins
	}
	ints := counts.RawInts()
	for xi := range 4 {
		for yi := range 2 {
			got := int(ints[xi*2+yi])
			if want[[2]int{xi, yi}] != got {
				t.Fatalf("count[%d][%d] = %d, want %d", xi, yi, got, want[[2]int{xi, yi}])
			}
		}
	}
	if _, _, _, err := Histogram2D(smallVector(t, []float64{1}), y, 2, 2); err == nil {
		t.Fatal("mismatched lengths accepted")
	}
	if _, _, _, err := Histogram2D(x, y, 0, 2); err == nil {
		t.Fatal("zero bins accepted")
	}
}

// TestWeightedLinearRegression checks the weighted fit on an exact
// line and on a sample whose one wild point the weights silence: the
// unweighted fit chases the outlier, the weighted one recovers the
// line the other three points share.
func TestWeightedLinearRegression(t *testing.T) {
	design, err := core.FromFloats([]float64{1, 1, 1, 2, 1, 3, 1, 4}, 4, 2)
	if err != nil {
		t.Fatalf("design: %v", err)
	}
	y := smallVector(t, []float64{3, 5, 7, 9}) // y = 2x + 1 exactly
	w := smallVector(t, []float64{1, 1, 1, 1})
	res, err := WeightedLinearRegression(design, y, w)
	if err != nil {
		t.Fatalf("WeightedLinearRegression: %v", err)
	}
	if math.Abs(res.Coefficients[0]-1) > 1e-9 || math.Abs(res.Coefficients[1]-2) > 1e-9 {
		t.Fatalf("coefficients = %g, %g, want 1, 2", res.Coefficients[0], res.Coefficients[1])
	}
	// The fourth observation is the outlier now; heavy weights on the
	// clean points pull the fit back onto y = 2x + 1.
	y2 := smallVector(t, []float64{3, 5, 7, 50})
	w2 := smallVector(t, []float64{1000, 1000, 1000, 0.001})
	res2, err := WeightedLinearRegression(design, y2, w2)
	if err != nil {
		t.Fatalf("WeightedLinearRegression: %v", err)
	}
	if math.Abs(res2.Coefficients[1]-2) > 1e-3 {
		t.Fatalf("slope = %.6g, want 2 within the weights' pull", res2.Coefficients[1])
	}
	// Residuals read in the original units: the clean points sit on
	// the line again.
	if math.Abs(res2.Residuals[0]) > 1e-2 {
		t.Fatalf("first residual = %.6g, want a clean-point zero", res2.Residuals[0])
	}
	if _, err := WeightedLinearRegression(design, y, smallVector(t, []float64{1, 1, 1, -1})); err == nil {
		t.Fatal("negative weight accepted")
	}
}

// TestHistogram2DParallelCounting pins the paired sweep's split: past
// the per-worker floor the sample is cut into chunks, each worker
// counts into a private matrix and the merge adds every chunk into the
// total. The counts are exact integers, so the split must produce the
// matrix the single-worker sweep produces, cell for cell.
func TestHistogram2DParallelCounting(t *testing.T) {
	const n, xBins, yBins = 1 << 16, 16, 16
	xv := make([]float64, n)
	yv := make([]float64, n)
	for i := range n {
		xv[i] = float64((i*7919)%1000)/1000*4 - 2
		yv[i] = float64((i*104729)%1000)/1000*6 - 3
	}
	xv[0], xv[1] = -2, 2
	yv[0], yv[1] = -3, 3
	x := smallVector(t, xv)
	y := smallVector(t, yv)

	prev := engine.SetNumWorkers(4)
	defer engine.SetNumWorkers(prev)
	if w := engine.WorkersFor(n); w < 2 {
		t.Fatalf("the sweep did not split across workers: %d", w)
	}
	split, xEdges, yEdges, err := Histogram2D(x, y, xBins, yBins)
	if err != nil {
		t.Fatalf("Histogram2D across workers: %v", err)
	}
	engine.SetNumWorkers(1)
	serial, _, _, err := Histogram2D(x, y, xBins, yBins)
	if err != nil {
		t.Fatalf("Histogram2D on one worker: %v", err)
	}
	if split.Len() != xBins*yBins {
		t.Fatalf("the matrix holds %d cells, want %d", split.Len(), xBins*yBins)
	}
	got, want := split.RawInts(), serial.RawInts()
	total := int64(0)
	for i := range got {
		if got[i] != want[i] {
			t.Fatalf("cell %d = %d across workers, %d on one, want the same", i, got[i], want[i])
		}
		total += got[i]
	}
	if total != n {
		t.Fatalf("the counts total %d, want %d", total, n)
	}
	// The binning arithmetic, recomputed from the returned edges.
	xLo, xWidth := xEdges[0], xEdges[1]-xEdges[0]
	yLo, yWidth := yEdges[0], yEdges[1]-yEdges[0]
	ref := make([]int64, xBins*yBins)
	for i := range n {
		xi := int((xv[i] - xLo) / xWidth)
		if xi >= xBins {
			xi = xBins - 1
		}
		if xi < 0 {
			xi = 0
		}
		yi := int((yv[i] - yLo) / yWidth)
		if yi >= yBins {
			yi = yBins - 1
		}
		if yi < 0 {
			yi = 0
		}
		ref[xi*yBins+yi]++
	}
	for i := range xBins * yBins {
		if got[i] != ref[i] {
			t.Fatalf("cell %d = %d, the binning arithmetic gives %d", i, got[i], ref[i])
		}
	}
}

// refWindowExtreme is the retired window rescan: the seed walks to the
// first non-NaN sample and the rest folds with a strict comparison, so
// a tie keeps the earlier sample, the ±0 pair included, and an all-NaN
// window answers its last element. The monotonic deque must select the
// same sample, bit for bit.
func refWindowExtreme(src []float64, window int, greater bool) []float64 {
	out := make([]float64, len(src)-window+1)
	for i := range out {
		w := src[i : i+window]
		best, k := w[0], 1
		for math.IsNaN(best) && k < len(w) {
			best = w[k]
			k++
		}
		for _, v := range w[k:] {
			if greater {
				if v > best {
					best = v
				}
			} else {
				if v < best {
					best = v
				}
			}
		}
		out[i] = best
	}
	return out
}

// checkRollingRescan folds both window extrema over the series at every
// window and compares every output bit against the rescan.
func checkRollingRescan(t *testing.T, series []float64) {
	t.Helper()
	windows := make([]int, len(series))
	for i := range windows {
		windows[i] = i + 1
	}
	checkRollingRescanAt(t, series, windows...)
}

// checkRollingRescanAt folds both window extrema over the series at the
// given windows and compares every output bit against the rescan.
func checkRollingRescanAt(t *testing.T, series []float64, windows ...int) {
	t.Helper()
	for _, window := range windows {
		for _, greater := range []bool{true, false} {
			var got *core.Array
			var err error
			if greater {
				got, err = RollingMax(smallVector(t, series), window)
			} else {
				got, err = RollingMin(smallVector(t, series), window)
			}
			if err != nil {
				t.Fatalf("window %d over %v: %v", window, series, err)
			}
			want := refWindowExtreme(series, window, greater)
			if got.Len() != len(want) {
				t.Fatalf("window %d: %d outputs, want %d", window, got.Len(), len(want))
			}
			for i := range want {
				if gb, wb := math.Float64bits(got.FloatAt(i)), math.Float64bits(want[i]); gb != wb {
					// A long series is not worth printing: the index, the
					// two values and the window say which sample the fold
					// picked instead.
					if len(series) > 16 {
						t.Fatalf("window %d: [%d] = %v (%#x), the rescan gives %v (%#x) over a %d-sample series",
							window, i, got.FloatAt(i), gb, want[i], wb, len(series))
					}
					t.Fatalf("window %d over %v: [%d] = %v (%#x), the rescan gives %v (%#x)",
						window, series, i, got.FloatAt(i), gb, want[i], wb)
				}
			}
		}
	}
}

// TestRollingRescanTiesAndNaN pins the deque's selection rule against
// the retired window rescan on the samples where the rule is visible:
// a tie between +0 and −0 (the earlier sample wins, so the sign of the
// zero is part of the answer), NaN sharing a window with a number, an
// all-NaN window and the infinities.
func TestRollingRescanTiesAndNaN(t *testing.T) {
	negZero, posZero := math.Copysign(0, -1), 0.0
	checkRollingRescan(t, []float64{negZero, posZero})
	checkRollingRescan(t, []float64{posZero, negZero})
	checkRollingRescan(t, []float64{negZero, posZero, negZero, posZero})
	checkRollingRescan(t, []float64{negZero, posZero, math.NaN()})
	checkRollingRescan(t, []float64{math.NaN(), negZero, posZero})
	checkRollingRescan(t, []float64{posZero, math.NaN(), negZero})
	checkRollingRescan(t, []float64{math.NaN(), math.NaN(), math.NaN()})
	checkRollingRescan(t, []float64{math.Inf(-1), math.Inf(1), negZero, posZero, math.NaN()})
}

// TestRollingRescanCompaction pins the deque's buffer compaction: once
// the dead prefix reaches rollingDequeCompact the buffer is compacted,
// and the compaction must drop dead indices only, never the live entry
// at the deque's front. The series is monotone with a NaN every
// threshold indices: the NaN leaves the front one index ahead of the
// window's own oldest, which is exactly the state a compaction that
// copies from the front instead of from the dead prefix destroys, and
// the window's extreme then comes back from the wrong sample.
func TestRollingRescanCompaction(t *testing.T) {
	const n = 4 * rollingDequeCompact
	series := make([]float64, n)
	for i := range n {
		if i > 0 && i%rollingDequeCompact == 0 {
			series[i] = math.NaN()
		} else {
			series[i] = -float64(i)
		}
	}
	checkRollingRescanAt(t, series, 2, 3, rollingDequeCompact)
}

func TestMedianSliceMatchesSort(t *testing.T) {
	// The selection must answer what the sort answered. The pool holds no
	// NaN: Go orders NaN nowhere, so "the sorted middle" is not a
	// meaningful reference once one is present, and the medians that
	// matter (slopes, residuals, deviations) are finite by the time they
	// reach here. The zeros are in, because the two zeros are the one
	// pair that compares equal while differing in bits, and that is the
	// case the selection cannot pin.
	pool := []float64{0, math.Copysign(0, -1), -1, 1, 2, -2, 0.5, -0.5, 3, -3, 7, 7, 7}
	lens := []int{1, 2, 3, 4, 5, 9, 14, 25, 26, 27, 50, 100, 1000, 4097}
	for _, n := range lens {
		for shift := range len(pool) {
			vals := make([]float64, n)
			for i := range vals {
				vals[i] = pool[(i+shift)%len(pool)]
			}
			want := append([]float64(nil), vals...)
			slices.Sort(want)
			var expected float64
			if n%2 == 1 {
				expected = want[n/2]
			} else {
				expected = want[n/2-1]/2 + want[n/2]/2
			}
			got := medianSlice(vals)
			if n%2 == 1 {
				// An odd length has no tie ambiguity at the middle: the
				// value is unique, so the bits must agree outright.
				if math.Float64bits(got) != math.Float64bits(expected) {
					t.Errorf("n=%d shift=%d: median %v, sorted %v", n, shift, got, expected)
				}
			} else if got != expected {
				t.Errorf("n=%d shift=%d: median %v, sorted %v", n, shift, got, expected)
			}
			// The selection leaves a permutation of the input. The two
			// sorts used for the check are themselves unstable, so the
			// comparison is by value and the negative zeros are counted
			// separately: that is the one distinction a value comparison
			// cannot see.
			gotSorted := append([]float64(nil), vals...)
			slices.Sort(gotSorted)
			for i := range gotSorted {
				if gotSorted[i] != want[i] {
					t.Fatalf("n=%d shift=%d: elements changed at %d: %v against %v", n, shift, i, gotSorted[i], want[i])
				}
			}
			if negZero(gotSorted) != negZero(want) {
				t.Fatalf("n=%d shift=%d: negative zeros changed: %d against %d", n, shift, negZero(gotSorted), negZero(want))
			}
		}
	}
}

func TestMedianSliceNaNStaysInTheInput(t *testing.T) {
	// A NaN in the input is not the selection's business: the multiset
	// must survive it and the median must be a value that was there.
	vals := []float64{1, math.NaN(), 2, 3, 4}
	got := medianSlice(vals)
	if got == got {
		found := false
		for _, v := range []float64{1, 2, 3, 4} {
			if got == v {
				found = true
			}
		}
		if !found {
			t.Errorf("median %v is not an input value", got)
		}
	}
	count := 0
	for _, v := range vals {
		if math.IsNaN(v) {
			count++
		}
	}
	if count != 1 {
		t.Errorf("the selection lost the NaN: %d of them remain", count)
	}
}

func TestMedianSliceOnDistinctValues(t *testing.T) {
	// Distinct values have no tie ambiguity at all: the median must be the
	// exact order statistic on every length.
	for n := 1; n <= 200; n++ {
		vals := make([]float64, n)
		for i := range vals {
			vals[i] = float64((i*37)%101) - 50
		}
		want := append([]float64(nil), vals...)
		slices.Sort(want)
		expected := want[n/2]
		if n%2 == 0 {
			expected = want[n/2-1]/2 + want[n/2]/2
		}
		if got := medianSlice(vals); got != expected {
			t.Fatalf("n=%d: median %v, sorted %v", n, got, expected)
		}
	}
}

// negZero counts the negative zeros in a slice: the one distinction a
// value comparison cannot see.
func negZero(vals []float64) int {
	n := 0
	for _, v := range vals {
		if v == 0 && math.Signbit(v) {
			n++
		}
	}
	return n
}