447 lines
14 KiB
Go
447 lines
14 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package linalg
|
||
|
||
import (
|
||
"math"
|
||
"math/cmplx"
|
||
"slices"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// The general eigenproblem. Eigen answers real symmetric matrices and
|
||
// EigenComplex Hermitian ones, where structure pays for itself. A
|
||
// general square matrix has no structure to exploit, so EigenGeneral
|
||
// goes through the complex Schur form: the input is promoted to
|
||
// complex128, reduced to upper Hessenberg form by complex Householder
|
||
// reflectors, and driven to upper triangular form by the shifted QR
|
||
// iteration, one Wilkinson-shifted sweep of Givens rotations per step,
|
||
// deflating whenever a subdiagonal entry falls under a scale-relative
|
||
// tolerance. Working over the complex plane is what keeps the
|
||
// iteration single-shift: conjugate eigenvalue pairs are ordinary
|
||
// points there, where a real iteration would need the double-shift
|
||
// bulge chase to reach them.
|
||
//
|
||
// Eigenvectors come from the triangular Schur factor by back
|
||
// substitution and one multiplication with the accumulated unitary
|
||
// similarity, which keeps them orthonormal to rounding level even
|
||
// when the eigenvalues themselves are ill conditioned.
|
||
|
||
// EigenGeneral returns the eigenvalues and eigenvectors of a square
|
||
// matrix of any dtype; real and integer inputs are promoted to
|
||
// complex128. Values is a complex128 vector sorted descending by
|
||
// magnitude, bitwise ties broken by descending real part and then
|
||
// descending imaginary part; conjugate pairs share a magnitude only
|
||
// to rounding, so their relative order follows the rounding noise
|
||
// rather than the tiebreak. Vectors is a complex128 (n, n) array
|
||
// whose column j is the unit eigenvector for values[j]. Symmetric
|
||
// real matrices get a faster answer from Eigen and Hermitian matrices
|
||
// from EigenComplex.
|
||
func EigenGeneral(a *core.Array) (values, vectors *core.Array, err error) {
|
||
if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
|
||
return nil, nil, base.Errf("EigenGeneral: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
|
||
}
|
||
n := a.Shape()[0]
|
||
if n == 0 {
|
||
return nil, nil, base.Errf("EigenGeneral: zero-sized matrix, got shape %s", base.ShapeText(a.Shape()))
|
||
}
|
||
h := make([]complex128, n*n)
|
||
switch {
|
||
case a.Dtype() == core.Complex && !a.Strided() && len(a.RawComplexes()) == n*n:
|
||
copy(h, a.RawComplexes())
|
||
case a.Dtype() == core.Float && !a.Strided() && len(a.RawFloats()) == n*n:
|
||
for i := range n * n {
|
||
h[i] = complex(a.RawFloats()[i], 0)
|
||
}
|
||
default:
|
||
for i := range n * n {
|
||
h[i] = a.ComplexAt(i)
|
||
}
|
||
}
|
||
// Outside the safe window the reflector norms, the norm the shift
|
||
// floors are taken against and the squared magnitudes the sweeps form
|
||
// all leave the normal range: a tiny matrix collapses to zeros and a
|
||
// huge one has vMax*vMax overflow, which zeroes beta so the
|
||
// Hessenberg reduction is skipped and the QR iteration then fails on
|
||
// a perfectly valid matrix. The matrix is moved into the window and
|
||
// the eigenvalues, which carry the scale, are moved back on return;
|
||
// the Schur vectors are scale-free.
|
||
ws := windowScale(maxMagComplex(h))
|
||
if ws != 1 {
|
||
scaleComplexes(h, ws)
|
||
}
|
||
scale := 0.0
|
||
for _, z := range h {
|
||
if m := base.AbsComplex(z); m > scale {
|
||
scale = m
|
||
}
|
||
}
|
||
// q accumulates the unitary similarity: the input satisfies
|
||
// A = q·H·qᴴ at every stage.
|
||
q := eyeComplex(n)
|
||
hessenbergComplex(h, q, n)
|
||
if err := hessenbergQr(h, q, n, scale); err != nil {
|
||
return nil, nil, err
|
||
}
|
||
|
||
vals := make([]complex128, n)
|
||
for i := range n {
|
||
vals[i] = h[i*n+i]
|
||
}
|
||
idx := sortMagDescIndices(vals)
|
||
outVals := make([]complex128, n)
|
||
for j := range n {
|
||
outVals[j] = vals[idx[j]]
|
||
}
|
||
// Eigenvector j depends on the position of λ_j inside the Schur
|
||
// factor, so they are built in place order and permuted afterwards.
|
||
inPlace := make([]complex128, n*n)
|
||
for j := range n {
|
||
v := schurEigenVector(h, q, n, j, scale)
|
||
for i := range n {
|
||
inPlace[i*n+j] = v[i]
|
||
}
|
||
}
|
||
outVecs := make([]complex128, n*n)
|
||
for j := range n {
|
||
for i := range n {
|
||
outVecs[i*n+j] = inPlace[i*n+idx[j]]
|
||
}
|
||
}
|
||
// Only the eigenvalues carry the matrix's scale.
|
||
if ws != 1 {
|
||
unscaleComplexes(outVals, ws)
|
||
}
|
||
valuesArr := core.New(core.Complex, []int{n}...)
|
||
copy(valuesArr.RawComplexes(), outVals)
|
||
vecArr := core.New(core.Complex, []int{n, n}...)
|
||
copy(vecArr.RawComplexes(), outVecs)
|
||
return valuesArr, vecArr, nil
|
||
}
|
||
|
||
// hessenbergComplex reduces h to upper Hessenberg form in place by
|
||
// complex Householder reflectors, accumulating the unitary similarity
|
||
// into q so that the original matrix satisfies A = q·H·qᴴ.
|
||
func hessenbergComplex(h, q []complex128, n int) {
|
||
for k := range n - 2 {
|
||
// The reflector zeroes column k below its subdiagonal; columns
|
||
// already in form are skipped. Squared magnitudes are summed
|
||
// relative to the column's largest entry so a column with
|
||
// entries near 1e154 cannot overflow on the way to its norm.
|
||
scale := 0.0
|
||
for i := k + 1; i < n; i++ {
|
||
if m := base.AbsComplex(h[i*n+k]); m > scale {
|
||
scale = m
|
||
}
|
||
}
|
||
below := 0.0
|
||
if scale > 0 {
|
||
for i := k + 2; i < n; i++ {
|
||
m := base.AbsComplex(h[i*n+k]) / scale
|
||
below += m * m
|
||
}
|
||
}
|
||
if below == 0 {
|
||
continue
|
||
}
|
||
x0 := h[(k+1)*n+k]
|
||
norm := scale * math.Hypot(base.AbsComplex(x0)/scale, math.Sqrt(below))
|
||
// alpha = −sign(x0)·‖x‖ lands the reflection on the far side
|
||
// of x0, away from the cancellation zone.
|
||
phase := complex(1, 0)
|
||
if m := base.AbsComplex(x0); m > 0 {
|
||
phase = x0 / complex(m, 0)
|
||
}
|
||
alpha := -phase * complex(norm, 0)
|
||
v := make([]complex128, n-k-1)
|
||
v[0] = x0 - alpha
|
||
for i := 1; i < n-k-1; i++ {
|
||
v[i] = h[(k+1+i)*n+k]
|
||
}
|
||
vMax := 0.0
|
||
for _, z := range v {
|
||
if m := base.AbsComplex(z); m > vMax {
|
||
vMax = m
|
||
}
|
||
}
|
||
vNorm2 := 0.0
|
||
if vMax > 0 {
|
||
for _, z := range v {
|
||
re := real(z) / vMax
|
||
im := imag(z) / vMax
|
||
vNorm2 += re*re + im*im
|
||
}
|
||
}
|
||
// beta = 2/(vᴴv) needs vMax*vMax to stay in the normal range. A
|
||
// column whose entries are extreme in the matrix's own units
|
||
// (far below it, say) overflows that square to +Inf, which makes
|
||
// beta +0, or underflows it to 0, which makes beta +Inf and the
|
||
// update NaN. Both are silent corruption of an otherwise valid
|
||
// reduction, so the reflector is re-expressed in window units
|
||
// with the scaling compensated in beta, exactly as
|
||
// householderVectorInto does for the real reflectors.
|
||
betaR := 2 / (vMax * vMax * vNorm2)
|
||
if betaR <= 0 || math.IsInf(betaR, 0) {
|
||
wsV := windowScale(vMax)
|
||
for i := range v {
|
||
v[i] = complex(real(v[i])*wsV, imag(v[i])*wsV)
|
||
}
|
||
w := vMax * wsV
|
||
betaR = 2 / (w * w * vNorm2)
|
||
}
|
||
beta := complex(betaR, 0)
|
||
|
||
// H = P·H with P = I − β·v·vᴴ, rows k+1..n−1, columns k..n−1.
|
||
for j := k; j < n; j++ {
|
||
s := complex(0, 0)
|
||
for i := range v {
|
||
s += cmplx.Conj(v[i]) * h[(k+1+i)*n+j]
|
||
}
|
||
s *= beta
|
||
for i := range v {
|
||
h[(k+1+i)*n+j] -= s * v[i]
|
||
}
|
||
}
|
||
// H = H·P, every row, columns k+1..n−1.
|
||
for i := range n {
|
||
s := complex(0, 0)
|
||
for j := range v {
|
||
s += h[i*n+k+1+j] * v[j]
|
||
}
|
||
s *= beta
|
||
for j := range v {
|
||
h[i*n+k+1+j] -= s * cmplx.Conj(v[j])
|
||
}
|
||
}
|
||
// q = q·P keeps the similarity: A = q·H·qᴴ.
|
||
for i := range n {
|
||
s := complex(0, 0)
|
||
for j := range v {
|
||
s += q[i*n+k+1+j] * v[j]
|
||
}
|
||
s *= beta
|
||
for j := range v {
|
||
q[i*n+k+1+j] -= s * cmplx.Conj(v[j])
|
||
}
|
||
}
|
||
// Column k below the subdiagonal is zero by construction; make
|
||
// it exactly so instead of carrying rounding residue.
|
||
for i := k + 2; i < n; i++ {
|
||
h[i*n+k] = 0
|
||
}
|
||
}
|
||
}
|
||
|
||
// hessenbergQr drives the shifted QR iteration on an upper Hessenberg
|
||
// matrix until every subdiagonal entry deflates, accumulating the
|
||
// Schur similarity into q. On return h is upper triangular: the Schur
|
||
// form of the original matrix, with its eigenvalues on the diagonal.
|
||
func hessenbergQr(h, q []complex128, n int, scale float64) error {
|
||
// Purely relative to the matrix magnitude: an absolute floor would
|
||
// treat a legitimate tiny-scale matrix (a 1e-20 rotation, say) as
|
||
// one big deflated block and return zeros.
|
||
floor := base.EpsF * scale
|
||
negligible := func(i int) bool {
|
||
local := base.EpsF * (base.AbsComplex(h[i*n+i]) + base.AbsComplex(h[(i-1)*n+i-1]))
|
||
return base.AbsComplex(h[i*n+i-1]) <= math.Max(local, floor)
|
||
}
|
||
// Rotation scratch reused across sweeps; each sweep uses the first
|
||
// hi−lo entries of each.
|
||
cs := make([]complex128, n)
|
||
sn := make([]complex128, n)
|
||
hi := n - 1
|
||
iter := 0
|
||
for hi > 0 {
|
||
for hi > 0 && negligible(hi) {
|
||
h[hi*n+hi-1] = 0
|
||
hi--
|
||
iter = 0
|
||
}
|
||
if hi == 0 {
|
||
return nil
|
||
}
|
||
lo := hi
|
||
for lo > 0 && !negligible(lo) {
|
||
lo--
|
||
}
|
||
if lo > 0 {
|
||
// A negligible entry splits the matrix; the block above it
|
||
// is settled in later rounds.
|
||
h[lo*n+lo-1] = 0
|
||
}
|
||
iter++
|
||
if iter > 100 {
|
||
return base.Errf("EigenGeneral: QR iteration failed to converge at row %d", hi)
|
||
}
|
||
var mu complex128
|
||
if iter%10 == 0 {
|
||
// An exceptional shift breaks the rare cyclic pattern the
|
||
// Wilkinson shift can settle into.
|
||
mu = h[hi*n+hi] + complex(0.75*base.AbsComplex(h[hi*n+hi-1]), 0)
|
||
} else {
|
||
mu = wilkinsonShift(h, n, hi)
|
||
}
|
||
qrSweepComplex(h, q, n, lo, hi, mu, cs, sn)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// wilkinsonShift returns the eigenvalue of the trailing 2×2 block
|
||
// closest to its bottom-right entry: the shift that deflates the
|
||
// subdiagonal under it fastest.
|
||
func wilkinsonShift(h []complex128, n, i int) complex128 {
|
||
a := h[(i-1)*n+i-1]
|
||
b := h[(i-1)*n+i]
|
||
c := h[i*n+i-1]
|
||
d := h[i*n+i]
|
||
delta := (a - d) / 2
|
||
disc := cmplx.Sqrt(delta*delta + b*c)
|
||
low := d + delta - disc
|
||
high := d + delta + disc
|
||
if base.AbsComplex(low-d) <= base.AbsComplex(high-d) {
|
||
return low
|
||
}
|
||
return high
|
||
}
|
||
|
||
// qrSweepComplex performs one explicit single-shift QR step over the
|
||
// active block [lo, hi]: H = G·(H − μI)·Gᴴ + μI for the sequence of
|
||
// Givens rotations G that triangularises H − μI, with every rotation
|
||
// folded into q so the similarity stays exact. Each rotation satisfies
|
||
// G·[a; b] = [r; 0] with r real, which is what zeroes the subdiagonal
|
||
// one entry per step. cs and sn are the caller's scratch for the
|
||
// rotations; each sweep fully overwrites its first hi−lo entries.
|
||
func qrSweepComplex(h, q []complex128, n, lo, hi int, mu complex128, cs, sn []complex128) {
|
||
for i := lo; i <= hi; i++ {
|
||
h[i*n+i] -= mu
|
||
}
|
||
for j := lo; j < hi; j++ {
|
||
a, b := h[j*n+j], h[(j+1)*n+j]
|
||
r := math.Hypot(base.AbsComplex(a), base.AbsComplex(b))
|
||
var c, s complex128
|
||
if r == 0 {
|
||
c, s = 1, 0
|
||
} else {
|
||
c, s = a/complex(r, 0), b/complex(r, 0)
|
||
}
|
||
cs[j-lo], sn[j-lo] = c, s
|
||
// Left: rows j and j+1 over columns j..n−1.
|
||
for k := j; k < n; k++ {
|
||
x, y := h[j*n+k], h[(j+1)*n+k]
|
||
h[j*n+k] = cmplx.Conj(c)*x + cmplx.Conj(s)*y
|
||
h[(j+1)*n+k] = -s*x + c*y
|
||
}
|
||
h[(j+1)*n+j] = 0
|
||
}
|
||
for j := lo; j < hi; j++ {
|
||
c, s := cs[j-lo], sn[j-lo]
|
||
// Right: columns j and j+1 down to row j+1, the deepest row
|
||
// the triangular factor can reach in either column.
|
||
for i := 0; i <= j+1; i++ {
|
||
x, y := h[i*n+j], h[i*n+j+1]
|
||
h[i*n+j] = x*c + y*s
|
||
h[i*n+j+1] = -x*cmplx.Conj(s) + y*cmplx.Conj(c)
|
||
}
|
||
for i := range n {
|
||
x, y := q[i*n+j], q[i*n+j+1]
|
||
q[i*n+j] = x*c + y*s
|
||
q[i*n+j+1] = -x*cmplx.Conj(s) + y*cmplx.Conj(c)
|
||
}
|
||
}
|
||
for i := lo; i <= hi; i++ {
|
||
h[i*n+i] += mu
|
||
}
|
||
}
|
||
|
||
// schurEigenVector builds the unit eigenvector for the diagonal entry
|
||
// j of a triangular Schur factor t with its similarity q (the input
|
||
// satisfies A = q·t·qᴴ): back substitution on the leading block gives
|
||
// the coordinates in Schur space, one multiplication with q lifts
|
||
// them back. A nearly multiple eigenvalue perturbs the denominator
|
||
// off exact zero, the standard defence against dividing by the
|
||
// spectrum's own degeneracy.
|
||
func schurEigenVector(t, q []complex128, n, j int, scale float64) []complex128 {
|
||
x := make([]complex128, n)
|
||
x[j] = 1
|
||
lam := t[j*n+j]
|
||
// Relative to the matrix magnitude for the same reason as the QR
|
||
// floor above: tiny-scale spectra deserve their eigenvectors too.
|
||
denFloor := base.EpsF * scale
|
||
for i := j - 1; i >= 0; i-- {
|
||
s := -t[i*n+j]
|
||
for k := i + 1; k < j; k++ {
|
||
s -= t[i*n+k] * x[k]
|
||
}
|
||
d := t[i*n+i] - lam
|
||
if base.AbsComplex(d) < denFloor {
|
||
d = complex(denFloor, 0)
|
||
}
|
||
if base.AbsComplex(d) == 0 {
|
||
// The zero matrix: every denominator vanishes, and the
|
||
// coordinate is unconstrained. Zero keeps the vector finite
|
||
// where 0/0 would poison it with NaN.
|
||
x[i] = 0
|
||
continue
|
||
}
|
||
x[i] = s / d
|
||
}
|
||
v := make([]complex128, n)
|
||
for i := range n {
|
||
acc := complex(0, 0)
|
||
for k := range j + 1 {
|
||
acc += q[i*n+k] * x[k]
|
||
}
|
||
v[i] = acc
|
||
}
|
||
norm := 0.0
|
||
for _, z := range v {
|
||
norm += real(z)*real(z) + imag(z)*imag(z)
|
||
}
|
||
norm = math.Sqrt(norm)
|
||
if norm > 0 {
|
||
for i := range v {
|
||
v[i] /= complex(norm, 0)
|
||
}
|
||
}
|
||
return v
|
||
}
|
||
|
||
// sortMagDescIndices sorts indices so the values come descending by
|
||
// magnitude, bitwise ties broken by descending real part and then
|
||
// descending imaginary part. The sort is stable, so ties beyond the
|
||
// comparator keep their original order exactly as before.
|
||
func sortMagDescIndices(vals []complex128) []int {
|
||
idx := make([]int, len(vals))
|
||
for i := range idx {
|
||
idx[i] = i
|
||
}
|
||
slices.SortStableFunc(idx, func(a, b int) int {
|
||
ma, mb := base.AbsComplex(vals[a]), base.AbsComplex(vals[b])
|
||
if ma != mb {
|
||
if ma > mb {
|
||
return -1
|
||
}
|
||
return 1
|
||
}
|
||
ra, rb := real(vals[a]), real(vals[b])
|
||
if ra != rb {
|
||
if ra > rb {
|
||
return -1
|
||
}
|
||
return 1
|
||
}
|
||
ia, ib := imag(vals[a]), imag(vals[b])
|
||
switch {
|
||
case ia > ib:
|
||
return -1
|
||
case ia < ib:
|
||
return 1
|
||
default:
|
||
return 0
|
||
}
|
||
})
|
||
return idx
|
||
}
|