feat: initial release

Assisted-by: GLM 5.3 Flash

linalg/expm.go

@@ -0,0 +1,328 @@
+// 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
+}