tensor/stats/example_test.go

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

package stats_test

// The godoc examples: the flagship workflows of the package as
// runnable, checked snippets. pkg.go.dev renders them beside the API,
// and `go test` executes them, so the documentation cannot rot.

import (
	"fmt"
	"log"

	tensor "sourcedock.dev/petrbalvin/tensor"
	"sourcedock.dev/petrbalvin/tensor/stats"
)

// The normal quantile inverts NormalCDF: these are the two-sided
// critical values at the 2 percent and the 10 percent level.
func ExampleNormalQuantile() {
	for _, q := range []float64{0.01, 0.05, 0.95, 0.99} {
		z, err := stats.NormalQuantile(q)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%.4f\n", z)
	}
	// Output:
	// -2.3263
	// -1.6449
	// 1.6449
	// 2.3263
}

// Welch's t-test compares two independent samples without assuming
// equal variances, and reports the two-sided p-value.
func ExampleWelchTTest() {
	a, _ := tensor.FromFloats([]float64{5.1, 4.9, 5.4, 5.0, 5.3}, 5)
	b, _ := tensor.FromFloats([]float64{6.2, 6.0, 5.8, 6.3, 6.1}, 5)
	t, df, p, err := stats.WelchTTest(a, b)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("t = %.3f, df = %.2f, p = %.2e\n", t, df, p)
	// Output:
	// t = -7.431, df = 7.96, p = 7.61e-05
}

// Ordinary least squares with the full classical inference: the
// coefficients, their standard errors, R² and the model F test.
func ExampleLinearRegression() {
	// y = 4 + 3x over x = 0..5, the intercept carried as the constant
	// first column, as every regression entry point of the package
	// expects it.
	design, _ := tensor.FromFloats([]float64{
		1, 0, 1, 1, 1, 2, 1, 3, 1, 4, 1, 5,
	}, 6, 2)
	y, _ := tensor.FromFloats([]float64{4.1, 7.0, 9.9, 13.2, 15.9, 19.1}, 6)
	fit, err := stats.LinearRegression(design, y)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("intercept %.3f ± %.3f (p = %.2e)\n",
		fit.Coefficients[0], fit.StandardErrors[0], fit.PValues[0])
	fmt.Printf("slope     %.3f ± %.3f (p = %.2e)\n",
		fit.Coefficients[1], fit.StandardErrors[1], fit.PValues[1])
	fmt.Printf("R² %.4f, F = %.1f on (%d, %d) df\n",
		fit.RSquared, fit.FStatistic, fit.DModel, fit.DResidual)
	// Output:
	// intercept 4.033 ± 0.098 (p = 2.08e-06)
	// slope     3.000 ± 0.032 (p = 8.12e-08)
	// R² 0.9995, F = 8590.9 on (1, 4) df
}

// Logistic regression fits a binary response by maximum likelihood
// through the logit link and reports Wald inference.
func ExampleLogisticRegression() {
	design, _ := tensor.FromFloats([]float64{
		1, 1, 1, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 7, 1, 8, 1, 9, 1, 10,
	}, 10, 2)
	y, _ := tensor.FromFloats([]float64{0, 0, 1, 0, 1, 1, 0, 1, 1, 1}, 10)
	fit, err := stats.LogisticRegression(design, y)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("intercept %.4f (SE %.4f)\n", fit.Coefficients[0], fit.StandardErrors[0])
	fmt.Printf("slope     %.4f (SE %.4f)\n", fit.Coefficients[1], fit.StandardErrors[1])
	fmt.Printf("log likelihood %.4f, %d Newton steps\n",
		fit.LogLikelihood, fit.Iterations)
	// Output:
	// intercept -2.2903 (SE 1.7992)
	// slope     0.5279 (SE 0.3377)
	// log likelihood -4.9014, 6 Newton steps
}

// PCA decomposes a correlated cloud onto its principal components and
// whitens it to unit covariance.
func ExamplePCA() {
	data, _ := tensor.FromFloats([]float64{
		1, 2.1, 2, 3.9, 3, 6.2, 4, 7.8, 5, 10.1,
	}, 5, 2)
	fit, err := stats.PCA(data)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("explained variance ratio: %.4f %.4f\n",
		fit.ExplainedVarianceRatio[0], fit.ExplainedVarianceRatio[1])
	whitened, err := fit.Whiten(data)
	if err != nil {
		log.Fatal(err)
	}
	back, err := fit.Unwhiten(whitened)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("whitened first row: %.4f %.4f\n",
		whitened.FloatAt(0), whitened.FloatAt(1))
	fmt.Printf("unwhitened recovers: %.4f %.4f\n", back.FloatAt(0), back.FloatAt(1))
	// Output:
	// explained variance ratio: 0.9996 0.0004
	// whitened first row: -1.2486 -0.4190
	// unwhitened recovers: 1.0000 2.1000
}

// KMeans partitions a sample into k clusters, deterministically for a
// given generator state.
func ExampleKMeans() {
	data, _ := tensor.FromFloats([]float64{
		0, 0, 0.2, 0, 0, 0.2,
		9, 9, 9.2, 9, 9, 9.2,
	}, 6, 2)
	fit, err := stats.KMeans(tensor.NewGenerator(7), data, 2)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("labels: %v\n", fit.Labels)
	fmt.Printf("inertia: %.4f\n", fit.Inertia)
	// Output:
	// labels: [0 0 0 1 1 1]
	// inertia: 0.1067
}

// KernelDensity smooths a sample into a continuous density; a
// non-positive bandwidth asks for Silverman's rule.
func ExampleKernelDensity() {
	sample, _ := tensor.FromFloats([]float64{-1, 0, 0.5, 1, 1.5}, 5)
	points, _ := tensor.FromFloats([]float64{-1, 0, 1}, 3)
	density, err := stats.KernelDensity(sample, 0, points)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%.4f %.4f %.4f\n",
		density.FloatAt(0), density.FloatAt(1), density.FloatAt(2))
	// Output:
	// 0.1852 0.3018 0.3772
}