feat: initial release

Assisted-by: GLM 5.3 Flash

linalg/sparsecholesky.go

@@ -0,0 +1,524 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package linalg
+
+import (
+	"cmp"
+	"math"
+	"slices"
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// The sparse Cholesky factorisation: the elimination tree (Liu's
+// construction with path compression) discovers the fill of every row
+// cheaply, and the row-oriented elimination that follows is the
+// up-looking form of the direct-methods standard (George and Liu;
+// Davis). The factor exists only for symmetric positive definite
+// matrices, where it is the direct solver of choice: no pivoting, no
+// fill beyond the pattern the ordering decides.
+
+// SparseOrdering selects the fill-reducing permutation a sparse
+// factorisation applies before it eliminates.
+type SparseOrdering int
+
+const (
+	// SparseOrderingNatural eliminates rows and columns in stored
+	// order. Only right for matrices whose pattern is already dense
+	// near the diagonal, and the reference point every ordering is
+	// measured against.
+	SparseOrderingNatural SparseOrdering = iota
+	// SparseOrderingReverseCuthillMcKee orders every component in
+	// reverse breadth first order from a pseudo-peripheral start: the
+	// classic choice for mesh-shaped patterns, where it keeps the
+	// factor banded and the fill small.
+	SparseOrderingReverseCuthillMcKee
+	// SparseOrderingMinimumDegree eliminates the vertex with the
+	// fewest remaining neighbours at every step, absorbing each
+	// eliminated pattern into its neighbours: the stronger choice on
+	// irregular patterns, where a bandwidth order leaves the factor
+	// dense in the middle.
+	SparseOrderingMinimumDegree
+)
+
+// rowValuePair carries a factor entry through a column sort: the row
+// index travels with its value.
+type rowValuePair struct {
+	row int
+	val float64
+}
+
+// SparseCholesky carries the factorisation of a symmetric positive
+// definite matrix: L with A permuted to the elimination order, so
+// that P·A·Pᵀ = L·Lᵀ. One factorisation solves any number of right
+// hand sides.
+type SparseCholesky struct {
+	// perm[k] is the original index of the row eliminated k-th;
+	// inversePerm maps an original index back to its position.
+	perm        []int
+	inversePerm []int
+	n           int
+	// L in CSC form: column j holds the off-diagonal rows i > j in
+	// ascending order, and diag carries L[j][j].
+	colStart []int
+	rowIdx   []int
+	values   []float64
+	diag     []float64
+}
+
+// NewSparseCholesky factors a symmetric positive definite matrix. The
+// lower triangle defines the matrix: a stored upper triangle entry is
+// refused unless its lower counterpart is stored with exactly the
+// same value, so a silently asymmetric input cannot be factored. The
+// ordering argument picks the fill-reducing permutation; see
+// SparseOrdering.
+func NewSparseCholesky(a *core.SparseCOO, ordering SparseOrdering) (*SparseCholesky, error) {
+	const name = "NewSparseCholesky"
+	if a.Values.Dtype() == core.Complex {
+		return nil, base.Errf("%s: complex sparse matrices are not supported", name)
+	}
+	if len(a.Shape) != 2 || a.Shape[0] != a.Shape[1] {
+		return nil, base.Errf("%s: needs a square 2-D sparse matrix, got shape %v", name, a.Shape)
+	}
+	c, err := CSCFromCOO(a)
+	if err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	for j := range c.Cols {
+		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
+			i := c.RowIdx[p]
+			v := c.Values[p]
+			if math.IsNaN(v) || math.IsInf(v, 0) {
+				return nil, base.Errf("%s: entry [%d,%d] is not finite", name, i, j)
+			}
+			// The lower triangle is the authority: an upper entry
+			// must mirror a stored lower entry bit for bit.
+			if i < j {
+				q, ok := findEntry(c, j, i)
+				if !ok || c.Values[q] != v {
+					return nil, base.Errf("%s: entry [%d,%d] = %g has no matching lower counterpart [%d,%d]",
+						name, i, j, v, j, i)
+				}
+			}
+		}
+	}
+	var perm []int
+	switch ordering {
+	case SparseOrderingNatural:
+		perm = make([]int, c.Cols)
+		for i := range perm {
+			perm[i] = i
+		}
+	case SparseOrderingReverseCuthillMcKee:
+		perm, err = reverseCuthillMcKee(c)
+		if err != nil {
+			return nil, base.Errf("%s: %w", name, err)
+		}
+	case SparseOrderingMinimumDegree:
+		perm, err = minimumDegree(c)
+		if err != nil {
+			return nil, base.Errf("%s: %w", name, err)
+		}
+	default:
+		return nil, base.Errf("%s: unknown ordering %d", name, int(ordering))
+	}
+	f := &SparseCholesky{perm: perm, n: c.Cols}
+	f.inversePerm = make([]int, c.Cols)
+	for k, i := range perm {
+		f.inversePerm[i] = k
+	}
+	lower := permutedLower(c, perm)
+	rows := lowerRows(lower)
+	if err := f.factor(lower, rows); err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	return f, nil
+}
+
+// lowerRows returns the row view of a lower triangular matrix: row k
+// holds its columns in ascending order, diagonal included. The full
+// transpose would carry the upper entries into the rows, which the
+// elimination must not see.
+func lowerRows(c *SparseCSC) *SparseCSR {
+	rows := &SparseCSR{Rows: c.Cols, Cols: c.Cols, RowStart: make([]int, c.Cols+1)}
+	for j := range c.Cols {
+		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
+			if c.RowIdx[p] >= j {
+				rows.RowStart[c.RowIdx[p]+1]++
+			}
+		}
+	}
+	for i := range c.Cols {
+		rows.RowStart[i+1] += rows.RowStart[i]
+	}
+	rows.ColIdx = make([]int, rows.RowStart[c.Cols])
+	rows.Values = make([]float64, rows.RowStart[c.Cols])
+	next := make([]int, c.Cols)
+	copy(next, rows.RowStart[:c.Cols])
+	for j := range c.Cols {
+		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
+			i := c.RowIdx[p]
+			if i < j {
+				continue
+			}
+			q := next[i]
+			rows.ColIdx[q] = j
+			rows.Values[q] = c.Values[p]
+			next[i] = q + 1
+		}
+	}
+	return rows
+}
+
+// findEntry looks a stored entry (i, j) up by binary search; the
+// canonical form keeps every column's row indices sorted.
+func findEntry(c *SparseCSC, i, j int) (int, bool) {
+	at, found := slices.BinarySearch(c.RowIdx[c.ColStart[j]:c.ColStart[j+1]], i)
+	if !found {
+		return 0, false
+	}
+	return c.ColStart[j] + at, true
+}
+
+// permutedLower builds the lower triangle of P·A·Pᵀ in CSC form: new
+// position k carries the original row perm[k]. Permuting does not
+// keep the triangle entry for entry, so every stored lower entry (i,
+// j) is placed at the unordered pair of new positions: row max of
+// inversePerm[i] and inversePerm[j], column min. The symmetry check
+// in NewSparseCholesky has already pinned each upper entry to its
+// lower counterpart, so walking the lower triangle alone visits
+// every pair once.
+func permutedLower(c *SparseCSC, perm []int) *SparseCSC {
+	inversePerm := make([]int, c.Cols)
+	for k, i := range perm {
+		inversePerm[i] = k
+	}
+	counts := make([]int, c.Cols+1)
+	for j := range c.Cols {
+		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
+			i := c.RowIdx[p]
+			if i < j {
+				continue
+			}
+			// The entry lands below the new diagonal: row max of the
+			// two new positions, column min.
+			b := min(inversePerm[j], inversePerm[i])
+			counts[b+1]++
+		}
+	}
+	for i := range c.Cols {
+		counts[i+1] += counts[i]
+	}
+	out := &SparseCSC{Rows: c.Cols, Cols: c.Cols, ColStart: counts}
+	out.RowIdx = make([]int, counts[c.Cols])
+	out.Values = make([]float64, counts[c.Cols])
+	next := make([]int, c.Cols)
+	copy(next, counts[:c.Cols])
+	for j := range c.Cols {
+		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
+			i := c.RowIdx[p]
+			if i < j {
+				continue
+			}
+			ki, kj := inversePerm[i], inversePerm[j]
+			b := ki
+			r := kj
+			if kj < ki {
+				b = kj
+				r = ki
+			}
+			q := next[b]
+			out.RowIdx[q] = r
+			out.Values[q] = c.Values[p]
+			next[b] = q + 1
+		}
+	}
+	// Each new column receives entries from several old columns, so
+	// the rows must be sorted into the canonical ascending order,
+	// values travelling with their rows. One scratch buffer serves
+	// every column: it grows to the longest unsorted column once
+	// instead of allocating per column. The row indices within a
+	// column are unique, so the sorted placement is unique and the
+	// scratch cannot move a value.
+	pairsScratch := make([]rowValuePair, 0, c.Cols)
+	for b := range c.Cols {
+		segment := out.RowIdx[counts[b]:counts[b+1]]
+		if slices.IsSorted(segment) {
+			continue
+		}
+		pairs := pairsScratch[:0]
+		for p := counts[b]; p < counts[b+1]; p++ {
+			pairs = append(pairs, rowValuePair{out.RowIdx[p], out.Values[p]})
+		}
+		slices.SortFunc(pairs, func(x, y rowValuePair) int { return cmp.Compare(x.row, y.row) })
+		for p, e := range pairs {
+			out.RowIdx[counts[b]+p] = e.row
+			out.Values[counts[b]+p] = e.val
+		}
+		pairsScratch = pairs
+	}
+	return out
+}
+
+// eliminationTree returns the parent of every column in the
+// elimination tree of the lower triangular pattern, Liu's
+// construction with path compression: for every stored entry (k, j),
+// j < k, the root of j's tree is hung under k, so parent[i] > i and
+// the chain from any j whose row reaches k ends exactly at k.
+// Unparented columns carry −1.
+func eliminationTree(n int, rows *SparseCSR) []int {
+	parent := make([]int, n)
+	ancestor := make([]int, n)
+	for i := range ancestor {
+		ancestor[i] = -1
+	}
+	for k := range n {
+		parent[k] = -1
+		for p := rows.RowStart[k]; p < rows.RowStart[k+1]; p++ {
+			j := rows.ColIdx[p]
+			if j >= k {
+				continue
+			}
+			r := j
+			for r != -1 && r != k {
+				next := ancestor[r]
+				ancestor[r] = k
+				if next == -1 {
+					parent[r] = k
+					break
+				}
+				r = next
+			}
+		}
+	}
+	return parent
+}
+
+// rowReach collects the update columns of row k: every stored entry
+// (k, j) walked up its parent chain until the walk reaches k. The
+// chain can also run through positions whose working entry stays
+// zero, which subtracts nothing later, so the superset costs
+// arithmetic, never correctness.
+//
+// The result is sorted ascending, and the walk keeps it that way as
+// it grows: a chain climbs, so its next vertex usually lands past the
+// tail and appends, and one that lands below inserts into its sorted
+// place. The set's order is part of the elimination's arithmetic: the
+// columns below are processed in it, and two of them can push into
+// one working entry.
+func rowReach(rows *SparseCSR, parent []int, k int, mark []bool, set []int) []int {
+	set = rowReachUnsorted(rows, parent, k, mark, set)
+	if !slices.IsSorted(set) {
+		slices.Sort(set)
+	}
+	return set
+}
+
+// rowReachUnsorted collects the same set without ordering it. The
+// counting pass reads the set's length alone, where the walk order
+// costs nothing, so half the reach walks pay no sort.
+func rowReachUnsorted(rows *SparseCSR, parent []int, k int, mark []bool, set []int) []int {
+	set = set[:0]
+	for p := rows.RowStart[k]; p < rows.RowStart[k+1]; p++ {
+		j := rows.ColIdx[p]
+		for j != k && j != -1 && !mark[j] {
+			mark[j] = true
+			set = append(set, j)
+			j = parent[j]
+		}
+	}
+	for _, j := range set {
+		mark[j] = false
+	}
+	return set
+}
+
+// cholColumnCounts returns the exact number of entries every column of
+// the factor receives, by running the same reach walk the numeric
+// elimination runs and counting the columns it visits. The numeric pass
+// appends an entry to column j for every j in the reach of a row k
+// except where a working entry has cancelled to zero, which only makes
+// the real count smaller, and both passes call the same pure walk on
+// the same elimination tree. The counts are therefore an exact upper
+// bound: the arena below needs no growth and leaves no gap between
+// columns. The walk runs unsorted here: the counts read the set's
+// size, which no order changes.
+//
+// The result is returned in prefix-sum form: column j's entries start
+// at counts[j], and counts[n] is the factor's off-diagonal entry count.
+func cholColumnCounts(n int, rows *SparseCSR, parent []int, mark []bool, set []int) []int {
+	counts := make([]int, n+1)
+	for k := range n {
+		set = rowReachUnsorted(rows, parent, k, mark, set)
+		for _, j := range set {
+			counts[j+1]++
+		}
+	}
+	for j := range n {
+		counts[j+1] += counts[j]
+	}
+	return counts
+}
+
+// factor runs the row-oriented elimination. Row k of L is its
+// working row divided through the square-rooted diagonal, the update
+// columns come from row k's stored entries walked through the
+// elimination tree, and each settled entry L[k,j] pushes its
+// contribution down column j of the factor, the update the left-looking
+// recurrence demands: L[k,i] loses L[k,j]·L[i,j] for every earlier
+// column j, so the columns are walked ascending and every working
+// entry has settled by the time its turn comes. The elimination
+// tree's guarantee makes the walk safe: every stored entry (k, j)
+// sits on j's parent chain, and a chain position whose working entry
+// stays zero simply subtracts nothing.
+//
+// The factor is written straight into one flat arena, a contiguous
+// block per column with a cursor, so the whole numeric factor costs two
+// allocations and the scatter into the working row reads one base array
+// instead of a slice header per column. A column's entries land in
+// ascending k, the order the solver reads.
+func (f *SparseCholesky) factor(lower *SparseCSC, rows *SparseCSR) error {
+	const name = "NewSparseCholesky"
+	n := lower.Cols
+	parent := eliminationTree(n, rows)
+	f.diag = make([]float64, n)
+	x := make([]float64, n)
+	set := make([]int, 0, n)
+	mark := make([]bool, n)
+	f.colStart = cholColumnCounts(n, rows, parent, mark, set)
+	arenaRows := make([]int, f.colStart[n])
+	arenaVals := make([]float64, f.colStart[n])
+	// next[j] is column j's cursor: where its next entry lands. It
+	// starts at the column's block and ends one past its last written
+	// entry, so next[j]−colStart[j] is the column's actual length.
+	next := make([]int, n)
+	copy(next, f.colStart[:n])
+	for k := range n {
+		// x holds row k of the working matrix: the stored entries
+		// minus the updates below.
+		rs, re := rows.RowStart[k], rows.RowStart[k+1]
+		ci := rows.ColIdx[rs:re:re]
+		cv := rows.Values[rs:re:re]
+		for i, c := range ci {
+			x[c] = cv[i]
+		}
+		set = rowReach(rows, parent, k, mark, set)
+		for _, j := range set {
+			dj := f.diag[j]
+			lkj := x[j] / dj
+			if lkj != 0 {
+				at := next[j]
+				arenaRows[at] = k
+				arenaVals[at] = lkj
+				next[j] = at + 1
+			}
+			// Push the entry down column j: the diagonal term zeroes
+			// x[j], the rows below carry the fill, and the k term is
+			// the diagonal sum of this very row. The column is walked
+			// through slices cut to its written range, so the loop
+			// carries no bound checks and no reload of the cursor
+			// next[j] the writes above could be thought to alias.
+			x[j] -= lkj * dj
+			cs, ce := f.colStart[j], next[j]
+			rw := arenaRows[cs:ce:ce]
+			rv := arenaVals[cs:ce:ce]
+			for i, r := range rw {
+				x[r] -= lkj * rv[i]
+			}
+		}
+		d := x[k]
+		if math.IsInf(d, 0) {
+			return base.Errf("%s: the pivot at [%d,%d] overflowed to %g; the elimination has left the float64 range",
+				name, f.perm[k], f.perm[k], d)
+		}
+		if math.IsNaN(d) || d <= 0 {
+			return base.Errf("%s: the pivot at [%d,%d] is %g; the matrix is not positive definite",
+				name, f.perm[k], f.perm[k], d)
+		}
+		f.diag[k] = math.Sqrt(d)
+		// Restore every x position the row touched: the update
+		// columns, their factor entries, and the diagonal.
+		for _, j := range set {
+			x[j] = 0
+			cs, ce := f.colStart[j], next[j]
+			for _, r := range arenaRows[cs:ce:ce] {
+				x[r] = 0
+			}
+		}
+		x[k] = 0
+	}
+	// Pack the columns: a column whose working entry cancelled along
+	// the way sits shorter than its block, so only the written range
+	// moves, and a column's entries keep their ascending-k order.
+	total := 0
+	for j := range n {
+		total += next[j] - f.colStart[j]
+	}
+	colStart := make([]int, n+1)
+	f.rowIdx = make([]int, total)
+	f.values = make([]float64, total)
+	for j := range n {
+		b, w := f.colStart[j], next[j]
+		colStart[j+1] = colStart[j] + (w - b)
+		copy(f.rowIdx[colStart[j]:], arenaRows[b:w])
+		copy(f.values[colStart[j]:], arenaVals[b:w])
+	}
+	f.colStart = colStart
+	return nil
+}
+
+// Solve computes x = A⁻¹·b for a dense vector b: permute, forward
+// substitution with L, backward substitution with Lᵀ, undo the
+// permutation.
+func (f *SparseCholesky) Solve(b *core.Array) (*core.Array, error) {
+	const name = "Solve"
+	if b.NDim() != 1 {
+		return nil, base.Errf("%s: the right hand side must be rank 1", name)
+	}
+	if b.Dtype() == core.Complex {
+		return nil, base.Errf("%s: complex right hand sides are not supported", name)
+	}
+	if b.Len() != f.n {
+		return nil, base.Errf("%s: right hand side length %d does not match %d rows", name, b.Len(), f.n)
+	}
+	x := core.New(core.Float, []int{f.n}...)
+	xf := x.RawFloats()
+	for k := range f.n {
+		xf[k] = b.FloatAt(f.perm[k])
+	}
+	for j := range f.n {
+		xj := xf[j] / f.diag[j]
+		xf[j] = xj
+		for p := f.colStart[j]; p < f.colStart[j+1]; p++ {
+			xf[f.rowIdx[p]] -= f.values[p] * xj
+		}
+	}
+	for j := f.n - 1; j >= 0; j-- {
+		sum := xf[j]
+		for p := f.colStart[j]; p < f.colStart[j+1]; p++ {
+			sum -= f.values[p] * xf[f.rowIdx[p]]
+		}
+		xf[j] = sum / f.diag[j]
+	}
+	// z solves P·A·Pᵀ·z = P·b, so z[k] is the answer at original
+	// position perm[k]: the scatter walks the same side as the
+	// gather.
+	out := core.New(core.Float, []int{f.n}...)
+	of := out.RawFloats()
+	for k := range f.n {
+		of[f.perm[k]] = xf[k]
+	}
+	return out, nil
+}
+
+// Permutation returns the elimination order: position k holds the
+// original index factored k-th, with P·A·Pᵀ = L·Lᵀ. The slice is a
+// copy, so the caller cannot move the factor's state.
+func (f *SparseCholesky) Permutation() []int {
+	return slices.Clone(f.perm)
+}
+
+// NNZ returns the count of stored non-zeros in the factor, diagonal
+// included.
+func (f *SparseCholesky) NNZ() int { return len(f.values) + f.n }