feat: initial release

Assisted-by: GLM 5.3 Flash

linalg/decomp.go

@@ -0,0 +1,596 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package linalg
+
+import (
+	"math"
+	"sync"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+	"sourcedock.dev/petrbalvin/tensor/internal/engine"
+)
+
+// Matrix decompositions. All routines operate on float64 in
+// and out; complex matrices are rejected. The QR routine uses
+// Householder reflections, numerically stable and standard. Cholesky
+// uses the Banachiewicz algorithm. LeastSquares builds on QR.
+
+// QR returns the QR decomposition a = Q * R of an m×n matrix a, with
+// m ≥ n. Q is m×m orthogonal and R is m×n upper triangular.
+func QR(a *core.Array) (q, r *core.Array, err error) {
+	if a.Dtype() == core.Complex {
+		return nil, nil, base.Errf("QR: complex matrices are not supported")
+	}
+	if a.NDim() != 2 {
+		return nil, nil, base.Errf("QR: needs a 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
+	}
+	m, n := a.Shape()[0], a.Shape()[1]
+	if m < n {
+		return nil, nil, base.Errf("QR: needs m ≥ n, got shape %s", base.ShapeText(a.Shape()))
+	}
+	aMat := denseFloats(a, m, n)
+	qMat, rMat, err := qrRaw(aMat, m, n)
+	if err != nil {
+		return nil, nil, err
+	}
+	return floatsToArray(qMat, []int{m, m}), floatsToArray(rMat, []int{m, n}), nil
+}
+
+// Cholesky returns the lower-triangular Cholesky factor L such that
+// a = L * Lᵀ for a symmetric positive definite matrix a.
+func Cholesky(a *core.Array) (*core.Array, error) {
+	if a.Dtype() == core.Complex {
+		return nil, base.Errf("Cholesky: complex matrices are not supported")
+	}
+	if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
+		return nil, base.Errf("Cholesky: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
+	}
+	n := a.Shape()[0]
+	lMat, err := denseCholFactor(denseFloats(a, n, n), n)
+	if err != nil {
+		return nil, err
+	}
+	return floatsToArray(lMat, []int{n, n}), nil
+}
+
+// denseCholBlock is the column-block width of the Cholesky sweep. A wider
+// block shortens the serial diagonal factorisation and lengthens the
+// parallel panel update's k run; 64 keeps the serial share near a
+// twentieth of the arithmetic at the sizes this package serves while
+// both panel passes stay inside a cache line's neighbourhood.
+const denseCholBlock = 64
+
+// denseCholMinWork is the element-touch floor one Cholesky panel worker must
+// receive before the dispatch is worth its spawn cost, counted in
+// update terms (rows times columns times the k run). Below it the panel
+// runs on the calling goroutine.
+const denseCholMinWork = 8192
+
+// denseCholFactor factors a flat n×n row-major matrix into its lower
+// triangular Cholesky factor, consuming the input.
+//
+// The sweep runs in column blocks of denseCholBlock: each block is first
+// updated from the finished columns, then factored serially, then
+// divided into the trapezoid below it. For an element (i, j) the three
+// stages subtract the columns k < j in the one ascending order the
+// unblocked dot always used, since the panel update covers exactly the
+// k < j0 prefix and the factorisation or the divide covers the rest,
+// so the running value the last stage divides is bit for bit the value
+// the unblocked loop held. Only the order in which independent elements
+// are visited changes, which is why the panel update and the trapezoid
+// divide, both disjoint per row, run over a crew.
+func denseCholFactor(mat []float64, n int) ([]float64, error) {
+	lMat := make([]float64, n*n)
+	for j0 := 0; j0 < n; j0 += denseCholBlock {
+		j1 := min(j0+denseCholBlock, n)
+		// Panel update: every element of rows j0..n−1 in columns
+		// j0..j1−1 loses the contribution of the finished columns k < j0.
+		// A row's elements depend on those columns alone, so rows split
+		// over the crew.
+		if j0 > 0 {
+			cols := j1 - j0
+			panel := func(start, end int) {
+				for i := start; i < end; i++ {
+					for j := j0; j < j1 && j <= i; j++ {
+						s := mat[i*n+j]
+						for k := range j0 {
+							s -= lMat[i*n+k] * lMat[j*n+k]
+						}
+						mat[i*n+j] = s
+					}
+				}
+			}
+			if rows := n - j0; rows*cols*j0 >= denseCholMinWork {
+				engine.ParallelMin(rows, denseCholRows(denseCholMinWork, cols*j0), func(start, end int) {
+					panel(j0+start, j0+end)
+				})
+			} else {
+				panel(j0, n)
+			}
+		}
+		// Diagonal block: columns in order, each finished before the
+		// next reads it. This is the only serial arithmetic left, and it
+		// shrinks with the block width: the block's own triangle is
+		// denseCholBlock³/6 of the matrix's n³/6.
+		for j := j0; j < j1; j++ {
+			for i := j; i < j1; i++ {
+				s := mat[i*n+j]
+				for k := j0; k < j; k++ {
+					s -= lMat[i*n+k] * lMat[j*n+k]
+				}
+				if i == j {
+					if s <= 0 {
+						return nil, base.Errf("Cholesky: matrix is not positive definite (pivot %d = %v)", i, s)
+					}
+					lMat[i*n+j] = math.Sqrt(s)
+				} else {
+					if lMat[j*n+j] == 0 {
+						return nil, base.Errf("Cholesky: zero diagonal at %d", j)
+					}
+					lMat[i*n+j] = s / lMat[j*n+j]
+				}
+			}
+		}
+		// Trapezoid below the block: row i of the panel needs row j of the
+		// diagonal block, which the sweep above has finished, and then its
+		// own earlier columns of the same block; both are in hand, so the
+		// rows divide independently. The blocks are disjoint per row, so
+		// one crew covers the divide.
+		if j1 < n {
+			divide := func(start, end int) {
+				for j := j0; j < j1; j++ {
+					d := lMat[j*n+j]
+					for i := start; i < end; i++ {
+						s := mat[i*n+j]
+						for k := j0; k < j; k++ {
+							s -= lMat[i*n+k] * lMat[j*n+k]
+						}
+						lMat[i*n+j] = s / d
+					}
+				}
+			}
+			if rows := n - j1; rows*(j1-j0)*(j1-j0)/2 >= denseCholMinWork {
+				engine.ParallelMin(rows, denseCholRows(denseCholMinWork, (j1-j0)*(j1-j0)/2), func(start, end int) {
+					divide(j1+start, j1+end)
+				})
+			} else {
+				divide(j1, n)
+			}
+		}
+	}
+	return lMat, nil
+}
+
+// denseCholRows returns the rows one worker needs for the given per-row
+// cost to reach the dispatch floor, never below one.
+func denseCholRows(floor, cost int) int {
+	if cost <= 0 {
+		return 1
+	}
+	return max((floor+cost-1)/cost, 1)
+}
+
+// LeastSquares solves Ax = b in the least-squares sense. a is m×n
+// with m ≥ n. b is m or m×k. Implemented via QR.
+func LeastSquares(a, b *core.Array) (*core.Array, error) {
+	if a.Dtype() == core.Complex {
+		return nil, base.Errf("LeastSquares: complex matrices are not supported")
+	}
+	// The QR route reads b through FloatAt as well: without the gate a
+	// complex right-hand side reaches the empty int payload and panics.
+	if b.Dtype() == core.Complex {
+		return nil, base.Errf("LeastSquares: complex right-hand sides are not supported")
+	}
+	if a.NDim() != 2 {
+		return nil, base.Errf("LeastSquares: 'a' must be 2-D, got shape %s", base.ShapeText(a.Shape()))
+	}
+	m, n := a.Shape()[0], a.Shape()[1]
+	if m < n {
+		return nil, base.Errf("LeastSquares: needs m ≥ n, got shape %s", base.ShapeText(a.Shape()))
+	}
+	if b.NDim() != 1 && b.NDim() != 2 {
+		return nil, base.Errf("LeastSquares: 'b' must be 1-D or 2-D, got shape %s", base.ShapeText(b.Shape()))
+	}
+	if b.Shape()[0] != m {
+		return nil, base.Errf("LeastSquares: b rows (%d) must match a rows (%d)", b.Shape()[0], m)
+	}
+	k := 1
+	if b.NDim() == 2 {
+		k = b.Shape()[1]
+	}
+	aMat := denseFloats(a, m, n)
+	bMat := make([]float64, m*k)
+	if b.Dtype() == core.Float && !b.Strided() && len(b.RawFloats()) == m*k {
+		copy(bMat, b.RawFloats())
+	} else if b.NDim() == 1 {
+		for i := range m {
+			bMat[i] = b.FloatAt(i)
+		}
+	} else {
+		for i := range m {
+			for j := range k {
+				bMat[i*k+j] = b.FloatAt(i*k + j)
+			}
+		}
+	}
+	// The reflectors go straight onto b; no Q is formed, because the
+	// only thing the solve wants from it is Qᵀ·b.
+	qtB, rMat := qrSolveRows(aMat, m, n, k, bMat)
+	x := make([]float64, n*k)
+	// Rank guard: with no column pivoting an exactly or nearly
+	// dependent column leaves a rounding-level R diagonal instead of an
+	// exact zero, which back-substitution would amplify into a huge
+	// silent answer. Pivots below the Householder-QR analogue of the
+	// SVD route's threshold, n·eps·max|R|, refuse the system.
+	rMax := 0.0
+	for i := range n {
+		if v := math.Abs(rMat[i*n+i]); v > rMax {
+			rMax = v
+		}
+	}
+	rankTol := float64(n) * base.EpsF * rMax
+	for col := 0; col < k; col++ {
+		for i := n - 1; i >= 0; i-- {
+			s := qtB[i*k+col]
+			for j := i + 1; j < n; j++ {
+				s -= rMat[i*n+j] * x[j*k+col]
+			}
+			if math.Abs(rMat[i*n+i]) <= rankTol {
+				return nil, base.Errf("LeastSquares: rank-deficient system (R pivot %d = %g is at the rounding floor %g)", i, rMat[i*n+i], rankTol)
+			}
+			x[i*k+col] = s / rMat[i*n+i]
+		}
+	}
+	if b.NDim() == 1 {
+		return floatsToArray(x, []int{n}), nil
+	}
+	return floatsToArray(x, []int{n, k}), nil
+}
+
+// householderReflector builds the reflector that zeroes the entries
+// below the diagonal of column k of the row-major m×n matrix r and
+// writes its scaled vector into v[:m−k]. It reports the vector length,
+// beta = 2/(vᵀv) and whether a reflector was needed at all: a zero
+// column, a zero norm or a zero vector leaves the column as the sweep
+// wants it and the step is skipped.
+//
+// The scaling is the one the sweep has always used: entries near 1e154
+// would overflow the raw squares to +Inf, which zeroes beta and forms
+// 0*Inf = NaN in the update, so the column is divided by its largest
+// magnitude first and beta absorbs the square of that scale, leaving
+// the update beta·(vᵀy)·v at the value the unscaled arithmetic would
+// have had, with every intermediate O(1).
+func householderReflector(r []float64, m, n, k int, v []float64) (ln int, beta float64, ok bool) {
+	scale := 0.0
+	for i := k; i < m; i++ {
+		if a := math.Abs(r[i*n+k]); a > scale {
+			scale = a
+		}
+	}
+	if scale == 0 {
+		return 0, 0, false
+	}
+	sum := 0.0
+	for i := k; i < m; i++ {
+		t := r[i*n+k] / scale
+		sum += t * t
+	}
+	norm := math.Sqrt(sum)
+	if norm == 0 {
+		return 0, 0, false
+	}
+	xk := r[k*n+k] / scale
+	s := -sign(xk)
+	if s == 0 {
+		s = -1
+	}
+	u0 := xk - s*norm
+	vv := v[:m-k]
+	vv[0] = u0
+	for i := 1; i < m-k; i++ {
+		vv[i] = r[(k+i)*n+k] / scale
+	}
+	vtv := u0 * u0
+	for i := 1; i < m-k; i++ {
+		vtv += vv[i] * vv[i]
+	}
+	if vtv == 0 {
+		return 0, 0, false
+	}
+	return m - k, 2.0 / vtv, true
+}
+
+// qrSolveRows runs the Householder sweep on a row-major m×n system and
+// applies every reflector to the k right-hand sides as it goes, so the
+// projection Qᵀ·b comes out without Q ever being formed: the sweep costs
+// O(m·n²) for R and O(m·n·k) for the right-hand sides, where forming Q
+// costs O(m²·n) and the dense product after it another O(m²·k). For the
+// tall systems this serves, m ≫ n, that product is nearly the whole
+// solve, and applying the reflectors is also the more accurate route:
+// the projection stays orthogonal by construction instead of carrying
+// the rounding of an explicitly accumulated Q. It returns the leading n
+// rows of the transformed right-hand sides, concatenated by column, and
+// R.
+func qrSolveRows(aMat []float64, m, n, k int, rhs []float64) (qtB, rMat []float64) {
+	rMat = make([]float64, m*n)
+	copy(rMat, aMat)
+	y := make([]float64, m*k)
+	copy(y, rhs)
+	v := make([]float64, m)
+	for kk := range n {
+		ln, beta, ok := householderReflector(rMat, m, n, kk, v)
+		if !ok {
+			continue
+		}
+		vv := v[:ln]
+		// The right-hand sides first: each column's dot runs over i
+		// ascending and its update repeats the same operand order, and
+		// the columns are independent outputs, so a split over them
+		// cannot move a value.
+		if k*ln >= qrRHSMinWork {
+			engine.ParallelMin(k, 1, func(start, end int) {
+				for j := start; j < end; j++ {
+					applyReflectorToColumn(y, j, k, kk, ln, beta, vv)
+				}
+			})
+		} else {
+			for j := range k {
+				applyReflectorToColumn(y, j, k, kk, ln, beta, vv)
+			}
+		}
+		// Then the trailing columns of R, blocked and dispatched exactly
+		// as the explicit-Q sweep does it, with the Q columns it also
+		// carries left out.
+		rBlocks := (n - kk + qrColumnBlock - 1) / qrColumnBlock
+		worker := func(gi, nw int) {
+			var d [qrColumnBlock]float64
+			for b := gi * rBlocks / nw; b < (gi+1)*rBlocks/nw; b++ {
+				j0 := kk + b*qrColumnBlock
+				j1 := min(j0+qrColumnBlock, n)
+				wd := j1 - j0
+				_ = d[wd-1] // bounds-check proof: the block never exceeds qrColumnBlock
+				clear(d[:wd])
+				for jj := range wd {
+					jcol := j0 + jj
+					for i := range ln {
+						d[jj] += vv[i] * rMat[(kk+i)*n+jcol]
+					}
+				}
+				for j := range wd {
+					d[j] = beta * d[j]
+				}
+				for i := range ln {
+					vi := vv[i]
+					base := (kk+i)*n + j0
+					for j := range wd {
+						rMat[base+j] -= vi * d[j]
+					}
+				}
+			}
+		}
+		weight := (n - kk) + k
+		w := min(weight*ln/qrWorkQuantum+1, engine.WorkersFor(rBlocks))
+		if w < 2 {
+			worker(0, 1)
+			continue
+		}
+		var wg sync.WaitGroup
+		for gi := range w {
+			wg.Go(func() {
+				worker(gi, w)
+			})
+		}
+		wg.Wait()
+	}
+	qtB = make([]float64, n*k)
+	for j := range k {
+		for i := range n {
+			qtB[i*k+j] = y[i*k+j]
+		}
+	}
+	return qtB, rMat
+}
+
+// applyReflectorToColumn applies one reflector to column j of an m×k
+// right-hand-side block: the dot runs over i ascending and the update
+// repeats the same operand order.
+func applyReflectorToColumn(y []float64, j, k, kk, ln int, beta float64, vv []float64) {
+	dot := 0.0
+	base := kk * k
+	for i := range ln {
+		dot += vv[i] * y[base+i*k+j]
+	}
+	w := beta * dot
+	for i := range ln {
+		y[base+i*k+j] -= vv[i] * w
+	}
+}
+
+// qrRHSMinWork is the element work one worker must carry before the
+// right-hand-side columns of a reflector are split across goroutines.
+const qrRHSMinWork = 1 << 15
+
+// qrWorkQuantum is the element-touch budget one worker receives in a
+// qrRaw reflector dispatch, counted as the reflector's reach (columns
+// of R plus columns of Q) times the reflector length. A reflector whose
+// application does not fill one quantum runs on the calling goroutine:
+// the dispatch would cost more than the work. A fixed 32-way split
+// measured 15 000 allocations and several milliseconds of scheduler
+// time per 256² QR, because 32 goroutines spawn and synchronise per
+// reflector however small the reflector's reach has grown; sizing the
+// crew by the work keeps each dispatch's spawn cost proportional to the
+// work it spreads.
+const qrWorkQuantum = 8192
+
+// qrColumnBlock is the number of R columns one blocked update pass
+// covers. R is row-major, so a single column's update walks a stride-n
+// pattern that touches one cache line per element; a block of columns
+// walks the same rows linearly and reuses each line across the block.
+// Eight columns span one 64-byte line of float64.
+const qrColumnBlock = 8
+
+// qrRaw performs Householder QR on a flat m×n matrix and returns the
+// flat Q (m, m) and R (m, n).
+//
+// The reflectors apply one after another: their ORDER defines Q and R
+// and is never changed. Within a single reflector every column of R
+// and every column of Q is updated exactly once, reading only the
+// frozen reflector v and beta, so the per-column work (dot along i
+// ascending, then the same-operand rank-1 subtraction) is dispatched
+// over the columns. Identical inputs, identical per-element sequence:
+// the result is bit-identical to the serial sweep, whichever crew size
+// or column blocking the dispatcher picks.
+func qrRaw(aMat []float64, m, n int) (qMat, rMat []float64, err error) {
+	rMat = make([]float64, m*n)
+	copy(rMat, aMat)
+	qMat = eye(m)
+	// One reflector scratch reused across the sweep; each step uses its
+	// first m−k entries.
+	v := make([]float64, m)
+	// Per-reflector state the crew reads: the sweep writes it on the
+	// calling goroutine and the crew only reads it, so one body closure
+	// and its captures live for the whole sweep instead of a fresh pair
+	// of closures escaping to the heap per reflector.
+	var (
+		k    int       // reflector's leading row and column
+		ln   int       // reflector length, m−k
+		beta float64   // 2/(vᵀv) of the current reflector
+		vv   []float64 // reflector prefix view, re-sliced per step
+	)
+	// worker runs the reflector's R columns and Q columns that fall to
+	// crew member gi of nw: a contiguous slab of the R blocks and a
+	// contiguous slab of the Q columns. Every member owns its own block
+	// scratch d, and slabs are disjoint, so no two members ever touch
+	// the same line; contiguous slabs keep each member's working set at
+	// its own rows of qMat and its own column stripe of rMat instead of
+	// walking the whole window. Each column's dot runs over i ascending
+	// and the subtraction repeats the same operand order, so the slab
+	// boundaries never move a bit.
+	worker := func(gi, nw int) {
+		rBlocks := (n - k + qrColumnBlock - 1) / qrColumnBlock
+		// R columns in blocks of qrColumnBlock. R is row-major, so a
+		// lone column's update walks a stride-n pattern that touches one
+		// cache line per element; a block keeps the hot window at one
+		// line per row. The per-column dot below keeps the exact shape
+		// the per-column sweep always had (same accumulator shape, same
+		// operand order), because reshaping the dot loop changes the
+		// compiler's fusing of the multiply into the add and that moves
+		// bits. Each column's w is beta·dot as before, and the blocked
+		// subtraction walks the same i order over the block linearly.
+		var d [qrColumnBlock]float64
+		for b := gi * rBlocks / nw; b < (gi+1)*rBlocks/nw; b++ {
+			j0 := k + b*qrColumnBlock
+			j1 := min(j0+qrColumnBlock, n)
+			wd := j1 - j0
+			_ = d[wd-1] // bounds-check proof: the block never exceeds qrColumnBlock
+			clear(d[:wd])
+			for jj := range wd {
+				jcol := j0 + jj
+				for i := range ln {
+					d[jj] += vv[i] * rMat[(k+i)*n+jcol]
+				}
+			}
+			for j := range wd {
+				d[j] = beta * d[j]
+			}
+			for i := range ln {
+				vi := vv[i]
+				base := (k+i)*n + j0
+				for j := range wd {
+					rMat[base+j] -= vi * d[j]
+				}
+			}
+		}
+		// Q columns, one at a time: per column the access is already
+		// linear in i (row-major, row col, columns k..m−1), and a slab
+		// of consecutive columns is a slab of consecutive rows.
+		for col := gi * m / nw; col < (gi+1)*m/nw; col++ {
+			dot := 0.0
+			base := col*m + k
+			for i := range ln {
+				dot += vv[i] * qMat[base+i]
+			}
+			w := beta * dot
+			for i := range ln {
+				qMat[base+i] -= vv[i] * w
+			}
+		}
+	}
+	for k = range n {
+		// The reflector's construction and its scaling live in
+		// householderReflector, shared with the explicit-Q-free solve
+		// below.
+		var ok bool
+		ln, beta, ok = householderReflector(rMat, m, n, k, v)
+		if !ok {
+			continue
+		}
+		vv = v[:ln]
+		// Columns k..n-1 of R (in blocks, one item per block) and all m
+		// columns of Q: disjoint writes, read-only v and beta, so they
+		// can run as one parallel crew. The work weight counts a block
+		// as the qrColumnBlock columns it covers, which keeps the crew
+		// size honest about R's share.
+		items := (n-k+qrColumnBlock-1)/qrColumnBlock + m
+		weight := (n - k) + m
+		w := min(weight*ln/qrWorkQuantum+1, engine.WorkersFor(items))
+		if w < 2 {
+			worker(0, 1)
+			continue
+		}
+		var wg sync.WaitGroup
+		for gi := range w {
+			wg.Go(func() {
+				worker(gi, w)
+			})
+		}
+		wg.Wait()
+	}
+	return qMat, rMat, nil
+}
+
+// denseFloats copies a 2-D array's payload into a flat m×n float64
+// matrix, converting from any numeric dtype via floatAt. Contiguous
+// float64 payloads take a direct-copy fast path; floatAt gives the
+// identical values for every dtype and layout either way.
+func denseFloats(a *core.Array, m, n int) []float64 {
+	out := make([]float64, m*n)
+	if a.Dtype() == core.Float && !a.Strided() && len(a.RawFloats()) == m*n {
+		copy(out, a.RawFloats())
+		return out
+	}
+	for i := range m {
+		for j := range n {
+			out[i*n+j] = a.FloatAt(i*n + j)
+		}
+	}
+	return out
+}
+
+// floatsToArray wraps a flat m×n float64 matrix as a core.Array.
+func floatsToArray(data []float64, shape []int) *core.Array {
+	out := core.New(core.Float, shape...)
+	copy(out.RawFloats(), data)
+	return out
+}
+
+func eye(n int) []float64 {
+	out := make([]float64, n*n)
+	for i := range n {
+		out[i*n+i] = 1
+	}
+	return out
+}
+
+func sign(x float64) float64 {
+	switch {
+	case x > 0:
+		return 1
+	case x < 0:
+		return -1
+	default:
+		return 0
+	}
+}