746 lines
22 KiB
Go
746 lines
22 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"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
)
|
||
|
||
// Complex sparse linear algebra. The real sparse surface
|
||
// stores float64 payloads; the complex side mirrors it with a complex
|
||
// CSR whose values are complex128, built from a core.SparseCOO holding
|
||
// complex values (NewSparseCOO accepts them; SparseFrom does not, so
|
||
// callers assemble the COO directly). The solvers are complex
|
||
// conjugate gradient for Hermitian positive-definite systems, complex
|
||
// BiCGSTAB for general ones, and a Hermitian Lanczos eigensolver; the
|
||
// electromagnetics Helmholtz problems that motivated this surface hit
|
||
// all three.
|
||
|
||
// complexCSR is the complex128 compressed sparse row form.
|
||
type complexCSR struct {
|
||
rowStart []int
|
||
colIdx []int
|
||
vals []complex128
|
||
n int
|
||
}
|
||
|
||
// cEntry is one coordinate entry on the way into CSR form.
|
||
type cEntry struct {
|
||
row, col int
|
||
val complex128
|
||
}
|
||
|
||
// stableOrderComplexEntries orders entries by (row, col) and returns
|
||
// the slice holding the result. Two stable counting passes reach the
|
||
// order a stable comparison sort by (row, col) reaches, in linear time:
|
||
// the first places by column, the second by row, so entries that share
|
||
// both coordinates keep their input order and their values accumulate
|
||
// in it.
|
||
func stableOrderComplexEntries(entries []cEntry, rows, cols int) []cEntry {
|
||
buf := make([]cEntry, len(entries))
|
||
count := make([]int, max(rows, cols)+1)
|
||
countingPlace(buf, entries, count, func(e cEntry) int { return e.col })
|
||
countingPlace(entries, buf, count, func(e cEntry) int { return e.row })
|
||
return entries
|
||
}
|
||
|
||
// cooToComplexCSR converts a complex-valued SparseCOO to CSR form,
|
||
// merging duplicate coordinates and dropping explicit zeros. The
|
||
// values must be complex128; anything else is an error naming the
|
||
// dtype.
|
||
func cooToComplexCSR(s *core.SparseCOO, name string) (*complexCSR, error) {
|
||
if s.Values.Dtype() != core.Complex {
|
||
return nil, base.Errf("%s: needs complex128 values, got %s", name, s.Values.Dtype())
|
||
}
|
||
if len(s.Shape) != 2 {
|
||
return nil, base.Errf("%s: needs a 2-D matrix, got rank %d", name, len(s.Shape))
|
||
}
|
||
rows, cols := s.Shape[0], s.Shape[1]
|
||
if rows != cols {
|
||
return nil, base.Errf("%s: needs a square matrix, got %d×%d", name, rows, cols)
|
||
}
|
||
nnz := s.Indices.Shape()[0]
|
||
idx := s.Indices.RawInts()
|
||
entries := make([]cEntry, nnz)
|
||
for i := range nnz {
|
||
r := int(idx[i*2])
|
||
c := int(idx[i*2+1])
|
||
if r < 0 || r >= rows || c < 0 || c >= cols {
|
||
return nil, base.Errf("%s: index [%d,%d] out of range for %d×%d", name, r, c, rows, cols)
|
||
}
|
||
entries[i] = cEntry{r, c, s.Values.ComplexAt(i)}
|
||
}
|
||
// Sort by (row, col) so duplicates merge and rows are contiguous,
|
||
// keeping equal coordinates in COO order.
|
||
sorted := stableOrderComplexEntries(entries, rows, cols)
|
||
cm := &complexCSR{
|
||
rowStart: make([]int, rows+1),
|
||
colIdx: make([]int, 0, nnz),
|
||
vals: make([]complex128, 0, nnz),
|
||
n: rows,
|
||
}
|
||
// The duplicates are adjacent after the sort, so each run merges into
|
||
// one accumulated value on the way into the compressed form.
|
||
for p := 0; p < len(sorted); {
|
||
e := sorted[p]
|
||
v := e.val
|
||
p++
|
||
for p < len(sorted) && sorted[p].row == e.row && sorted[p].col == e.col {
|
||
v += sorted[p].val
|
||
p++
|
||
}
|
||
if v == 0 {
|
||
continue
|
||
}
|
||
cm.colIdx = append(cm.colIdx, e.col)
|
||
cm.vals = append(cm.vals, v)
|
||
cm.rowStart[e.row+1]++
|
||
}
|
||
for i := range cm.n {
|
||
cm.rowStart[i+1] += cm.rowStart[i]
|
||
}
|
||
return cm, nil
|
||
}
|
||
|
||
// checkHermitian verifies every stored entry has a conjugate mirror,
|
||
// the property the complex CG solver and the Hermitian Lanczos rely
|
||
// on. The tolerance is purely relative to the largest magnitude present,
|
||
// deliberately without an absolute floor: a floor would approve a matrix
|
||
// of small scale whose asymmetry is a large fraction of it, and the
|
||
// Jacobi and Lanczos routes would then answer for a matrix that is
|
||
// neither A nor Aᴴ. Relative-only keeps the approval and the arithmetic
|
||
// in agreement. A non-finite entry is refused in its own right: NaN
|
||
// compares unequal to everything, so the mirror test alone would wave
|
||
// it through as Hermitian. The wording of the refusal is unchanged.
|
||
func (c *complexCSR) checkHermitian(name string) error {
|
||
scale := 0.0
|
||
for _, v := range c.vals {
|
||
if a := cmplxAbs(v); a > scale {
|
||
scale = a
|
||
}
|
||
}
|
||
tol := 1e-12 * scale
|
||
for i := range c.n {
|
||
for p := c.rowStart[i]; p < c.rowStart[i+1]; p++ {
|
||
j := c.colIdx[p]
|
||
v := c.vals[p]
|
||
if math.IsNaN(real(v)) || math.IsNaN(imag(v)) ||
|
||
math.IsInf(real(v), 0) || math.IsInf(imag(v), 0) {
|
||
return base.Errf("%s: entry [%d,%d] is not finite", name, i, j)
|
||
}
|
||
mirror, ok := c.at(j, i)
|
||
if !ok || cmplxAbs(mirror-complexConj(v)) > tol {
|
||
return base.Errf("%s: matrix is not Hermitian within 1e-12 tolerance", name)
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func complexConj(z complex128) complex128 { return complex(real(z), -imag(z)) }
|
||
|
||
func cmplxAbs(z complex128) float64 { return math.Hypot(real(z), imag(z)) }
|
||
|
||
// at reads the entry (row, col), reporting whether it is stored. The
|
||
// row's column indices are sorted, so the lookup is a binary search:
|
||
// the Hermitian check calls it once per stored entry, and a linear
|
||
// scan made that check quadratic in the row length.
|
||
func (c *complexCSR) at(row, col int) (complex128, bool) {
|
||
lo, hi := c.rowStart[row], c.rowStart[row+1]
|
||
for lo < hi {
|
||
mid := int(uint(lo+hi) >> 1)
|
||
if c.colIdx[mid] < col {
|
||
lo = mid + 1
|
||
} else {
|
||
hi = mid
|
||
}
|
||
}
|
||
if lo < c.rowStart[row+1] && c.colIdx[lo] == col {
|
||
return c.vals[lo], true
|
||
}
|
||
return 0, false
|
||
}
|
||
|
||
// complexMatVecRange computes y[i] = A[i,:]·x for the rows in [s, e) of
|
||
// a matrix in compressed row form. It is the whole body of the product,
|
||
// so the serial path calls it directly and the parallel path calls it
|
||
// per chunk: taking the structure slices as arguments rather than as
|
||
// captured variables keeps the serial path free of the closure a row
|
||
// split needs, which the complex solvers call once per iteration.
|
||
func complexMatVecRange(s, e int, rowStart, colIdx []int, values, x, y []complex128) {
|
||
for i := s; i < e; i++ {
|
||
sum := 0i
|
||
for p := rowStart[i]; p < rowStart[i+1]; p++ {
|
||
sum += values[p] * x[colIdx[p]]
|
||
}
|
||
y[i] = sum
|
||
}
|
||
}
|
||
|
||
// matVec computes y = A·x over complex vectors. Output rows are
|
||
// independent, so the row range splits across workers once a worker's
|
||
// share of the stored entries pays for the split, and runs on the
|
||
// calling goroutine below that.
|
||
func (c *complexCSR) matVec(x, y []complex128) {
|
||
// The three structure slices are taken once: every worker's rows walk
|
||
// them per stored entry, and the row loop holds no other state.
|
||
rowStart := c.rowStart
|
||
colIdx := c.colIdx
|
||
vals := c.vals
|
||
if sparseMatVecSplit(c.n, len(vals)) {
|
||
engine.ParallelMin(c.n, 1, func(s, e int) {
|
||
complexMatVecRange(s, e, rowStart, colIdx, vals, x, y)
|
||
})
|
||
return
|
||
}
|
||
complexMatVecRange(0, c.n, rowStart, colIdx, vals, x, y)
|
||
}
|
||
|
||
// complexVecF64 flattens a rank-1 complex array, rejecting anything
|
||
// else loudly.
|
||
func complexVecF64(b *core.Array, n int, name string) ([]complex128, error) {
|
||
if b.NDim() != 1 || b.Len() != n {
|
||
return nil, base.Errf("%s: b must be a rank-1 vector of length %d", name, n)
|
||
}
|
||
if b.Dtype() != core.Complex {
|
||
return nil, base.Errf("%s: b must be complex128, got %s", name, b.Dtype())
|
||
}
|
||
out := make([]complex128, n)
|
||
if !b.Strided() {
|
||
copy(out, b.RawComplexes()[:n])
|
||
return out, nil
|
||
}
|
||
for i := range n {
|
||
out[i] = b.ComplexAt(i)
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// complexFromArray builds the result array from a complex vector.
|
||
func complexFromArray(v []complex128) (*core.Array, error) {
|
||
out, err := zeros(core.Complex, []int{len(v)})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
copy(out.RawComplexes(), v)
|
||
return out, nil
|
||
}
|
||
|
||
// SpSolveComplexCG returns x solving A·x = b for a Hermitian
|
||
// positive-definite sparse A with complex entries, by conjugate
|
||
// gradient with a complex Jacobi (diagonal) preconditioner. The
|
||
// stopping rule mirrors SpSolve: ‖b − A·x‖₂ ≤ tol·‖b‖₂, tol ≤ 0 uses
|
||
// 1e-10, maxIter ≤ 0 uses n, and an unconverged solve is an error
|
||
// naming the achieved residual rather than a silent approximation. A
|
||
// zero or missing diagonal entry is refused, exactly as the real
|
||
// solver refuses it.
|
||
func SpSolveComplexCG(a *core.SparseCOO, b *core.Array, tol float64, maxIter int) (*core.Array, error) {
|
||
const name = "SpSolveComplexCG"
|
||
c, err := cooToComplexCSR(a, name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := c.checkHermitian(name); err != nil {
|
||
return nil, err
|
||
}
|
||
n := c.n
|
||
bv, err := complexVecF64(b, n, name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if tol <= 0 {
|
||
tol = spSolveTol
|
||
}
|
||
if maxIter <= 0 {
|
||
maxIter = n
|
||
}
|
||
diag := make([]complex128, n)
|
||
for i := range n {
|
||
d, ok := c.at(i, i)
|
||
if !ok || d == 0 {
|
||
return nil, base.Errf("%s: zero diagonal entry at row %d", name, i)
|
||
}
|
||
diag[i] = d
|
||
}
|
||
// The system is solved in scaled units, which leaves the solution x
|
||
// untouched: A·x = b and (f·A)·x = (f·b) have the same x. The factor
|
||
// is 1 for any ordinary magnitude, so the recurrence stays
|
||
// bit-identical there, and outside that range it is what keeps the
|
||
// dot products and the squared magnitudes in the recurrence from
|
||
// overflowing or underflowing to nothing.
|
||
ws := scaleSystem(c, bv, diag)
|
||
bNorm := normC(bv)
|
||
if bNorm == 0 {
|
||
return complexFromArray(make([]complex128, n))
|
||
}
|
||
x := make([]complex128, n)
|
||
r := append([]complex128(nil), bv...)
|
||
p := make([]complex128, n)
|
||
z := make([]complex128, n)
|
||
Ap := make([]complex128, n)
|
||
for i := range n {
|
||
z[i] = r[i] / diag[i]
|
||
}
|
||
copy(p, z)
|
||
rz := dotC(r, z)
|
||
for iter := range maxIter {
|
||
c.matVec(p, Ap)
|
||
pAp := dotC(p, Ap)
|
||
if !cFinite(pAp) || pAp == 0 {
|
||
return nil, base.Errf("%s: breakdown at step %d (p·A·p vanished)", name, iter+1)
|
||
}
|
||
alpha := rz / pAp
|
||
for i := range n {
|
||
x[i] += alpha * p[i]
|
||
r[i] -= alpha * Ap[i]
|
||
}
|
||
if normC(r) <= tol*bNorm {
|
||
return complexFromArray(x)
|
||
}
|
||
for i := range n {
|
||
z[i] = r[i] / diag[i]
|
||
}
|
||
rzNext := dotC(r, z)
|
||
if rzNext == 0 {
|
||
return nil, base.Errf("%s: breakdown at step %d (r·z vanished)", name, iter+1)
|
||
}
|
||
beta := rzNext / rz
|
||
for i := range n {
|
||
p[i] = z[i] + beta*p[i]
|
||
}
|
||
rz = rzNext
|
||
}
|
||
// The residual and tolerance are reported in the caller's units, not
|
||
// in the scaled units the recurrence ran in (ws is 1 unless the
|
||
// system was extreme).
|
||
return nil, base.Errf("%s: no convergence in %d steps, residual %.3g (tolerance %.3g)",
|
||
name, maxIter, normC(r)/ws, tol*bNorm/ws)
|
||
}
|
||
|
||
// dotC is the complex inner product aᴴ·b (conjugating the left side,
|
||
// the convention every complex Krylov method uses).
|
||
func dotC(a, b []complex128) complex128 {
|
||
s := 0i
|
||
for i := range a {
|
||
s += complexConj(a[i]) * b[i]
|
||
}
|
||
return s
|
||
}
|
||
|
||
// normC is the Euclidean norm of a complex vector. The squares are
|
||
// summed directly while the sum stays finite and non-zero, which is what
|
||
// the stopping rules of the solvers were tuned on and is bit-for-bit the
|
||
// magnitude of every ordinary vector. A sum that overflows to +Inf or
|
||
// underflows to 0 (the whole vector below about 1e-162) carries no
|
||
// magnitude at all, so it falls back to a max-scaled accumulation that
|
||
// does: the magnitude has to exist for a relative convergence test to
|
||
// mean anything.
|
||
func normC(a []complex128) float64 {
|
||
s := 0.0
|
||
for _, v := range a {
|
||
s += real(v)*real(v) + imag(v)*imag(v)
|
||
}
|
||
if s == 0 || math.IsInf(s, 1) {
|
||
return normCScaled(a)
|
||
}
|
||
return math.Sqrt(s)
|
||
}
|
||
|
||
// normCScaled sums the squares relative to the largest magnitude, so no
|
||
// intermediate leaves the normal float64 range.
|
||
func normCScaled(a []complex128) float64 {
|
||
maxAbs := 0.0
|
||
for _, v := range a {
|
||
if m := cmplxAbs(v); m > maxAbs {
|
||
maxAbs = m
|
||
}
|
||
}
|
||
if maxAbs == 0 {
|
||
return 0
|
||
}
|
||
s := 0.0
|
||
for _, v := range a {
|
||
re := real(v) / maxAbs
|
||
im := imag(v) / maxAbs
|
||
s += re*re + im*im
|
||
}
|
||
return maxAbs * math.Sqrt(s)
|
||
}
|
||
|
||
// scaleSystem moves a system's matrix, right-hand side and preconditioner
|
||
// into the magnitude window the Krylov recurrences are safe in, by one
|
||
// shared exact power of two, and reports that factor. A·x = b and
|
||
// (f·A)·x = (f·b) have the same solution x, so the factor cancels: it
|
||
// only has to come off the residual and tolerance the error paths print
|
||
// back. Inside the window the dot products, the squared magnitudes and
|
||
// the stopping rule's tol*bNorm all mean what they claim; outside it a
|
||
// huge b makes bNorm +Inf (every affine check passes) and a tiny one
|
||
// makes it 0 (the zero solution is returned as exact).
|
||
func scaleSystem(c *complexCSR, bv, diag []complex128) float64 {
|
||
ws := windowScale(math.Max(maxMagComplex(c.vals), maxMagComplex(bv)))
|
||
if ws != 1 {
|
||
scaleComplexes(c.vals, ws)
|
||
scaleComplexes(bv, ws)
|
||
scaleComplexes(diag, ws)
|
||
}
|
||
return ws
|
||
}
|
||
|
||
// SpSolveComplexBiCGSTAB returns x solving A·x = b for a general
|
||
// (nonsymmetric, non-Hermitian) sparse A with complex entries, by
|
||
// BiCGSTAB with a complex Jacobi preconditioner. The contract mirrors
|
||
// SpSolveBiCGSTAB, including the affirmative convergence checks a NaN
|
||
// residual can never pass and the breakdown errors that name the step.
|
||
func SpSolveComplexBiCGSTAB(a *core.SparseCOO, b *core.Array, tol float64, maxIter int) (*core.Array, error) {
|
||
const name = "SpSolveComplexBiCGSTAB"
|
||
c, err := cooToComplexCSR(a, name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
n := c.n
|
||
bv, err := complexVecF64(b, n, name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if tol <= 0 {
|
||
tol = spSolveTol
|
||
}
|
||
if maxIter <= 0 {
|
||
maxIter = n
|
||
}
|
||
diag := make([]complex128, n)
|
||
for i := range n {
|
||
d, ok := c.at(i, i)
|
||
if !ok || d == 0 {
|
||
return nil, base.Errf("%s: zero diagonal entry at row %d", name, i)
|
||
}
|
||
diag[i] = d
|
||
}
|
||
// The system is solved in scaled units, which leaves the solution x
|
||
// untouched, exactly as in SpSolveComplexCG above.
|
||
ws := scaleSystem(c, bv, diag)
|
||
bNorm := normC(bv)
|
||
if bNorm == 0 {
|
||
return complexFromArray(make([]complex128, n))
|
||
}
|
||
r := append([]complex128(nil), bv...)
|
||
rHat := append([]complex128(nil), r...)
|
||
x := make([]complex128, n)
|
||
v := make([]complex128, n)
|
||
p := make([]complex128, n)
|
||
s := make([]complex128, n)
|
||
t := make([]complex128, n)
|
||
precond := func(dst, src []complex128) {
|
||
for i := range n {
|
||
dst[i] = src[i] / diag[i]
|
||
}
|
||
}
|
||
rho := 1 + 0i
|
||
alpha := 1 + 0i
|
||
omega := 1 + 0i
|
||
z := make([]complex128, n)
|
||
y := make([]complex128, n)
|
||
mt := make([]complex128, n)
|
||
for iter := range maxIter {
|
||
rhoNext := dotC(rHat, r)
|
||
if !cFinite(rhoNext) || rhoNext == 0 {
|
||
return nil, base.Errf("%s: breakdown at step %d (rho vanished)", name, iter+1)
|
||
}
|
||
if iter > 0 {
|
||
beta := rhoNext / rho * (alpha / omega)
|
||
for i := range n {
|
||
p[i] = r[i] + beta*(p[i]-omega*v[i])
|
||
}
|
||
} else {
|
||
copy(p, r)
|
||
}
|
||
precond(z, p)
|
||
c.matVec(z, v)
|
||
rhatV := dotC(rHat, v)
|
||
if rhatV == 0 {
|
||
return nil, base.Errf("%s: breakdown at step %d (direction orthogonal to shadow)", name, iter+1)
|
||
}
|
||
alpha = rhoNext / rhatV
|
||
for i := range n {
|
||
s[i] = r[i] - alpha*v[i]
|
||
}
|
||
if !vecFiniteC(s) {
|
||
return nil, base.Errf("%s: breakdown at step %d (non-finite residual)", name, iter+1)
|
||
}
|
||
// The half step may already be exact. The check is affirmative
|
||
// so a NaN norm (which fails every comparison) can never read
|
||
// as converged.
|
||
if normC(s) > tol*bNorm {
|
||
precond(y, s)
|
||
c.matVec(y, t)
|
||
precond(mt, t)
|
||
mtS := dotC(mt, s)
|
||
mtT := dotC(mt, t)
|
||
if mtT == 0 {
|
||
return nil, base.Errf("%s: breakdown at step %d (stabiliser vanished)", name, iter+1)
|
||
}
|
||
omega = mtS / mtT
|
||
if !cFinite(omega) {
|
||
return nil, base.Errf("%s: breakdown at step %d (stabiliser vanished)", name, iter+1)
|
||
}
|
||
for i := range n {
|
||
x[i] += alpha*z[i] + omega*y[i]
|
||
r[i] = s[i] - omega*t[i]
|
||
}
|
||
if !vecFiniteC(r) {
|
||
return nil, base.Errf("%s: breakdown at step %d (non-finite residual)", name, iter+1)
|
||
}
|
||
if normC(r) <= tol*bNorm {
|
||
return complexFromArray(x)
|
||
}
|
||
if omega == 0 {
|
||
return nil, base.Errf("%s: stagnation at step %d (omega zero)", name, iter+1)
|
||
}
|
||
} else {
|
||
for i := range n {
|
||
x[i] += alpha * z[i]
|
||
}
|
||
return complexFromArray(x)
|
||
}
|
||
rho = rhoNext
|
||
}
|
||
// The residual and tolerance are reported in the caller's units, not
|
||
// in the scaled units the recurrence ran in (ws is 1 unless the
|
||
// system was extreme).
|
||
return nil, base.Errf("%s: no convergence in %d steps, residual %.3g (tolerance %.3g)",
|
||
name, maxIter, normC(r)/ws, tol*bNorm/ws)
|
||
}
|
||
|
||
func cFinite(z complex128) bool {
|
||
return !math.IsNaN(real(z)) && !math.IsInf(real(z), 0) &&
|
||
!math.IsNaN(imag(z)) && !math.IsInf(imag(z), 0)
|
||
}
|
||
|
||
func vecFiniteC(v []complex128) bool {
|
||
for _, z := range v {
|
||
if !cFinite(z) {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// SpEigenComplex returns the k eigenvalues of largest magnitude of a
|
||
// Hermitian sparse matrix with complex entries, each with its unit
|
||
// eigenvector, by the Lanczos recurrence over complex vectors. The
|
||
// contract mirrors SpEigen: values ordered by descending magnitude,
|
||
// eigenvectors as columns of an (n, k) complex array, approximate
|
||
// Ritz pairs whose accuracy improves with the iteration budget. A
|
||
// non-Hermitian matrix is refused; the general complex eigenproblem
|
||
// has no short recurrence and belongs to dense methods.
|
||
func SpEigenComplex(s *core.SparseCOO, k int, gen *core.Generator) (values, vectors *core.Array, err error) {
|
||
const name = "SpEigenComplex"
|
||
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)
|
||
}
|
||
n := s.Shape[0]
|
||
if n == 0 {
|
||
return nil, nil, base.Errf("%s: zero-sized matrix, got shape %v", name, s.Shape)
|
||
}
|
||
if k < 1 || k > n {
|
||
return nil, nil, base.Errf("%s: k must be in [1, %d], got %d", name, n, k)
|
||
}
|
||
c, err := cooToComplexCSR(s, name)
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
if err := c.checkHermitian(name); err != nil {
|
||
return nil, nil, err
|
||
}
|
||
// The recurrence is projected in scaled units: the alphas and betas
|
||
// carry the matrix's magnitude, the basis vectors do not, and a
|
||
// power-of-two factor multiplies the first two exactly while leaving
|
||
// the third bit-identical, so only the Ritz values have to be moved
|
||
// back. Without it the projected tridiagonal's squared magnitudes
|
||
// overflow at the top of the range and the recurrence collapses to
|
||
// NaN, and at the bottom the deflation floor sees zeros.
|
||
ws := windowScale(maxMagComplex(c.vals))
|
||
if ws != 1 {
|
||
scaleComplexes(c.vals, ws)
|
||
}
|
||
if gen == nil {
|
||
gen = core.NewGenerator(spEigenSeed)
|
||
}
|
||
alphas, betas, basis := c.lanczosComplex(k, gen)
|
||
r := len(alphas)
|
||
|
||
tMat := make([]float64, r*r)
|
||
for i := range r {
|
||
tMat[i*r+i] = alphas[i]
|
||
if i+1 < r {
|
||
tMat[i*r+i+1] = betas[i]
|
||
tMat[(i+1)*r+i] = betas[i]
|
||
}
|
||
}
|
||
tVec := eye(r)
|
||
if err := symmetricQr(tMat, tVec, r); err != nil {
|
||
return nil, nil, base.Errf("SpEigenComplex: %w", err)
|
||
}
|
||
|
||
idx := make([]int, r)
|
||
for i := range r {
|
||
idx[i] = i
|
||
}
|
||
// Rank by descending magnitude, the same comparator SpEigen uses.
|
||
for i := 1; i < r; i++ {
|
||
for j := i; j > 0; j-- {
|
||
a, b := idx[j-1], idx[j]
|
||
da, db := math.Abs(tMat[a*r+a]), math.Abs(tMat[b*r+b])
|
||
if da > db || (da == db && a <= b) {
|
||
break
|
||
}
|
||
idx[j-1], idx[j] = idx[j], idx[j-1]
|
||
}
|
||
}
|
||
sel := idx[:k]
|
||
|
||
outVals := make([]float64, k)
|
||
outVecs := make([]complex128, n*k)
|
||
tmp := make([]complex128, n)
|
||
for j, si := range sel {
|
||
outVals[j] = tMat[si*r+si]
|
||
// Ritz vector: lift the real tridiagonal eigenvector through
|
||
// the complex Lanczos basis, v = Q·y.
|
||
clear(tmp)
|
||
for p := range r {
|
||
y := tVec[p*r+si]
|
||
if y == 0 {
|
||
continue
|
||
}
|
||
row := basis[p*n : (p+1)*n]
|
||
for i := range n {
|
||
tmp[i] += complex(y, 0) * row[i]
|
||
}
|
||
}
|
||
norm := normC(tmp)
|
||
if norm > 0 {
|
||
for i := range n {
|
||
tmp[i] /= complex(norm, 0)
|
||
}
|
||
}
|
||
for i := range n {
|
||
outVecs[i*k+j] = tmp[i]
|
||
}
|
||
}
|
||
if ws != 1 {
|
||
unscaleFloats(outVals, ws)
|
||
}
|
||
vecs, err := complexFromArray2D(outVecs, n, k)
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
return floatsToArray(outVals, []int{k}), vecs, nil
|
||
}
|
||
|
||
// complexFromArray2D builds an (n, k) complex array from row-major
|
||
// values.
|
||
func complexFromArray2D(v []complex128, n, k int) (*core.Array, error) {
|
||
out, err := zeros(core.Complex, []int{n, k})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
copy(out.RawComplexes(), v)
|
||
return out, nil
|
||
}
|
||
|
||
// lanczosComplex runs the Hermitian Lanczos recurrence over complex
|
||
// vectors: alphas stay real (qᴴAq of a Hermitian A), betas stay real
|
||
// (vector norms), so the projected problem reuses the real symmetric
|
||
// tridiagonal eigensolver unchanged. Full reorthogonalisation and the
|
||
// purely scale-relative deflation floor mirror the real lanczos.
|
||
func (c *complexCSR) lanczosComplex(k int, gen *core.Generator) (alphas, betas []float64, basis []complex128) {
|
||
n := c.n
|
||
steps := min(n, max(2*k, k+spEigenBlock))
|
||
w := make([]complex128, n)
|
||
q := make([]complex128, n)
|
||
prev := make([]complex128, n)
|
||
// The budget bounds the recurrence: one basis row and at most one
|
||
// coefficient per step, so all three are sized once rather than grown.
|
||
alphas = make([]float64, 0, steps)
|
||
betas = make([]float64, 0, steps)
|
||
basis = make([]complex128, 0, steps*n)
|
||
scale := 0.0
|
||
addScale := func(v float64) {
|
||
if v > scale {
|
||
scale = v
|
||
}
|
||
}
|
||
start := func() {
|
||
for i := range n {
|
||
q[i] = complex(gen.NormalUnit(), gen.NormalUnit())
|
||
}
|
||
for p := range len(alphas) {
|
||
row := basis[p*n : (p+1)*n]
|
||
d := dotC(row, q)
|
||
for i := range n {
|
||
q[i] -= d * row[i]
|
||
}
|
||
}
|
||
if norm := normC(q); norm > 0 {
|
||
for i := range n {
|
||
q[i] /= complex(norm, 0)
|
||
}
|
||
}
|
||
}
|
||
start()
|
||
for range steps {
|
||
basis = append(basis, q...)
|
||
c.matVec(q, w)
|
||
alpha := real(dotC(q, w))
|
||
alphas = append(alphas, alpha)
|
||
addScale(math.Abs(alpha))
|
||
a := complex(alpha, 0)
|
||
for i := range n {
|
||
w[i] -= a * q[i]
|
||
}
|
||
if len(alphas) > 1 {
|
||
b := complex(betas[len(betas)-1], 0)
|
||
for i := range n {
|
||
w[i] -= b * prev[i]
|
||
}
|
||
}
|
||
// Full reorthogonalisation, twice: one pass removes the
|
||
// accumulated loss, the second what the first reintroduces.
|
||
for range 2 {
|
||
for p := range len(alphas) {
|
||
row := basis[p*n : (p+1)*n]
|
||
d := dotC(row, w)
|
||
for i := range n {
|
||
w[i] -= d * row[i]
|
||
}
|
||
}
|
||
}
|
||
beta := normC(w)
|
||
// A purely scale-relative exhaustion threshold; an absolute
|
||
// floor would collapse the recurrence for tiny-norm matrices.
|
||
if beta <= float64(n)*base.EpsF*scale {
|
||
if len(alphas) >= n {
|
||
break
|
||
}
|
||
betas = append(betas, 0)
|
||
start()
|
||
continue
|
||
}
|
||
betas = append(betas, beta)
|
||
addScale(beta)
|
||
copy(prev, q)
|
||
for i := range n {
|
||
q[i] = w[i] / complex(beta, 0)
|
||
}
|
||
}
|
||
if len(betas) >= len(alphas) {
|
||
betas = betas[:len(alphas)-1]
|
||
}
|
||
return alphas, betas, basis
|
||
}
|