tensor/linalg/example_test.go

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

package linalg_test

// Runnable examples for the flagship workflows of the package. `go
// test` executes them against the Output comments below, so the
// documentation cannot drift from the code.

import (
	"fmt"
	"log"
	"math"

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

// A dense solve, and the same matrix factored once for reuse.
func ExampleSolve() {
	a, err := tensor.FromFloats([]float64{4, 1, 1, 3}, 2, 2)
	if err != nil {
		log.Fatal(err)
	}
	b, err := tensor.FromFloats([]float64{1, 2}, 2)
	if err != nil {
		log.Fatal(err)
	}
	x, err := linalg.Solve(a, b)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("x = %.4f %.4f\n", x.FloatAt(0), x.FloatAt(1))

	// The same matrix through Cholesky: a = L·Lᵀ, one factor for any
	// number of right-hand sides.
	l, err := linalg.Cholesky(a)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("L = [%.1f %.1f; %.1f %.4f]\n",
		l.FloatAt(0), l.FloatAt(1), l.FloatAt(2), l.FloatAt(3))
	// Output:
	// x = 0.0909 0.6364
	// L = [2.0 0.0; 0.5 1.6583]
}

// A symmetric eigenproblem: ascending eigenvalues and orthonormal
// eigenvector columns.
func ExampleEigen() {
	a, err := tensor.FromFloats([]float64{1, 2, 2, 1}, 2, 2)
	if err != nil {
		log.Fatal(err)
	}
	values, vectors, err := linalg.Eigen(a)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("values = %.4f %.4f\n", values.FloatAt(0), values.FloatAt(1))

	// The sign of an eigenvector is arbitrary, so the magnitudes are
	// what a stable example can print. Here |V[i][j]| = 1/√2.
	fmt.Printf("|V| = [%.4f %.4f; %.4f %.4f]\n",
		math.Abs(vectors.FloatAt(0)), math.Abs(vectors.FloatAt(1)),
		math.Abs(vectors.FloatAt(2)), math.Abs(vectors.FloatAt(3)))
	// Output:
	// values = -1.0000 3.0000
	// |V| = [0.7071 0.7071; 0.7071 0.7071]
}

// A matrix function: the exponential of a rotation generator is the
// rotation itself.
func ExampleMatrixExp() {
	// A = [[0, -1], [1, 0]] generates a quarter-turn rotation rate, so
	// exp(A) = [[cos 1, -sin 1], [sin 1, cos 1]].
	a, err := tensor.FromFloats([]float64{0, -1, 1, 0}, 2, 2)
	if err != nil {
		log.Fatal(err)
	}
	e, err := linalg.MatrixExp(a)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%.4f %.4f %.4f %.4f\n",
		e.FloatAt(0), e.FloatAt(1), e.FloatAt(2), e.FloatAt(3))
	// Output: 0.5403 -0.8415 0.8415 0.5403
}

// A sparse symmetric positive definite solve by conjugate gradient.
func ExampleSpSolve() {
	a := laplacian1D()
	// The right-hand side A·1, so the exact solution is the vector of
	// ones and the printed digits show the residual the iteration left.
	b, err := tensor.FromFloats([]float64{1, 0, 0, 1}, 4)
	if err != nil {
		log.Fatal(err)
	}
	x, err := linalg.SpSolve(a, b, 0, 0)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%.4f %.4f %.4f %.4f\n",
		x.FloatAt(0), x.FloatAt(1), x.FloatAt(2), x.FloatAt(3))
	// Output: 1.0000 1.0000 1.0000 1.0000
}

// A sparse direct factorisation under a fill-reducing ordering.
func ExampleNewSparseCholesky() {
	a := laplacian1D()
	b, err := tensor.FromFloats([]float64{1, 0, 0, 1}, 4)
	if err != nil {
		log.Fatal(err)
	}
	f, err := linalg.NewSparseCholesky(a, linalg.SparseOrderingReverseCuthillMcKee)
	if err != nil {
		log.Fatal(err)
	}
	x, err := f.Solve(b)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%.4f %.4f %.4f %.4f\n",
		x.FloatAt(0), x.FloatAt(1), x.FloatAt(2), x.FloatAt(3))

	// The factor of a tridiagonal matrix keeps the band: 4 diagonal
	// entries and 3 subdiagonal ones.
	fmt.Println("factor non-zeros:", f.NNZ())
	// Output:
	// 1.0000 1.0000 1.0000 1.0000
	// factor non-zeros: 7
}

// A natural cubic spline through samples of a straight line.
func ExampleNewCubicSpline() {
	xs, err := tensor.FromFloats([]float64{0, 1, 2}, 3)
	if err != nil {
		log.Fatal(err)
	}
	ys, err := tensor.FromFloats([]float64{1, 3, 5}, 3)
	if err != nil {
		log.Fatal(err)
	}
	s, err := linalg.NewCubicSpline(xs, ys)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%.4f %.4f\n", s.At(0.5), s.At(1.5))

	// Extrapolation has no boundary condition here, so it is undefined.
	fmt.Printf("%.4f\n", s.At(2.5))
	// Output:
	// 2.0000 4.0000
	// NaN
}

// A sequence of element-wise steps read as one expression.
func ExamplePipe() {
	a, err := tensor.FromFloats([]float64{1, 2, 3}, 3)
	if err != nil {
		log.Fatal(err)
	}
	out, err := linalg.Pipe(a).AddF(1).Sqrt().MulF(2).Result()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%.4f %.4f %.4f\n", out.FloatAt(0), out.FloatAt(1), out.FloatAt(2))

	// A step that fails is recorded, and the steps after it are no-ops:
	// adding a length-2 array to a length-3 one is a shape error.
	short, err := tensor.FromFloats([]float64{1, 2}, 2)
	if err != nil {
		log.Fatal(err)
	}
	_, err = linalg.Pipe(a).Add(short).MulF(2).Result()
	fmt.Println("error:", err != nil)
	// Output:
	// 2.8284 3.4641 4.0000
	// error: true
}

// laplacian1D builds the 4×4 second-difference matrix with -1 on both
// off-diagonals and 2 on the diagonal, as a coordinate matrix.
func laplacian1D() *tensor.SparseCOO {
	indices, err := tensor.FromInts([]int64{
		0, 0, 1, 0, 0, 1, 1, 1,
		2, 1, 1, 2, 2, 2, 3, 2,
		2, 3, 3, 3,
	}, 10, 2)
	if err != nil {
		panic(err)
	}
	values, err := tensor.FromFloats([]float64{
		2, -1, -1, 2, -1, -1, 2, -1, -1, 2,
	}, 10)
	if err != nil {
		panic(err)
	}
	a, err := tensor.NewSparseCOO(indices, values, []int{4, 4})
	if err != nil {
		panic(err)
	}
	return a
}