328 lines
11 KiB
Go
328 lines
11 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package linalg
|
||
|
||
import (
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// The general (non-Hermitian) sparse eigenproblem. Lanczos
|
||
// needs symmetry: its three-term recurrence exists because the Krylov
|
||
// space of a symmetric operator comes with an orthogonal basis for
|
||
// free. A general operator needs the full Arnoldi recurrence, one
|
||
// explicit orthogonalisation per column, and its projected matrix is
|
||
// upper Hessenberg rather than tridiagonal. Resonance problems in
|
||
// electromagnetics and damped oscillation systems produce exactly
|
||
// these operators, complex and non-Hermitian included.
|
||
//
|
||
// Like SpEigen, this is a one-shot Krylov method: the basis runs to
|
||
// min(n, max(2k, k+40)) columns, the small Hessenberg eigenproblem is
|
||
// solved exactly by the dense EigenGeneral, and the returned pairs are
|
||
// Ritz approximations whose accuracy improves with the budget. The
|
||
// restart-on-breakdown scheme mirrors lanczos: an exhausted direction
|
||
// closes its block with a zero coupling and a fresh orthogonal random
|
||
// direction opens the next.
|
||
|
||
// SpEigenGeneral returns the k eigenvalues of largest magnitude of a
|
||
// general real sparse matrix, each with its unit eigenvector: values
|
||
// is a complex128 vector ordered by descending magnitude (a real
|
||
// matrix may carry complex conjugate pairs), vectors a complex128
|
||
// (n, k) array whose column j belongs to values[j]. The matrix must be
|
||
// square and real; complex input belongs to SpEigenGeneralComplex,
|
||
// symmetric input gets a cheaper answer from SpEigen. gen seeds the
|
||
// start vector (nil uses a fixed seed). Like the dense EigenGeneral,
|
||
// and unlike the symmetry-checking SpEigen, entries are not screened
|
||
// for finiteness: a non-finite entry answers NaN Ritz pairs.
|
||
func SpEigenGeneral(s *core.SparseCOO, k int, gen *core.Generator) (values, vectors *core.Array, err error) {
|
||
const name = "SpEigenGeneral"
|
||
if s.Values.Dtype() == core.Complex {
|
||
return nil, nil, base.Errf("%s: complex matrices belong to SpEigenGeneralComplex", name)
|
||
}
|
||
// The same square 2-D contract SpEigen enforces, and for the same
|
||
// reason: the Krylov product indexes x by the stored column, so a
|
||
// column beyond the row count reads out of range, and the read
|
||
// happens inside a worker goroutine the caller cannot recover.
|
||
// Ranks above two are refused too: the index walk would silently
|
||
// ignore every dimension past the first two.
|
||
if len(s.Shape) != 2 || s.Shape[0] != s.Shape[1] {
|
||
return nil, nil, base.Errf("%s: needs a square 2-D sparse matrix, got shape %v", name, s.Shape)
|
||
}
|
||
c, err := cooToCSR(s, name)
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
return arnoldiEigen(c.matVec, c.n, k, gen, name, krylovF64)
|
||
}
|
||
|
||
// SpEigenGeneralComplex is SpEigenGeneral for complex128 sparse
|
||
// matrices: non-Hermitian operators included, which is the
|
||
// electromagnetics case.
|
||
func SpEigenGeneralComplex(s *core.SparseCOO, k int, gen *core.Generator) (values, vectors *core.Array, err error) {
|
||
const name = "SpEigenGeneralComplex"
|
||
c, err := cooToComplexCSR(s, name)
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
return arnoldiEigen(c.matVec, c.n, k, gen, name, krylovC128)
|
||
}
|
||
|
||
// arnoldiEigen runs the shared pipeline: Arnoldi basis, dense
|
||
// eigensolve of the projected Hessenberg, Ritz lift. The matvec
|
||
// closure hides whether the operator is real or complex, and kern
|
||
// supplies the element-type-specific primitives the basis arithmetic
|
||
// needs.
|
||
func arnoldiEigen[T scalar](matvec func(x, y []T), n, k int, gen *core.Generator, name string, kern krylovKernel[T]) (*core.Array, *core.Array, error) {
|
||
if n == 0 {
|
||
return nil, nil, base.Errf("%s: zero-sized matrix", name)
|
||
}
|
||
if k < 1 || k > n {
|
||
return nil, nil, base.Errf("%s: k must be in [1, %d], got %d", name, n, k)
|
||
}
|
||
if gen == nil {
|
||
gen = core.NewGenerator(spEigenSeed)
|
||
}
|
||
m := min(n, max(2*k, k+spEigenBlock))
|
||
h, v := arnoldi(matvec, n, m, gen, kern)
|
||
|
||
// The projected problem: the leading m×m block of the Hessenberg,
|
||
// dense and small enough for the general eigensolver.
|
||
hArr, err := zeros(core.Complex, []int{m, m})
|
||
if err != nil {
|
||
return nil, nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
hc := hArr.RawComplexes()
|
||
for j := range m {
|
||
for i := 0; i <= j+1 && i < m; i++ {
|
||
hc[i*m+j] = toComplex(h[i*m+j])
|
||
}
|
||
}
|
||
vals, vecs, err := EigenGeneral(hArr)
|
||
if err != nil {
|
||
return nil, nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
if k > vals.Len() {
|
||
k = vals.Len()
|
||
}
|
||
outVals, err := zeros(core.Complex, []int{k})
|
||
if err != nil {
|
||
return nil, nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
vc := outVals.RawComplexes()
|
||
outVecs := make([]complex128, n*k)
|
||
tmp := make([]complex128, n)
|
||
for j := range k {
|
||
vc[j] = vals.ComplexAt(j)
|
||
// Ritz vector: lift the Hessenberg eigenvector through the
|
||
// Arnoldi basis, u = V·y.
|
||
clear(tmp)
|
||
for p := range m {
|
||
y := vecs.ComplexAt(p*m + j)
|
||
if y == 0 {
|
||
continue
|
||
}
|
||
kern.widenAxpy(tmp, v[p*n:(p+1)*n], y)
|
||
}
|
||
norm := 0.0
|
||
for _, z := range tmp {
|
||
norm += real(z)*real(z) + imag(z)*imag(z)
|
||
}
|
||
norm = math.Sqrt(norm)
|
||
if norm > 0 {
|
||
for i := range n {
|
||
tmp[i] /= complex(norm, 0)
|
||
}
|
||
}
|
||
for i := range n {
|
||
outVecs[i*k+j] = tmp[i]
|
||
}
|
||
}
|
||
ritzVecs, err := complexFromArray2D(outVecs, n, k)
|
||
if err != nil {
|
||
return nil, nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
return outVals, ritzVecs, nil
|
||
}
|
||
|
||
// toComplex widens a scalar kernel value for the complex lift.
|
||
func toComplex[T scalar](v T) complex128 {
|
||
switch x := any(v).(type) {
|
||
case float64:
|
||
return complex(x, 0)
|
||
case complex128:
|
||
return x
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// arnoldi builds the Krylov basis v (columns of length n, m+1 of
|
||
// them) and the upper Hessenberg h ((m+1)×m, row-major) of the
|
||
// operator behind matvec. Full reorthogonalisation, twice per column,
|
||
// keeps the basis orthonormal to rounding level; a collapsed direction
|
||
// closes its block with a zero coupling and restarts in a fresh
|
||
// orthogonal random direction.
|
||
func arnoldi[T scalar](matvec func(x, y []T), n, m int, gen *core.Generator, kern krylovKernel[T]) (h, v []T) {
|
||
// The basis never exceeds the space it spans: m above n would let
|
||
// the truncation branch below hand back a shorter h than the m×m
|
||
// read loop expects.
|
||
if m > n {
|
||
panic("arnoldi: the block size exceeds the matrix order")
|
||
}
|
||
h = make([]T, (m+1)*m)
|
||
v = make([]T, n*(m+1))
|
||
w := make([]T, n)
|
||
scale := 0.0
|
||
addScale := func(z T) {
|
||
if a := absOf(z); a > scale {
|
||
scale = a
|
||
}
|
||
}
|
||
// normT is the one accumulation in the recurrence that squares its
|
||
// inputs. A column whose entries are above about 1.3e154 overflows
|
||
// the raw sum of squares to +Inf (the projected coupling becomes
|
||
// infinite and the Ritz values come back wrong without a word), and
|
||
// one below about 1.5e-162 underflows it to zero (the exhaustion
|
||
// test reads that as a collapsed Krylov block and closes the block
|
||
// on a direction the operator never exhausted). The sum therefore
|
||
// runs raw while it lands in the normal range, which keeps the
|
||
// ordinary-scale arithmetic bit-for-bit, and falls back to a
|
||
// max-scaled accumulation outside it, the way norm2F64 and normC
|
||
// do. NaN is carried through the raw path so a poisoned input stays
|
||
// poisoned rather than reading as an exhausted block.
|
||
normT := func(a []T) float64 {
|
||
if s := absOf(kern.dot(a, a)); s <= math.MaxFloat64 && (s >= 0x1p-1022 || math.IsNaN(s)) {
|
||
return math.Sqrt(s)
|
||
}
|
||
maxAbs := 0.0
|
||
for _, z := range a {
|
||
if m := absOf(z); m > maxAbs {
|
||
maxAbs = m
|
||
}
|
||
}
|
||
if maxAbs == 0 {
|
||
return 0
|
||
}
|
||
s := 0.0
|
||
for _, z := range a {
|
||
r := absOf(z) / maxAbs
|
||
s += r * r
|
||
}
|
||
return maxAbs * math.Sqrt(s)
|
||
}
|
||
start := func(col int) {
|
||
q := v[col*n : (col+1)*n]
|
||
kern.fillRand(q, gen)
|
||
for range 2 {
|
||
for p := range col {
|
||
row := v[p*n : (p+1)*n]
|
||
d := kern.dot(row, q)
|
||
for i := range n {
|
||
q[i] -= d * row[i]
|
||
}
|
||
}
|
||
}
|
||
if norm := normT(q); norm > 0 {
|
||
kern.scaleInto(q, q, 1/norm)
|
||
}
|
||
}
|
||
start(0)
|
||
for j := range m {
|
||
q := v[j*n : (j+1)*n]
|
||
matvec(q, w)
|
||
for i := 0; i <= j; i++ {
|
||
row := v[i*n : (i+1)*n]
|
||
d := kern.dot(row, w)
|
||
h[i*m+j] = d
|
||
for t := range n {
|
||
w[t] -= d * row[t]
|
||
}
|
||
}
|
||
// Full reorthogonalisation, twice: the first pass removes the
|
||
// accumulated loss, the second what the first reintroduces.
|
||
for range 2 {
|
||
for i := 0; i <= j; i++ {
|
||
row := v[i*n : (i+1)*n]
|
||
d := kern.dot(row, w)
|
||
h[i*m+j] += d
|
||
for t := range n {
|
||
w[t] -= d * row[t]
|
||
}
|
||
}
|
||
}
|
||
beta := normT(w)
|
||
addScale(h[j*m+j])
|
||
if j+1 < m {
|
||
if beta <= float64(n)*base.EpsF*scale {
|
||
// The Krylov space of this block is exhausted; close it
|
||
// with a zero coupling and restart in a fresh
|
||
// direction. m never exceeds n (asserted at the top),
|
||
// so the basis always has room for the restart column.
|
||
start(j + 1)
|
||
continue
|
||
}
|
||
h[(j+1)*m+j] = realT[T](beta)
|
||
kern.scaleInto(v[(j+1)*n:(j+2)*n], w, 1/beta)
|
||
}
|
||
}
|
||
return h, v
|
||
}
|
||
|
||
// krylovKernel holds the element-type-specific primitives the shared
|
||
// Arnoldi recurrence calls once per vector pass: the conjugated inner
|
||
// product, the real scaling of a vector, the start-vector fill and the
|
||
// widening lift of a real basis row into the complex Ritz vector.
|
||
// Passing them in keeps one copy of the recurrence's arithmetic while
|
||
// the element type is dispatched on once per pass instead of once per
|
||
// element.
|
||
type krylovKernel[T scalar] struct {
|
||
dot func(a, b []T) T
|
||
scaleInto func(dst, src []T, s float64)
|
||
fillRand func(dst []T, gen *core.Generator)
|
||
widenAxpy func(dst []complex128, src []T, y complex128)
|
||
}
|
||
|
||
// krylovF64 is the kernel of a real operator.
|
||
var krylovF64 = krylovKernel[float64]{
|
||
dot: dotF64,
|
||
scaleInto: func(dst, src []float64, s float64) {
|
||
for i := range dst {
|
||
dst[i] = src[i] * s
|
||
}
|
||
},
|
||
fillRand: func(dst []float64, gen *core.Generator) {
|
||
for i := range dst {
|
||
dst[i] = gen.NormalUnit()
|
||
}
|
||
},
|
||
widenAxpy: func(dst []complex128, src []float64, y complex128) {
|
||
for i, v := range src {
|
||
dst[i] += y * complex(v, 0)
|
||
}
|
||
},
|
||
}
|
||
|
||
// krylovC128 is the kernel of a complex operator: the inner product
|
||
// conjugates its left operand, and the start vector draws a real and an
|
||
// imaginary part from the generator, in that order.
|
||
var krylovC128 = krylovKernel[complex128]{
|
||
dot: dotC,
|
||
scaleInto: func(dst, src []complex128, s float64) {
|
||
for i := range dst {
|
||
dst[i] = src[i] * complex(s, 0)
|
||
}
|
||
},
|
||
fillRand: func(dst []complex128, gen *core.Generator) {
|
||
for i := range dst {
|
||
dst[i] = complex(gen.NormalUnit(), gen.NormalUnit())
|
||
}
|
||
},
|
||
widenAxpy: func(dst []complex128, src []complex128, y complex128) {
|
||
for i, v := range src {
|
||
dst[i] += y * v
|
||
}
|
||
},
|
||
}
|