feat: initial release

Assisted-by: GLM 5.3 Flash

linalg/sparsesolve.go

@@ -0,0 +1,207 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package linalg
+
+import (
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// Sparse symmetric positive-definite solve. The dense `Solve` factors
+// A once and back-substitutes, which costs O(n³) time and O(n²) memory
+// and is the right answer for a modest, dense system. When A is large
+// and sparse that factorisation is both unaffordable and wasteful: the
+// fill-in alone can exceed the memory the non-zeros needed. SpSolve
+// runs the preconditioned conjugate gradient instead, whose every step
+// is one sparse-matrix-vector product plus a handful of vector
+// operations, so the cost tracks the non-zero count rather than the
+// dimension.
+//
+// Conjugate gradient applies only to a symmetric positive-definite
+// matrix, and it is not a drop-in for `Solve`: the answer is an
+// approximation within a stated residual tolerance, not the exact
+// solution a factorisation returns. Symmetry is verified the same way
+// `SpEigen` verifies it. Positive-definiteness cannot be checked
+// cheaply up front, so it is caught in flight: a non-positive search
+// curvature is the matrix declaring itself indefinite.
+
+// spSolveTol is the default relative residual when the caller passes
+// zero. The iteration default is the dimension rather than a constant:
+// an exact arithmetic conjugate gradient terminates in at most n
+// steps, so a larger budget cannot buy convergence, only round-off
+// work.
+const spSolveTol = 1e-10
+
+// checkSparseSquare validates the inputs every real sparse solver
+// shares: a is a square 2-D matrix of non-complex dtype and, when b is
+// not nil, b is a real right-hand side vector of the matrix dimension.
+// It returns the dimension n.
+func checkSparseSquare(name string, a *core.SparseCOO, b *core.Array) (int, error) {
+	if a.Values.Dtype() == core.Complex {
+		return 0, base.Errf("%s: complex sparse matrices are not supported", name)
+	}
+	if len(a.Shape) != 2 || a.Shape[0] != a.Shape[1] {
+		return 0, base.Errf("%s: needs a square 2-D sparse matrix, got shape %v", name, a.Shape)
+	}
+	n := a.Shape[0]
+	if n == 0 {
+		return 0, base.Errf("%s: zero-sized matrix, got shape %v", name, a.Shape)
+	}
+	if b != nil {
+		if b.NDim() != 1 || b.Shape()[0] != n {
+			return 0, base.Errf("%s: right-hand side must be a vector of length %d, got shape %s",
+				name, n, base.ShapeText(b.Shape()))
+		}
+		if b.Dtype() == core.Complex {
+			return 0, base.Errf("%s: complex right-hand side is not supported", name)
+		}
+	}
+	return n, nil
+}
+
+// pickPreconditioner resolves the optional ILU argument of the Krylov
+// solvers, falling back to the Jacobi diagonal of c when none was
+// given. An ILU built for another dimension is refused: applying it
+// would index past its vectors, or precondition with the wrong
+// factorisation without a word.
+func pickPreconditioner(name string, c *sparseCSR, precond []*SparseILU) (*SparseILU, []float64, error) {
+	if len(precond) > 0 {
+		if precond[0] == nil {
+			return nil, nil, base.Errf("%s: the preconditioner is nil", name)
+		}
+		if precond[0].n != c.n {
+			return nil, nil, base.Errf("%s: the preconditioner was built for dimension %d, the system is %d",
+				name, precond[0].n, c.n)
+		}
+		return precond[0], nil, nil
+	}
+	diag, err := c.diagonal(name)
+	if err != nil {
+		return nil, nil, err
+	}
+	return nil, diag, nil
+}
+
+// vectorF64 flattens a rank-1 array into a plain float64 working
+// vector of the expected length n.
+func vectorF64(b *core.Array, n int) []float64 {
+	r := make([]float64, n)
+	for i := range n {
+		r[i] = b.FloatAt(i)
+	}
+	return r
+}
+
+// SpSolve returns the vector x solving A·x = b for a real symmetric
+// positive-definite sparse A, by preconditioned conjugate gradient
+// with a Jacobi (diagonal) preconditioner. The right-hand side b is a
+// rank-1 vector of length n.
+//
+// The stopping rule is the relative residual ‖b − A·x‖₂ ≤ tol·‖b‖₂.
+// A tol ≤ 0 uses 1e-10, and maxIter ≤ 0 uses n steps. x starts at
+// zero, so the initial residual is b itself and the loop needs no
+// separate first matvec. The default preconditioner is Jacobi
+// scaling; an ILU(0) factorisation from NewSparseILU may be passed to
+// replace it.
+//
+// An unconverged solve is an error naming the residual achieved, with
+// no estimate returned, rather than a silent approximation.
+// The matrix must be symmetric, and a zero diagonal entry is refused:
+// the Jacobi preconditioner divides by it, and for a symmetric
+// positive-definite matrix the diagonal is necessarily positive.
+func SpSolve(a *core.SparseCOO, b *core.Array, tol float64, maxIter int, precond ...*SparseILU) (*core.Array, error) {
+	const name = "SpSolve"
+	n, err := checkSparseSquare(name, a, b)
+	if err != nil {
+		return nil, err
+	}
+	c, err := symmetricCSR(a, name)
+	if err != nil {
+		return nil, err
+	}
+	ilu, diag, err := pickPreconditioner(name, c, precond)
+	if err != nil {
+		return nil, err
+	}
+
+	if tol <= 0 {
+		tol = spSolveTol
+	}
+	if maxIter <= 0 {
+		maxIter = n
+	}
+
+	// x starts at zero, so r = b − A·x is just b and the loop needs no
+	// opening matvec.
+	r := vectorF64(b, n)
+	bNorm := norm2F64(r)
+	if bNorm == 0 {
+		// The exact solution of A·0 = 0 is the zero vector; a relative
+		// stopping rule would otherwise divide by zero.
+		return core.Zeros(core.Float, n)
+	}
+
+	z := make([]float64, n)
+	iluPrecondition(z, r, ilu, diag)
+	p := append([]float64(nil), z...)
+	rz := dotF64(r, z)
+	x := make([]float64, n)
+	ap := make([]float64, n)
+
+	for iter := range maxIter {
+		c.matVec(p, ap)
+		pAp := dotF64(p, ap)
+		if !finiteF64(pAp) || pAp <= 0 {
+			// A non-positive or non-finite curvature means the matrix
+			// is not positive-definite, or has overflowed the
+			// representable range; either way conjugate gradient
+			// cannot proceed. A non-finite value must be an error,
+			// not a NaN that slips through the <= 0 test.
+			return nil, base.Errf("%s: matrix is not positive-definite (curvature %.3g at step %d)",
+				name, pAp, iter+1)
+		}
+		alpha := rz / pAp
+		for i := range n {
+			x[i] += alpha * p[i]
+			r[i] -= alpha * ap[i]
+		}
+		// norm2F64 skips NaN entries, so an all-NaN residual reads as a
+		// zero norm; a non-finite state is a breakdown before the
+		// convergence test can mistake it for an exact solve.
+		if !vecFinite(r) {
+			return nil, base.Errf("%s: non-finite residual at step %d", name, iter+1)
+		}
+		if norm2F64(r) <= tol*bNorm {
+			return floatsToArray(x, []int{n}), nil
+		}
+		iluPrecondition(z, r, ilu, diag)
+		rzNext := dotF64(r, z)
+		if rz == 0 {
+			return nil, base.Errf("%s: breakdown at step %d (preconditioned residual vanished)", name, iter+1)
+		}
+		beta := rzNext / rz
+		for i := range n {
+			p[i] = z[i] + beta*p[i]
+		}
+		rz = rzNext
+	}
+	return nil, base.Errf("%s: no convergence in %d steps, residual %.3g (tolerance %.3g)",
+		name, maxIter, norm2F64(r), tol*bNorm)
+}
+
+// diagonal returns the main diagonal, which the Jacobi preconditioner
+// divides by. A missing or zero entry makes the preconditioner
+// undefined and, for a symmetric positive-definite matrix, signals
+// that the input is not one.
+func (c *sparseCSR) diagonal(name string) ([]float64, error) {
+	d := make([]float64, c.n)
+	for i := range c.n {
+		v, ok := c.at(i, i)
+		if !ok || v == 0 {
+			return nil, base.Errf("%s: zero or missing diagonal entry at %d", name, i)
+		}
+		d[i] = v
+	}
+	return d, nil
+}