252 lines
7.3 KiB
Go
252 lines
7.3 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/core"
|
|
"testing"
|
|
)
|
|
|
|
// denseExpApply forms the whole exponential densely and multiplies it
|
|
// by v, the reference the Krylov projection is compared against.
|
|
func denseExpApply(t *testing.T, vals, v []float64, n int) []float64 {
|
|
t.Helper()
|
|
a, err := core.FromFloats(vals, n, n)
|
|
if err != nil {
|
|
t.Fatalf("FromFloats: %v", err)
|
|
}
|
|
ea, err := MatrixExp(a)
|
|
if err != nil {
|
|
t.Fatalf("MatrixExp: %v", err)
|
|
}
|
|
out := make([]float64, n)
|
|
for i := range n {
|
|
s := 0.0
|
|
for j := range n {
|
|
s += ea.FloatAt(i*n+j) * v[j]
|
|
}
|
|
out[i] = s
|
|
}
|
|
return out
|
|
}
|
|
|
|
// TestSpExpApplyMatchesDense compares the Krylov projection against
|
|
// the dense exponential on symmetric matrices, at both a full and a
|
|
// truncated Krylov dimension.
|
|
func TestSpExpApplyMatchesDense(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
n int
|
|
diag float64
|
|
off float64
|
|
steps int
|
|
}{
|
|
// n ≤ 40 gets the whole Krylov space, so the answer is exact.
|
|
{name: "exact_small", n: 6, diag: 3, off: -1, steps: 0},
|
|
{name: "exact_at_budget", n: 40, diag: 4, off: -1, steps: 0},
|
|
// A truncated dimension on a larger matrix is approximate.
|
|
{name: "truncated", n: 100, diag: 5, off: -1, steps: 40},
|
|
// A near-diagonal matrix, where the action is close to
|
|
// elementwise and the projection converges immediately.
|
|
{name: "weakly_coupled", n: 30, diag: 2, off: -0.01, steps: 0},
|
|
}
|
|
for _, tt := range cases {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
vals := spdTridiagonal(tt.n, tt.diag, tt.off)
|
|
sp := sparseFromDense(t, vals, tt.n)
|
|
v := make([]float64, tt.n)
|
|
for i := range tt.n {
|
|
v[i] = math.Sin(float64(i+1)) * 0.5
|
|
}
|
|
vArr, err := core.FromFloats(v, tt.n)
|
|
if err != nil {
|
|
t.Fatalf("FromFloats v: %v", err)
|
|
}
|
|
got, err := SpExpApply(sp, vArr, tt.steps)
|
|
if err != nil {
|
|
t.Fatalf("SpExpApply: %v", err)
|
|
}
|
|
want := denseExpApply(t, vals, v, tt.n)
|
|
for i := range tt.n {
|
|
g, w := got.FloatAt(i), want[i]
|
|
if math.Abs(g-w) > 1e-8*(1+math.Abs(w)) {
|
|
t.Fatalf("element %d = %.12g, want %.12g", i, g, w)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestSpExpApplyDiagonal pins the closed form: for a diagonal matrix,
|
|
// exp(A)·v is the elementwise exponential of the diagonal times v.
|
|
func TestSpExpApplyDiagonal(t *testing.T) {
|
|
const n = 5
|
|
diag := []float64{1, 2, -1, 0.5, 3}
|
|
vals := make([]float64, n*n)
|
|
for i := range n {
|
|
vals[i*n+i] = diag[i]
|
|
}
|
|
sp := sparseFromDense(t, vals, n)
|
|
v := []float64{1, 1, 1, 1, 1}
|
|
vArr, err := core.FromFloats(v, n)
|
|
if err != nil {
|
|
t.Fatalf("FromFloats: %v", err)
|
|
}
|
|
got, err := SpExpApply(sp, vArr, 0)
|
|
if err != nil {
|
|
t.Fatalf("SpExpApply: %v", err)
|
|
}
|
|
for i := range n {
|
|
want := math.Exp(diag[i])
|
|
if math.Abs(got.FloatAt(i)-want) > 1e-12*(1+math.Abs(want)) {
|
|
t.Fatalf("element %d = %.12g, want %.12g", i, got.FloatAt(i), want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestSpExpApplyZeroVector pins the degenerate contract: exp(A)·0 is
|
|
// the zero vector, returned without dividing by a zero norm.
|
|
func TestSpExpApplyZeroVector(t *testing.T) {
|
|
const n = 4
|
|
vals := spdTridiagonal(n, 3, -1)
|
|
sp := sparseFromDense(t, vals, n)
|
|
v, err := core.Zeros(core.Float, n)
|
|
if err != nil {
|
|
t.Fatalf("Zeros: %v", err)
|
|
}
|
|
got, err := SpExpApply(sp, v, 0)
|
|
if err != nil {
|
|
t.Fatalf("SpExpApply: %v", err)
|
|
}
|
|
for i := range n {
|
|
if got.FloatAt(i) != 0 {
|
|
t.Fatalf("element %d = %.12g, want 0 for a zero vector", i, got.FloatAt(i))
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestSpExpApplyDeterminism checks reproducibility: the projection
|
|
// starts from v itself and draws nothing random.
|
|
func TestSpExpApplyDeterminism(t *testing.T) {
|
|
const n = 50
|
|
vals := spdTridiagonal(n, 4, -1)
|
|
sp := sparseFromDense(t, vals, n)
|
|
v := make([]float64, n)
|
|
for i := range n {
|
|
v[i] = float64(i+1) / float64(n)
|
|
}
|
|
vArr, err := core.FromFloats(v, n)
|
|
if err != nil {
|
|
t.Fatalf("FromFloats: %v", err)
|
|
}
|
|
g1, err := SpExpApply(sp, vArr, 0)
|
|
if err != nil {
|
|
t.Fatalf("SpExpApply #1: %v", err)
|
|
}
|
|
g2, err := SpExpApply(sp, vArr, 0)
|
|
if err != nil {
|
|
t.Fatalf("SpExpApply #2: %v", err)
|
|
}
|
|
for i := range n {
|
|
if g1.FloatAt(i) != g2.FloatAt(i) {
|
|
t.Fatalf("element %d = %.12g vs %.12g across runs", i, g1.FloatAt(i), g2.FloatAt(i))
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestSpExpApplyStepsClamped checks that a step count past the
|
|
// dimension is clamped to it, rather than overrunning the space.
|
|
func TestSpExpApplyStepsClamped(t *testing.T) {
|
|
const n = 4
|
|
vals := spdTridiagonal(n, 3, -1)
|
|
sp := sparseFromDense(t, vals, n)
|
|
v := []float64{1, 2, 3, 4}
|
|
vArr, err := core.FromFloats(v, n)
|
|
if err != nil {
|
|
t.Fatalf("FromFloats: %v", err)
|
|
}
|
|
got, err := SpExpApply(sp, vArr, 1000)
|
|
if err != nil {
|
|
t.Fatalf("SpExpApply: %v", err)
|
|
}
|
|
want := denseExpApply(t, vals, v, n)
|
|
for i := range n {
|
|
if math.Abs(got.FloatAt(i)-want[i]) > 1e-10*(1+math.Abs(want[i])) {
|
|
t.Fatalf("element %d = %.12g, want %.12g", i, got.FloatAt(i), want[i])
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestSpExpApplyRejectsInvalid pins the error contract for every input
|
|
// the projection cannot honestly answer.
|
|
func TestSpExpApplyRejectsInvalid(t *testing.T) {
|
|
vector := func(n int) *core.Array {
|
|
v, err := core.FromFloats(make([]float64, n), n)
|
|
if err != nil {
|
|
t.Fatalf("FromFloats: %v", err)
|
|
}
|
|
return v
|
|
}
|
|
t.Run("complex_sparse", func(t *testing.T) {
|
|
idx, _ := core.FromInts([]int64{0, 0}, 1, 2)
|
|
vals, _ := core.FromComplexes([]complex128{1}, 1)
|
|
sp, err := core.NewSparseCOO(idx, vals, []int{1, 1})
|
|
if err != nil {
|
|
t.Fatalf("NewSparseCOO: %v", err)
|
|
}
|
|
if _, err := SpExpApply(sp, vector(1), 0); err == nil {
|
|
t.Fatal("expected an error for a complex sparse matrix")
|
|
}
|
|
})
|
|
t.Run("not_square", func(t *testing.T) {
|
|
idx, _ := core.FromInts([]int64{0, 0}, 1, 2)
|
|
vals, _ := core.FromFloats([]float64{1}, 1)
|
|
sp, err := core.NewSparseCOO(idx, vals, []int{1, 2})
|
|
if err != nil {
|
|
t.Fatalf("NewSparseCOO: %v", err)
|
|
}
|
|
if _, err := SpExpApply(sp, vector(1), 0); err == nil {
|
|
t.Fatal("expected an error for a non-square matrix")
|
|
}
|
|
})
|
|
t.Run("zero_sized", func(t *testing.T) {
|
|
idx, _ := core.FromInts(nil, 0, 2)
|
|
vals, _ := core.FromFloats(nil, 0)
|
|
sp := &core.SparseCOO{Indices: idx, Values: vals, Shape: []int{0, 0}}
|
|
if _, err := SpExpApply(sp, vector(0), 0); err == nil {
|
|
t.Fatal("expected an error for a zero-sized matrix")
|
|
}
|
|
})
|
|
t.Run("vector_wrong_length", func(t *testing.T) {
|
|
vals := spdTridiagonal(3, 3, -1)
|
|
sp := sparseFromDense(t, vals, 3)
|
|
if _, err := SpExpApply(sp, vector(2), 0); err == nil {
|
|
t.Fatal("expected an error for a vector of the wrong length")
|
|
}
|
|
})
|
|
t.Run("vector_not_rank1", func(t *testing.T) {
|
|
vals := spdTridiagonal(2, 3, -1)
|
|
sp := sparseFromDense(t, vals, 2)
|
|
m, _ := core.FromFloats([]float64{1, 0, 0, 1}, 2, 2)
|
|
if _, err := SpExpApply(sp, m, 0); err == nil {
|
|
t.Fatal("expected an error for a rank-2 vector")
|
|
}
|
|
})
|
|
t.Run("complex_vector", func(t *testing.T) {
|
|
vals := spdTridiagonal(2, 3, -1)
|
|
sp := sparseFromDense(t, vals, 2)
|
|
v, _ := core.FromComplexes([]complex128{1, 1}, 2)
|
|
if _, err := SpExpApply(sp, v, 0); err == nil {
|
|
t.Fatal("expected an error for a complex vector")
|
|
}
|
|
})
|
|
t.Run("asymmetric", func(t *testing.T) {
|
|
sp := sparseFromDense(t, []float64{4, 1, 2, 4}, 2)
|
|
if _, err := SpExpApply(sp, vector(2), 0); err == nil {
|
|
t.Fatal("expected an error for an asymmetric matrix")
|
|
}
|
|
})
|
|
}
|