tensor/stats/bench_test.go

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

package stats

import (
	"math"
	"testing"

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

// Benchmarks for the package's heavy paths: the O(n·m) kernel-density
// sweep, the sort-backed summaries, the regression and GLM normal
// equations and the incomplete-function machinery under the CDFs.

// benchSeries builds a deterministic float vector of length n mixing a
// few incommensurate frequencies, so sorts and moment sums see a
// realistic spread without a generator.
func benchSeries(n int) []float64 {
	vals := make([]float64, n)
	for i := range vals {
		x := float64(i)
		vals[i] = math.Sin(0.001*x)*10 + math.Sin(0.013*x)*3 + math.Cos(0.11*x) + float64(i%7)*0.125
	}
	return vals
}

func benchArray(b *testing.B, n int, shape ...int) *core.Array {
	b.Helper()
	if len(shape) == 0 {
		shape = []int{n}
	}
	a, err := core.FromFloats(benchSeries(n), shape...)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

func BenchmarkKernelDensity(b *testing.B) {
	sample := benchArray(b, 1000)
	points := benchArray(b, 256)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := KernelDensity(sample, 0.5, points); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkKernelDensitySilverman(b *testing.B) {
	sample := benchArray(b, 2000)
	points := benchArray(b, 512)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := KernelDensity(sample, 0, points); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkHistogram(b *testing.B) {
	a := benchArray(b, 100000)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := Histogram(a, 64); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkRollingMean(b *testing.B) {
	a := benchArray(b, 8192)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := RollingMean(a, 32); err != nil {
			b.Fatal(err)
		}
	}
}

// benchmarkRolling1M drives one rolling reduction over a million-sample
// series at the short and the long window, the pair the carried-update
// rescan trade-off is measured on.
func benchmarkRolling1M(b *testing.B, window int, sum bool) {
	b.Helper()
	a := benchArray(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		var err error
		if sum {
			_, err = RollingSum(a, window)
		} else {
			_, err = RollingMean(a, window)
		}
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkRollingSumWindow8(b *testing.B) {
	benchmarkRolling1M(b, 8, true)
}

func BenchmarkRollingSumWindow4096(b *testing.B) {
	benchmarkRolling1M(b, 4096, true)
}

func BenchmarkRollingMeanWindow8(b *testing.B) {
	benchmarkRolling1M(b, 8, false)
}

func BenchmarkRollingMeanWindow4096(b *testing.B) {
	benchmarkRolling1M(b, 4096, false)
}

func BenchmarkRollingMax(b *testing.B) {
	a := benchArray(b, 8192)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := RollingMax(a, 32); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkMedian(b *testing.B) {
	a := benchArray(b, 50000)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Median(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkCovarianceMatrix(b *testing.B) {
	a := benchArray(b, 1000*16, 1000, 16)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := CovarianceMatrix(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkLinearRegression(b *testing.B) {
	n, p := 2000, 8
	xv := make([]float64, n*p)
	yv := make([]float64, n)
	for r := range n {
		yv[r] = 0
		for c := range p {
			v := float64((r*31+c*17)%97)/97 - 0.5
			xv[r*p+c] = v
			yv[r] += float64(c+1) * v
		}
	}
	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 := LinearRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkLogisticRegression(b *testing.B) {
	n, p := 1000, 6
	xv := make([]float64, n*p)
	yv := make([]float64, n)
	for r := range n {
		eta := -1.5
		for c := range p {
			v := float64((r*13+c*29)%53)/53 - 0.5
			xv[r*p+c] = v
			eta += float64(c) * v
		}
		if r%2 == 0 {
			eta += 0.5
		}
		if eta > 0 {
			yv[r] = 1
		}
	}
	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 := LogisticRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkGammaCDF(b *testing.B) {
	b.ReportAllocs()
	for b.Loop() {
		if _, err := GammaCDF(2.5, 3.0, 1.2); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkBetaIncomplete(b *testing.B) {
	b.ReportAllocs()
	for b.Loop() {
		if _, err := BetaIncomplete(0.4, 2.5, 3.5); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkStudentTQuantile(b *testing.B) {
	b.ReportAllocs()
	for b.Loop() {
		if _, err := StudentTQuantile(0.975, 12); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNormalQuantile(b *testing.B) {
	b.ReportAllocs()
	for b.Loop() {
		if _, err := NormalQuantile(0.975); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkGammaQuantile(b *testing.B) {
	b.ReportAllocs()
	for b.Loop() {
		if _, err := GammaQuantile(0.99, 3, 1.2); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkKolmogorovSmirnovTest(b *testing.B) {
	a := benchArray(b, 4000)
	c := benchArray(b, 5000)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := KolmogorovSmirnovTest(a, c); err != nil {
			b.Fatal(err)
		}
	}
}