feat: initial release

Assisted-by: GLM 5.3 Flash
2026-09-03 10:00:00 +02:00
commit af4ee19703
617 changed files with 191195 additions and 0 deletions
@@ -0,0 +1,318 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package linalg
+
+import (
+	"fmt"
+	"testing"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// Benchmarks for the sparse surface: the matrix-vector and matrix-matrix
+// products, the coordinate-to-compressed construction, and the Krylov
+// solvers that drive them. Every input is a fixed literal pattern, so a
+// run is reproducible and two builds compare like with like.
+
+// benchSink keeps a kernel's result reachable so the compiler cannot
+// elide the call it came from.
+var (
+	benchSinkCSR *SparseCSR
+	benchSinkCSC *SparseCSC
+	benchSinkArr *core.Array
+)
+
+// laplacianTriples returns the flat (row, col, value) triples of the
+// n×n symmetric banded Laplacian of the given half-bandwidth: 2·width
+// on the diagonal and -1 at every offset up to width. It is symmetric
+// positive definite by strict diagonal dominance, which is what the
+// Lanczos and conjugate-gradient benchmarks need, and its non-zeros per
+// row are 2·width+1 at every row.
+func laplacianTriples(n, width int) []float64 {
+	tri := make([]float64, 0, n*(2*width+1)*3)
+	for i := range n {
+		tri = append(tri, float64(i), float64(i), float64(2*width))
+		for d := 1; d <= width; d++ {
+			j := i + d
+			if j >= n {
+				break
+			}
+			tri = append(tri, float64(i), float64(j), -1)
+			tri = append(tri, float64(j), float64(i), -1)
+		}
+	}
+	return tri
+}
+
+// duplicateTriples returns the triples of a matrix whose coordinates
+// repeat: entry k carries row (k·13) mod n and column (k·37) mod n, so
+// each coordinate appears about len/n times and the construction path
+// has to merge duplicates rather than only sort distinct pairs. The
+// values step along a literal modular sequence, so no merge cancels by
+// accident.
+func duplicateTriples(n, len_ int) []float64 {
+	tri := make([]float64, 0, len_*3)
+	for k := range len_ {
+		row := (k * 13) % n
+		col := (k * 37) % n
+		v := float64((k*7)%23) - 11
+		if v == 0 {
+			v = 1.5
+		}
+		tri = append(tri, float64(row), float64(col), v)
+	}
+	return tri
+}
+
+// benchCOO builds a sparse matrix from flat (row, col, value) triples.
+func benchCOO(b *testing.B, rows, cols int, tri []float64) *core.SparseCOO {
+	b.Helper()
+	nnz := len(tri) / 3
+	idx := make([]int64, 0, nnz*2)
+	vals := make([]float64, 0, nnz)
+	for i := range nnz {
+		idx = append(idx, int64(tri[i*3]), int64(tri[i*3+1]))
+		vals = append(vals, tri[i*3+2])
+	}
+	indices, err := core.FromInts(idx, nnz, 2)
+	if err != nil {
+		b.Fatalf("FromInts: %v", err)
+	}
+	values, err := core.FromFloats(vals, nnz)
+	if err != nil {
+		b.Fatalf("FromFloats: %v", err)
+	}
+	coo, err := core.NewSparseCOO(indices, values, []int{rows, cols})
+	if err != nil {
+		b.Fatalf("NewSparseCOO: %v", err)
+	}
+	return coo
+}
+
+// benchVector returns a dense vector of length n whose entries come
+// from a literal modular sequence, deliberately without zeros so a
+// scaled product cannot skip work.
+func benchVector(b *testing.B, n int) *core.Array {
+	b.Helper()
+	v := make([]float64, n)
+	for i := range n {
+		v[i] = float64((i*11)%17)/8 - 1
+		if v[i] == 0 {
+			v[i] = 0.25
+		}
+	}
+	arr, err := core.FromFloats(v, n)
+	if err != nil {
+		b.Fatalf("FromFloats: %v", err)
+	}
+	return arr
+}
+
+func BenchmarkSparseCSRMatVec(b *testing.B) {
+	for _, n := range []int{256, 131072} {
+		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+			coo := benchCOO(b, n, n, laplacianTriples(n, 1))
+			csr, err := CSRFromCOO(coo)
+			if err != nil {
+				b.Fatalf("CSRFromCOO: %v", err)
+			}
+			x := benchVector(b, n)
+			for b.Loop() {
+				benchSinkArr, _ = csr.MatVec(x)
+			}
+		})
+	}
+}
+
+func BenchmarkSparseCSCMatVec(b *testing.B) {
+	for _, n := range []int{256, 131072} {
+		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+			coo := benchCOO(b, n, n, laplacianTriples(n, 1))
+			csc, err := CSCFromCOO(coo)
+			if err != nil {
+				b.Fatalf("CSCFromCOO: %v", err)
+			}
+			x := benchVector(b, n)
+			for b.Loop() {
+				benchSinkArr, _ = csc.MatVec(x)
+			}
+		})
+	}
+}
+
+func BenchmarkSparseCSRMatMulSparse(b *testing.B) {
+	const n = 256
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		coo := benchCOO(b, n, n, laplacianTriples(n, 4))
+		csr, err := CSRFromCOO(coo)
+		if err != nil {
+			b.Fatalf("CSRFromCOO: %v", err)
+		}
+		for b.Loop() {
+			benchSinkCSR, _ = csr.MatMulSparse(csr)
+		}
+	})
+}
+
+func BenchmarkSparseCSRMatMulDense(b *testing.B) {
+	const n = 512
+	const cols = 4
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		coo := benchCOO(b, n, n, laplacianTriples(n, 2))
+		csr, err := CSRFromCOO(coo)
+		if err != nil {
+			b.Fatalf("CSRFromCOO: %v", err)
+		}
+		v := make([]float64, n*cols)
+		for i := range v {
+			v[i] = float64((i*11)%17)/8 - 1
+		}
+		x, err := core.FromFloats(v, n, cols)
+		if err != nil {
+			b.Fatalf("FromFloats: %v", err)
+		}
+		for b.Loop() {
+			benchSinkArr, _ = csr.MatMulDense(x)
+		}
+	})
+}
+
+func BenchmarkCSRFromCOODuplicates(b *testing.B) {
+	const n = 1500
+	tri := duplicateTriples(n, 6000)
+	b.Run("nnz=6000", func(b *testing.B) {
+		coo := benchCOO(b, n, n, tri)
+		for b.Loop() {
+			benchSinkCSR, _ = CSRFromCOO(coo)
+		}
+	})
+}
+
+func BenchmarkCSCFromCOODuplicates(b *testing.B) {
+	const n = 1500
+	tri := duplicateTriples(n, 6000)
+	b.Run("nnz=6000", func(b *testing.B) {
+		coo := benchCOO(b, n, n, tri)
+		for b.Loop() {
+			benchSinkCSC, _ = CSCFromCOO(coo)
+		}
+	})
+}
+
+// BenchmarkSpEigenBanded measures the real symmetric Krylov solve end
+// to end: the coordinate construction, the banded matrix-vector
+// products and the full reorthogonalisation of every Lanczos step.
+func BenchmarkSpEigenBanded(b *testing.B) {
+	const n = 800
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		coo := benchCOO(b, n, n, laplacianTriples(n, 1))
+		for b.Loop() {
+			benchSinkArr, _, _ = SpEigen(coo, 6, core.NewGenerator(7))
+		}
+	})
+}
+
+// BenchmarkSpEigenGeneralBanded measures the Arnoldi solve on the same
+// matrix: one explicit orthogonalisation per basis column instead of
+// the three-term recurrence.
+func BenchmarkSpEigenGeneralBanded(b *testing.B) {
+	const n = 400
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		coo := benchCOO(b, n, n, laplacianTriples(n, 1))
+		for b.Loop() {
+			benchSinkArr, _, _ = SpEigenGeneral(coo, 4, core.NewGenerator(7))
+		}
+	})
+}
+
+func BenchmarkSpEigenComplexBanded(b *testing.B) {
+	const n = 400
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		idx := make([]int64, 0, n*6)
+		vals := make([]complex128, 0, n*3)
+		for i := range n {
+			idx = append(idx, int64(i), int64(i))
+			vals = append(vals, 2+0i)
+			if i+1 < n {
+				idx = append(idx, int64(i), int64(i+1))
+				vals = append(vals, 1i)
+				idx = append(idx, int64(i+1), int64(i))
+				vals = append(vals, -1i)
+			}
+		}
+		indices, err := core.FromInts(idx, len(vals), 2)
+		if err != nil {
+			b.Fatalf("FromInts: %v", err)
+		}
+		values, err := core.FromComplexes(vals, len(vals))
+		if err != nil {
+			b.Fatalf("FromComplexes: %v", err)
+		}
+		coo, err := core.NewSparseCOO(indices, values, []int{n, n})
+		if err != nil {
+			b.Fatalf("NewSparseCOO: %v", err)
+		}
+		for b.Loop() {
+			benchSinkArr, _, _ = SpEigenComplex(coo, 4, core.NewGenerator(7))
+		}
+	})
+}
+
+// BenchmarkSpExpApplyBanded measures the Krylov projection of the
+// matrix exponential action, which shares the Lanczos kernel with the
+// symmetric eigensolver.
+func BenchmarkSpExpApplyBanded(b *testing.B) {
+	const n = 400
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		coo := benchCOO(b, n, n, laplacianTriples(n, 1))
+		v := benchVector(b, n)
+		for b.Loop() {
+			benchSinkArr, _ = SpExpApply(coo, v, 0)
+		}
+	})
+}
+
+// BenchmarkSpSolveComplexCGBanded measures the conjugate-gradient solve
+// whose every step is a complex sparse matrix-vector product with a
+// Jacobi preconditioner.
+func BenchmarkSpSolveComplexCGBanded(b *testing.B) {
+	const n = 256
+	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+		idx := make([]int64, 0, n*6)
+		vals := make([]complex128, 0, n*3)
+		for i := range n {
+			idx = append(idx, int64(i), int64(i))
+			vals = append(vals, 4+0i)
+			if i+1 < n {
+				idx = append(idx, int64(i), int64(i+1))
+				vals = append(vals, 1i)
+				idx = append(idx, int64(i+1), int64(i))
+				vals = append(vals, -1i)
+			}
+		}
+		indices, err := core.FromInts(idx, len(vals), 2)
+		if err != nil {
+			b.Fatalf("FromInts: %v", err)
+		}
+		values, err := core.FromComplexes(vals, len(vals))
+		if err != nil {
+			b.Fatalf("FromComplexes: %v", err)
+		}
+		coo, err := core.NewSparseCOO(indices, values, []int{n, n})
+		if err != nil {
+			b.Fatalf("NewSparseCOO: %v", err)
+		}
+		ones := make([]complex128, n)
+		for i := range ones {
+			ones[i] = 1
+		}
+		rhs, err := core.FromComplexes(ones, n)
+		if err != nil {
+			b.Fatalf("FromComplexes: %v", err)
+		}
+		for b.Loop() {
+			benchSinkArr, _ = SpSolveComplexCG(coo, rhs, 1e-12, 400)
+		}
+	})
+}