329 lines
9.2 KiB
Go
329 lines
9.2 KiB
Go
|
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|||
|
|
// SPDX-License-Identifier: MIT
|
|||
|
|
|
|||
|
|
package linalg
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
|||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
import "math"
|
|||
|
|
|
|||
|
|
// Matrix exponential, evaluated by scaling and squaring over a
|
|||
|
|
// diagonal Pade approximant (Higham, 2005). That is the construction
|
|||
|
|
// LAPACK's *ex* routines and every mainstream runtime use, because it
|
|||
|
|
// is backward stable and needs nothing but matrix products and one LU
|
|||
|
|
// solve, both of which the library already has.
|
|||
|
|
//
|
|||
|
|
// The kernel is generic over the same scalar family the LU kernel runs
|
|||
|
|
// on, so the complex path stays in complex128 end to end instead of
|
|||
|
|
// splitting off into a parallel implementation.
|
|||
|
|
|
|||
|
|
// expmTheta holds the 1-norm bound below which each Pade degree keeps
|
|||
|
|
// its approximation error inside double precision (Higham 2005,
|
|||
|
|
// Table 10.1).
|
|||
|
|
var expmTheta = map[int]float64{
|
|||
|
|
3: 1.495585217958292e-2,
|
|||
|
|
5: 2.539398330063230e-1,
|
|||
|
|
7: 9.504178996162932e-1,
|
|||
|
|
9: 2.097847961257068e+0,
|
|||
|
|
13: 5.371920351148152e+0,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// expmCoeffs returns the numerator coefficients b_j of the degree-m
|
|||
|
|
// diagonal Pade approximant, as the integers of Higham's Table 10.1.
|
|||
|
|
// The common factor across the row cancels in (V-U)^{-1}(V+U), so the
|
|||
|
|
// unnormalised values are used directly, exactly as the reference
|
|||
|
|
// implementations do.
|
|||
|
|
func expmCoeffs(deg int) []float64 {
|
|||
|
|
switch deg {
|
|||
|
|
case 3:
|
|||
|
|
return []float64{120, 60, 12, 1}
|
|||
|
|
case 5:
|
|||
|
|
return []float64{30240, 15120, 3360, 420, 30, 1}
|
|||
|
|
case 7:
|
|||
|
|
return []float64{17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1}
|
|||
|
|
case 9:
|
|||
|
|
return []float64{17643225600, 8821612800, 2075673600, 302702400,
|
|||
|
|
30270240, 2162160, 110880, 3960, 90, 1}
|
|||
|
|
default:
|
|||
|
|
return []float64{64764752532480000, 32382376266240000, 7771770303897600,
|
|||
|
|
1187353796428800, 129060195264000, 10559470521600, 670442572800,
|
|||
|
|
33522128640, 1323241920, 40840800, 960960, 16380, 182, 1}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// padeDegree picks the approximant degree and the number of squarings
|
|||
|
|
// for a given 1-norm: the smallest degree whose bound covers the norm,
|
|||
|
|
// else degree 13 with the matrix scaled down by a power of two.
|
|||
|
|
func padeDegree(nA float64) (deg, s int) {
|
|||
|
|
for _, d := range []int{3, 5, 7, 9, 13} {
|
|||
|
|
if nA <= expmTheta[d] {
|
|||
|
|
return d, 0
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
s = max(int(math.Ceil(math.Log2(nA/expmTheta[13]))), 0)
|
|||
|
|
return 13, s
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// MatrixExp returns the matrix exponential exp(A) of a square matrix.
|
|||
|
|
// Ints and float32 promote to float64; a complex matrix answers in
|
|||
|
|
// complex128. A zero matrix gives the identity, and exp(A) for
|
|||
|
|
// diagonal A is the exponential of the diagonal.
|
|||
|
|
func MatrixExp(a *core.Array) (*core.Array, error) {
|
|||
|
|
if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
|
|||
|
|
return nil, base.Errf("MatrixExp: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
|
|||
|
|
}
|
|||
|
|
n := a.Shape()[0]
|
|||
|
|
if n == 0 {
|
|||
|
|
return nil, base.Errf("MatrixExp: zero-sized matrix, got shape %s", base.ShapeText(a.Shape()))
|
|||
|
|
}
|
|||
|
|
if a.Dtype() == core.Complex {
|
|||
|
|
m, err := squareComplexMatrix(a, "MatrixExp")
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
r, err := expmKernel(m)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
flat := make([]complex128, n*n)
|
|||
|
|
for i := range n {
|
|||
|
|
copy(flat[i*n:(i+1)*n], r[i])
|
|||
|
|
}
|
|||
|
|
return core.FromComplexes(flat, n, n)
|
|||
|
|
}
|
|||
|
|
m, err := squareFloatMatrix(a, "MatrixExp")
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
r, err := expmKernel(m)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
flat := make([]float64, n*n)
|
|||
|
|
for i := range n {
|
|||
|
|
copy(flat[i*n:(i+1)*n], r[i])
|
|||
|
|
}
|
|||
|
|
return floatsToArray(flat, []int{n, n}), nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// expmKernel computes exp(a) by scaling and squaring over a Pade
|
|||
|
|
// approximant. The approximant is r(A) = (V-U)^{-1}(V+U) where U holds
|
|||
|
|
// the odd part and V the even part of the numerator polynomial, both
|
|||
|
|
// evaluated by Horner in A^2; s squarings then undo the scaling.
|
|||
|
|
func expmKernel[T scalar](a [][]T) ([][]T, error) {
|
|||
|
|
n := len(a)
|
|||
|
|
nA := matrixNorm1(a)
|
|||
|
|
if math.IsNaN(nA) || math.IsInf(nA, 0) {
|
|||
|
|
// The theta ladder divides by the norm and int(Ceil(...)) of a
|
|||
|
|
// non-finite value is implementation-defined; the Padé solve on
|
|||
|
|
// Inf data would return an all-NaN matrix with a nil error.
|
|||
|
|
return nil, base.Errf("MatrixExp: the matrix holds a non-finite entry (norm %g)", nA)
|
|||
|
|
}
|
|||
|
|
deg, s := padeDegree(nA)
|
|||
|
|
scaled := a
|
|||
|
|
if s > 0 {
|
|||
|
|
scaled = matScaleT(a, math.Ldexp(1, -s))
|
|||
|
|
}
|
|||
|
|
b := expmCoeffs(deg)
|
|||
|
|
a2 := matMulT(scaled, scaled)
|
|||
|
|
// The Horner sweeps alternate between two buffers instead of landing
|
|||
|
|
// a fresh matrix per coefficient: the product reads both operands
|
|||
|
|
// and writes the buffer the previous step left free, so the entries
|
|||
|
|
// the next multiply reads are the ones this one wrote.
|
|||
|
|
v, vFree := eyeScaled[T](n, b[deg-1]), newSquare[T](n)
|
|||
|
|
uIn, uFree := eyeScaled[T](n, b[deg]), newSquare[T](n)
|
|||
|
|
for j := deg - 3; j >= 0; j -= 2 {
|
|||
|
|
matMulAddEyeInto(vFree, a2, v, b[j])
|
|||
|
|
v, vFree = vFree, v
|
|||
|
|
}
|
|||
|
|
for j := deg - 2; j >= 1; j -= 2 {
|
|||
|
|
matMulAddEyeInto(uFree, a2, uIn, b[j])
|
|||
|
|
uIn, uFree = uFree, uIn
|
|||
|
|
}
|
|||
|
|
u := matMulT(scaled, uIn)
|
|||
|
|
r, err := solveRight("MatrixExp", matSubT(v, u), matAddT(v, u))
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
if s > 0 {
|
|||
|
|
sq := newSquare[T](n)
|
|||
|
|
for range s {
|
|||
|
|
matMulTInto(sq, r, r)
|
|||
|
|
r, sq = sq, r
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return r, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// newSquare allocates an n×n matrix in the kernel's element type as n
|
|||
|
|
// row views over one flat backing slice: two allocations instead of
|
|||
|
|
// n+1. The rows stay disjoint, so in-place row swaps remain valid.
|
|||
|
|
func newSquare[T scalar](n int) [][]T {
|
|||
|
|
back := make([]T, n*n)
|
|||
|
|
rows := make([][]T, n)
|
|||
|
|
for i := range n {
|
|||
|
|
rows[i] = back[i*n : (i+1)*n]
|
|||
|
|
}
|
|||
|
|
return rows
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// realT widens a float64 into the kernel's element type, giving a zero
|
|||
|
|
// imaginary part when the kernel runs on complex128.
|
|||
|
|
func realT[T scalar](f float64) T {
|
|||
|
|
var zero T
|
|||
|
|
if _, ok := any(zero).(complex128); ok {
|
|||
|
|
return any(complex(f, 0)).(T)
|
|||
|
|
}
|
|||
|
|
return any(f).(T)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// eyeScaled returns c times the n×n identity.
|
|||
|
|
func eyeScaled[T scalar](n int, c float64) [][]T {
|
|||
|
|
out := newSquare[T](n)
|
|||
|
|
cc := realT[T](c)
|
|||
|
|
for i := range n {
|
|||
|
|
out[i][i] = cc
|
|||
|
|
}
|
|||
|
|
return out
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matMulAddEyeInto writes a·m + c·I into out, which must be a square
|
|||
|
|
// matrix of the same order aliasing neither operand.
|
|||
|
|
func matMulAddEyeInto[T scalar](out, a, m [][]T, c float64) {
|
|||
|
|
matMulTInto(out, a, m)
|
|||
|
|
cc := realT[T](c)
|
|||
|
|
for i := range out {
|
|||
|
|
out[i][i] += cc
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matMulT multiplies two square matrices in the kernel's element type
|
|||
|
|
// and returns a fresh matrix.
|
|||
|
|
func matMulT[T scalar](a, b [][]T) [][]T {
|
|||
|
|
out := newSquare[T](len(a))
|
|||
|
|
matMulTInto(out, a, b)
|
|||
|
|
return out
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matMulTInto writes the product a·b into out, which must alias neither
|
|||
|
|
// operand. out arrives holding an earlier product, so it is cleared
|
|||
|
|
// first: a fresh matrix is zero by construction and the accumulation
|
|||
|
|
// below depends on that. The loop order keeps the inner product over p
|
|||
|
|
// contiguous in both operands, matching the dense kernel's cache
|
|||
|
|
// behaviour, and the zero multiplier is skipped the way the product has
|
|||
|
|
// always skipped it, so a zero times a non-finite entry contributes
|
|||
|
|
// nothing rather than a NaN.
|
|||
|
|
func matMulTInto[T scalar](out, a, b [][]T) {
|
|||
|
|
n := len(a)
|
|||
|
|
var zero T
|
|||
|
|
for i := range n {
|
|||
|
|
oi := out[i]
|
|||
|
|
clear(oi)
|
|||
|
|
ai := a[i]
|
|||
|
|
for p := range n {
|
|||
|
|
aip := ai[p]
|
|||
|
|
if aip == zero {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
bp := b[p]
|
|||
|
|
for j := range n {
|
|||
|
|
oi[j] += aip * bp[j]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matAddT returns the element-wise sum of two square matrices.
|
|||
|
|
func matAddT[T scalar](a, b [][]T) [][]T {
|
|||
|
|
n := len(a)
|
|||
|
|
out := newSquare[T](n)
|
|||
|
|
for i := range n {
|
|||
|
|
oi, ai, bi := out[i], a[i], b[i]
|
|||
|
|
for j := range n {
|
|||
|
|
oi[j] = ai[j] + bi[j]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return out
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matSubT returns the element-wise difference of two square matrices.
|
|||
|
|
func matSubT[T scalar](a, b [][]T) [][]T {
|
|||
|
|
n := len(a)
|
|||
|
|
out := newSquare[T](n)
|
|||
|
|
for i := range n {
|
|||
|
|
oi, ai, bi := out[i], a[i], b[i]
|
|||
|
|
for j := range n {
|
|||
|
|
oi[j] = ai[j] - bi[j]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return out
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matScaleT multiplies every element by the real scalar f.
|
|||
|
|
func matScaleT[T scalar](a [][]T, f float64) [][]T {
|
|||
|
|
n := len(a)
|
|||
|
|
fT := realT[T](f)
|
|||
|
|
out := newSquare[T](n)
|
|||
|
|
for i := range n {
|
|||
|
|
oi, ai := out[i], a[i]
|
|||
|
|
for j := range n {
|
|||
|
|
oi[j] = ai[j] * fT
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return out
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// matrixNorm1 returns the maximum absolute column sum, the norm the
|
|||
|
|
// scaling step is measured against.
|
|||
|
|
func matrixNorm1[T scalar](m [][]T) float64 {
|
|||
|
|
n := len(m)
|
|||
|
|
maxCol := 0.0
|
|||
|
|
for j := range n {
|
|||
|
|
sum := 0.0
|
|||
|
|
for i := range n {
|
|||
|
|
sum += absOf(m[i][j])
|
|||
|
|
}
|
|||
|
|
// A NaN column sum can never win the strict max below, so it
|
|||
|
|
// propagates explicitly: the caller reads the norm as the
|
|||
|
|
// finiteness screen.
|
|||
|
|
if math.IsNaN(sum) {
|
|||
|
|
return math.NaN()
|
|||
|
|
}
|
|||
|
|
if sum > maxCol {
|
|||
|
|
maxCol = sum
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return maxCol
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// solveRight solves a·x = b for row-major square matrices and returns
|
|||
|
|
// x row-major. a is consumed by the factorisation, so callers hand
|
|||
|
|
// over a freshly built matrix; the column-oriented solveSystem is
|
|||
|
|
// reused underneath rather than a second LU being written.
|
|||
|
|
func solveRight[T scalar](name string, a, b [][]T) ([][]T, error) {
|
|||
|
|
n := len(a)
|
|||
|
|
colBack := make([]T, n*n)
|
|||
|
|
cols := make([][]T, n)
|
|||
|
|
for j := range n {
|
|||
|
|
col := colBack[j*n : (j+1)*n]
|
|||
|
|
for i := range n {
|
|||
|
|
col[i] = b[i][j]
|
|||
|
|
}
|
|||
|
|
cols[j] = col
|
|||
|
|
}
|
|||
|
|
sol, err := base.SolveSystem(name, a, cols)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
x := newSquare[T](n)
|
|||
|
|
for i := range n {
|
|||
|
|
for j := range n {
|
|||
|
|
x[i][j] = sol[j][i]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return x, nil
|
|||
|
|
}
|