feat: initial release
Assisted-by: GLM 5.3 Flash
This commit is contained in:
+422
@@ -0,0 +1,422 @@
|
||||
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package linalg
|
||||
|
||||
import (
|
||||
"math"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||||
)
|
||||
|
||||
// Rank-revealing QR with column pivoting (Businger and Golub). The
|
||||
// plain Householder QR of `QR` answers a factorisation but hides the
|
||||
// rank: an exactly or nearly dependent column leaves a rounding-level
|
||||
// R diagonal and `LeastSquares` has to refuse the system outright.
|
||||
// Pivoting sweeps the column of largest remaining norm to the front at
|
||||
// every step, which forces the dependence into the trailing block: the
|
||||
// |R| diagonal comes out non-increasing and its decay is the rank
|
||||
// decision, the same tolerance convention `LeastSquares` applies to
|
||||
// its unpivoted diagonal.
|
||||
|
||||
// RRQR returns the rank-revealing QR factorisation with column
|
||||
// pivoting of an m×n matrix a with m ≥ n: A·P = Q·R, with Q m×m
|
||||
// orthogonal, R m×n upper triangular, and perm the column permutation
|
||||
// P, meaning the j-th column of A·P is column perm[j] of A. The rank
|
||||
// is the count of |R[i,i]| above n·eps·max|R|, the tolerance
|
||||
// `LeastSquares` applies; with column pivoting the diagonal decays
|
||||
// monotonically, so the count is where the decay crosses the floor.
|
||||
// Complex matrices, non-2-D inputs and underdetermined shapes are
|
||||
// refused, matching the dense surface.
|
||||
func RRQR(a *core.Array) (q, r *core.Array, perm []int, rank int, err error) {
|
||||
const name = "RRQR"
|
||||
if a.Dtype() == core.Complex {
|
||||
return nil, nil, nil, 0, base.Errf("%s: complex matrices are not supported", name)
|
||||
}
|
||||
if a.NDim() != 2 {
|
||||
return nil, nil, nil, 0, base.Errf("%s: needs a 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
m, n := a.Shape()[0], a.Shape()[1]
|
||||
if m == 0 || n == 0 {
|
||||
return nil, nil, nil, 0, base.Errf("%s: zero-sized matrix, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
if m < n {
|
||||
return nil, nil, nil, 0, base.Errf("%s: needs m ≥ n, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
rMat, qMat, perm := rrqrFactor(denseFloats(a, m, n), m, n, true)
|
||||
return floatsToArray(qMat, []int{m, m}), floatsToArray(rMat, []int{m, n}), perm, rrqrRankOf(rMat, n), nil
|
||||
}
|
||||
|
||||
// RRQRRank returns the numerical rank a pivoted factorisation of a
|
||||
// reports under the package's tolerance convention: the count of
|
||||
// |R[i,i]| above n·eps·max|R|. It is the rank `RRQR` returns, computed
|
||||
// without the orthogonal factor for callers that only want the count.
|
||||
func RRQRRank(a *core.Array) (int, error) {
|
||||
const name = "RRQRRank"
|
||||
if a.Dtype() == core.Complex {
|
||||
return 0, base.Errf("%s: complex matrices are not supported", name)
|
||||
}
|
||||
if a.NDim() != 2 {
|
||||
return 0, base.Errf("%s: needs a 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
m, n := a.Shape()[0], a.Shape()[1]
|
||||
if m == 0 || n == 0 {
|
||||
return 0, base.Errf("%s: zero-sized matrix, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
if m < n {
|
||||
return 0, base.Errf("%s: needs m ≥ n, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
rMat, _, _ := rrqrFactor(denseFloats(a, m, n), m, n, false)
|
||||
return rrqrRankOf(rMat, n), nil
|
||||
}
|
||||
|
||||
// rrqrRankOf counts the diagonal entries of the flat m×n R above
|
||||
// n·eps·max|R|, the rank convention `LeastSquares` states for its
|
||||
// unpivoted diagonal.
|
||||
func rrqrRankOf(rMat []float64, n int) int {
|
||||
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
|
||||
rank := 0
|
||||
for i := range n {
|
||||
if math.Abs(rMat[i*n+i]) > rankTol {
|
||||
rank++
|
||||
}
|
||||
}
|
||||
return rank
|
||||
}
|
||||
|
||||
// rrqrColumnBlock is the number of R columns one blocked reflector pass
|
||||
// covers, the blocking qrRaw applies to the same sweep and for the same
|
||||
// reason: R is row-major, so a lone column's update walks a stride-n
|
||||
// pattern that touches one cache line per element, where a block of
|
||||
// rrqrColumnBlock columns spans one line of float64 and walks each row
|
||||
// once.
|
||||
const rrqrColumnBlock = 8
|
||||
|
||||
// rrqrMinWork is the element-touch floor below which one worker's share
|
||||
// of a reflector application or of a column-norm scan falls back to the
|
||||
// calling goroutine: below it the dispatch costs more than the work.
|
||||
const rrqrMinWork = 2048
|
||||
|
||||
// rrqrFactor runs the pivoted Householder sweep on a flat m×n matrix.
|
||||
// The column of largest remaining norm moves to the front at every
|
||||
// step; norms are recomputed each step rather than downdated, so no
|
||||
// cancellation drift and no squared-overflow window: the cost is one
|
||||
// pass over the working block per step, the same order as the
|
||||
// reflector sweep itself. The Householder reflectors and their
|
||||
// application follow `qrRaw`: scale before squaring, one reflector per
|
||||
// step, strictly in sweep order, so the serial result is the
|
||||
// factorisation the parallel plain QR also answers.
|
||||
func rrqrFactor(aMat []float64, m, n int, wantQ bool) (rMat, qMat []float64, perm []int) {
|
||||
rMat = make([]float64, m*n)
|
||||
copy(rMat, aMat)
|
||||
perm = make([]int, n)
|
||||
for i := range perm {
|
||||
perm[i] = i
|
||||
}
|
||||
if wantQ {
|
||||
qMat = eye(m)
|
||||
}
|
||||
col := make([]float64, m)
|
||||
norms := make([]float64, n)
|
||||
for k := range n {
|
||||
best := rrqrPivot(rMat, m, n, k, perm, norms)
|
||||
if best == 0 {
|
||||
// Every remaining column is exactly zero: the trailing
|
||||
// block of R stays zero and the sweep is done.
|
||||
break
|
||||
}
|
||||
lv := col[:m-k]
|
||||
for i := range lv {
|
||||
lv[i] = rMat[(k+i)*n+k]
|
||||
}
|
||||
hh := householderVectorInto(lv, lv)
|
||||
if hh.beta == 0 {
|
||||
continue
|
||||
}
|
||||
rrqrApply(rMat, m, n, k, hh, qMat, nil)
|
||||
}
|
||||
return rMat, qMat, perm
|
||||
}
|
||||
|
||||
// rrqrPivot selects the column of largest remaining norm among columns
|
||||
// k..n−1, swaps it into column k and records the swap in perm. The
|
||||
// norms themselves are computed one worker per contiguous column slab,
|
||||
// each column by the serial two-pass accumulation stridedNorm2 always
|
||||
// ran, and the largest is then picked by a serial scan whose strict
|
||||
// comparison keeps the first of equal maxima, so the tie rule is
|
||||
// unchanged. A zero return means every remaining column is exactly zero.
|
||||
//
|
||||
// A slab is a run of consecutive columns, which is what makes the scan
|
||||
// affordable at all: column j and column j+1 share seven of every eight
|
||||
// cache lines their strided walks touch, so a slab re-reads them from
|
||||
// the cache a lone column at a time cannot. The crew only forms when one
|
||||
// worker's slab holds a full dispatch floor of entries.
|
||||
func rrqrPivot(rMat []float64, m, n, k int, perm []int, norms []float64) float64 {
|
||||
cols, rows := n-k, m-k
|
||||
engine.ParallelMin(cols, max((rrqrMinWork+rows-1)/rows, 1), func(start, end int) {
|
||||
for j := start; j < end; j++ {
|
||||
norms[j] = stridedNorm2(rMat, k*n+k+j, n, rows)
|
||||
}
|
||||
})
|
||||
best, bestJ := 0.0, k
|
||||
for j := range cols {
|
||||
if nrm := norms[j]; nrm > best {
|
||||
best, bestJ = nrm, k+j
|
||||
}
|
||||
}
|
||||
if best == 0 {
|
||||
return 0
|
||||
}
|
||||
if bestJ != k {
|
||||
rrqrSwapCols(rMat, m, n, bestJ, k)
|
||||
perm[bestJ], perm[k] = perm[k], perm[bestJ]
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// rrqrApply applies the reflector to the R columns k..n−1 in column
|
||||
// blocks and, when qMat is non-nil, to the m rows of the orthonormal
|
||||
// accumulator, or when bv is non-nil to the tail of the right-hand side.
|
||||
// The targets are disjoint buffers, so they share one crew; the R blocks
|
||||
// and the accumulator rows take contiguous slices, and every dot keeps
|
||||
// its own i-ascending chain with the subtraction repeating the same
|
||||
// operand order, so no crew size and no block boundary moves a bit.
|
||||
func rrqrApply(rMat []float64, m, n, k int, hh householder, qMat, bv []float64) {
|
||||
ln := m - k
|
||||
blocks := (n - k + rrqrColumnBlock - 1) / rrqrColumnBlock
|
||||
items := blocks
|
||||
if qMat != nil {
|
||||
items += m
|
||||
}
|
||||
if bv != nil {
|
||||
items++
|
||||
}
|
||||
weight := (n - k) + m
|
||||
w := min(weight*ln/rrqrMinWork+1, engine.WorkersFor(items))
|
||||
if w < 2 {
|
||||
rrqrApplyWorker(rMat, m, n, k, hh, qMat, bv, blocks, items, 0, 1)
|
||||
return
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for gi := range w {
|
||||
wg.Go(func() {
|
||||
rrqrApplyWorker(rMat, m, n, k, hh, qMat, bv, blocks, items, gi, w)
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// rrqrApplyWorker runs crew member gi of nw over its share of the
|
||||
// reflector's targets: a contiguous slab of the R column blocks, a
|
||||
// contiguous slab of the accumulator rows, and the right-hand side.
|
||||
func rrqrApplyWorker(rMat []float64, m, n, k int, hh householder, qMat, bv []float64, blocks, items, gi, nw int) {
|
||||
ln := m - k
|
||||
// Target layout in the item index space: the R blocks first, then the
|
||||
// accumulator rows when there is an accumulator, then the right-hand
|
||||
// side when there is one.
|
||||
qEnd := blocks
|
||||
if qMat != nil {
|
||||
qEnd += m
|
||||
}
|
||||
var d [rrqrColumnBlock]float64
|
||||
for it := gi * items / nw; it < (gi+1)*items/nw; it++ {
|
||||
switch {
|
||||
case it < blocks:
|
||||
j0 := k + it*rrqrColumnBlock
|
||||
j1 := min(j0+rrqrColumnBlock, n)
|
||||
wd := j1 - j0
|
||||
_ = d[wd-1] // bounds-check proof: the block never exceeds rrqrColumnBlock
|
||||
clear(d[:wd])
|
||||
for jj := range wd {
|
||||
jcol := j0 + jj
|
||||
for i := range ln {
|
||||
d[jj] += hh.v[i] * rMat[(k+i)*n+jcol]
|
||||
}
|
||||
}
|
||||
for j := range wd {
|
||||
d[j] = hh.beta * d[j]
|
||||
}
|
||||
for i := range ln {
|
||||
vi := hh.v[i]
|
||||
base := (k+i)*n + j0
|
||||
for j := range wd {
|
||||
rMat[base+j] -= vi * d[j]
|
||||
}
|
||||
}
|
||||
case it < qEnd:
|
||||
row := it - blocks
|
||||
base := row*m + k
|
||||
dot := 0.0
|
||||
for i := range ln {
|
||||
dot += hh.v[i] * qMat[base+i]
|
||||
}
|
||||
wq := hh.beta * dot
|
||||
for i := range ln {
|
||||
qMat[base+i] -= hh.v[i] * wq
|
||||
}
|
||||
default:
|
||||
dot := 0.0
|
||||
for i := range ln {
|
||||
dot += hh.v[i] * bv[k+i]
|
||||
}
|
||||
wb := hh.beta * dot
|
||||
for i := range ln {
|
||||
bv[k+i] -= hh.v[i] * wb
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// rrqrSwapCols exchanges columns j1 and j2 of a flat m×n matrix.
|
||||
func rrqrSwapCols(rMat []float64, m, n, j1, j2 int) {
|
||||
for i := range m {
|
||||
rMat[i*n+j1], rMat[i*n+j2] = rMat[i*n+j2], rMat[i*n+j1]
|
||||
}
|
||||
}
|
||||
|
||||
// stridedNorm2 returns the L2 norm of count entries walking stride
|
||||
// step from start, summed relative to the largest magnitude so a
|
||||
// column near 1e154 does not overflow on the way to its norm.
|
||||
func stridedNorm2(vals []float64, start, step, count int) float64 {
|
||||
maxAbs := 0.0
|
||||
for i := range count {
|
||||
if a := math.Abs(vals[start+i*step]); a > maxAbs {
|
||||
maxAbs = a
|
||||
}
|
||||
}
|
||||
if maxAbs == 0 {
|
||||
return 0
|
||||
}
|
||||
s := 0.0
|
||||
for i := range count {
|
||||
t := vals[start+i*step] / maxAbs
|
||||
s += t * t
|
||||
}
|
||||
return maxAbs * math.Sqrt(s)
|
||||
}
|
||||
|
||||
// rrqrValidPerm reports whether p is a permutation of 0..n−1, the
|
||||
// property the returned permutation's contract rests on.
|
||||
func rrqrValidPerm(p []int) bool {
|
||||
seen := make([]bool, len(p))
|
||||
for _, v := range p {
|
||||
if v < 0 || v >= len(p) || seen[v] {
|
||||
return false
|
||||
}
|
||||
seen[v] = true
|
||||
}
|
||||
return !slices.Contains(seen, false)
|
||||
}
|
||||
|
||||
// SolveRRQR solves min ‖A·x − b‖₂ through the pivoted factorisation.
|
||||
// On full rank the answer is the ordinary back-substitution through R
|
||||
// against c = Qᵀb. On rank-deficient input it is the minimum-norm
|
||||
// least-squares solution: the rank r is read off the R diagonal decay,
|
||||
// and the leading r×n trapezoidal block R₁ answers its underdetermined
|
||||
// system in minimum norm through x = R₁ᵀ(R₁R₁ᵀ)⁻¹c₁, which squares the
|
||||
// condition of the leading block exactly the way the package's
|
||||
// BᵀB-based SVD squares the spectrum: honest for a rank decision taken
|
||||
// from a decay that has already been observed.
|
||||
func SolveRRQR(a, b *core.Array) (*core.Array, error) {
|
||||
const name = "SolveRRQR"
|
||||
if a.Dtype() == core.Complex {
|
||||
return nil, base.Errf("%s: complex matrices are not supported", name)
|
||||
}
|
||||
if b.Dtype() == core.Complex {
|
||||
return nil, base.Errf("%s: complex right-hand sides are not supported", name)
|
||||
}
|
||||
if a.NDim() != 2 {
|
||||
return nil, base.Errf("%s: needs a 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
m, n := a.Shape()[0], a.Shape()[1]
|
||||
if m == 0 || n == 0 {
|
||||
return nil, base.Errf("%s: zero-sized matrix, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
if m < n {
|
||||
return nil, base.Errf("%s: needs m ≥ n, got shape %s", name, base.ShapeText(a.Shape()))
|
||||
}
|
||||
if b.NDim() != 1 || b.Len() != m {
|
||||
return nil, base.Errf("%s: right-hand side must be a vector of length %d, got shape %s",
|
||||
name, m, base.ShapeText(b.Shape()))
|
||||
}
|
||||
// The sweep below carries b through the reflectors as they are
|
||||
// built, so c = Qᵀb falls out of the factorisation without the
|
||||
// orthogonal factor ever being materialised.
|
||||
rMat := make([]float64, m*n)
|
||||
copy(rMat, denseFloats(a, m, n))
|
||||
bv := vectorF64(b, m)
|
||||
perm := make([]int, n)
|
||||
for i := range perm {
|
||||
perm[i] = i
|
||||
}
|
||||
col := make([]float64, m)
|
||||
norms := make([]float64, n)
|
||||
for k := range n {
|
||||
if rrqrPivot(rMat, m, n, k, perm, norms) == 0 {
|
||||
break
|
||||
}
|
||||
lv := col[:m-k]
|
||||
for i := range lv {
|
||||
lv[i] = rMat[(k+i)*n+k]
|
||||
}
|
||||
hh := householderVectorInto(lv, lv)
|
||||
if hh.beta == 0 {
|
||||
continue
|
||||
}
|
||||
rrqrApply(rMat, m, n, k, hh, nil, bv)
|
||||
}
|
||||
x := make([]float64, n)
|
||||
switch rank := rrqrRankOf(rMat, n); {
|
||||
case rank == n:
|
||||
// Full rank: back-substitution through R.
|
||||
for i := n - 1; i >= 0; i-- {
|
||||
s := bv[i]
|
||||
for j := i + 1; j < n; j++ {
|
||||
s -= rMat[i*n+j] * x[j]
|
||||
}
|
||||
x[i] = s / rMat[i*n+i]
|
||||
}
|
||||
default:
|
||||
// Minimum norm through the leading trapezoidal block R₁:
|
||||
// R₁R₁ᵀw = c₁, x = R₁ᵀw. R₁R₁ᵀ is symmetric positive definite
|
||||
// at the rank the decay showed, and the dense `Solve` carries
|
||||
// the small r×r system.
|
||||
if rank > 0 {
|
||||
mm := make([]float64, rank*rank)
|
||||
for i := range rank {
|
||||
for j := range rank {
|
||||
s := 0.0
|
||||
for l := range n {
|
||||
s += rMat[i*n+l] * rMat[j*n+l]
|
||||
}
|
||||
mm[i*rank+j] = s
|
||||
}
|
||||
}
|
||||
w, err := Solve(floatsToArray(mm, []int{rank, rank}), floatsToArray(append([]float64(nil), bv[:rank]...), []int{rank}))
|
||||
if err != nil {
|
||||
return nil, base.Errf("%s: %w", name, err)
|
||||
}
|
||||
for i := range rank {
|
||||
for j := range n {
|
||||
x[j] += rMat[i*n+j] * w.FloatAt(i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// x lives in permuted coordinates; undo the permutation.
|
||||
out := make([]float64, n)
|
||||
for j := range n {
|
||||
out[perm[j]] = x[j]
|
||||
}
|
||||
return floatsToArray(out, []int{n}), nil
|
||||
}
|
||||
Reference in New Issue
Block a user