feat: initial release

Assisted-by: GLM 5.3 Flash

linalg/matrixfunc.go

@@ -0,0 +1,403 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package linalg
+
+import (
+	"math"
+	"math/cmplx"
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// Matrix functions through the eigen-decomposition. MatrixSqrt and
+// MatrixLog evaluate √A and ln(A) for a symmetric positive
+// semi-definite real matrix by diagonalising: A = VΛVᵀ, so
+// f(A) = V·f(Λ)·Vᵀ. Every other input, a Hermitian complex matrix
+// included, runs through the complex Schur decomposition described at
+// the individual functions. The positive semi-definite requirement
+// matters: the eigenvalues of such a matrix are non-negative, which
+// makes the square root and the logarithm real. A matrix with a
+// significantly negative eigenvalue is refused rather than silently
+// producing complex or NaN entries.
+
+// spectrumRelTol is the relative tolerance MatrixSqrt and the Schur
+// routes of both matrix functions use to judge their spectrum: an
+// eigenvalue counts as non-positive or singular when it sits within
+// this fraction of the spectrum's magnitude. Purely relative on
+// purpose: an absolute floor would refuse legitimate small-scale
+// inputs such as diag(1e-12, 2e-12). The symmetric route of MatrixLog
+// judges negativity against the rounding floor n·eps·λmax instead,
+// which admits far smaller positive eigenvalues; see MatrixLog.
+const spectrumRelTol = 1e-10
+
+// MatrixSqrt returns the principal square root √A of a symmetric
+// positive semi-definite real matrix, by diagonalisation: A = VΛVᵀ,
+// so √A = V·√Λ·Vᵀ. Eigenvalues below −tol·‖A‖ are refused (the
+// matrix is not positive semi-definite); small negative values from
+// rounding are clamped to zero.
+// For nonsymmetric matrices the answer runs through the complex Schur
+// decomposition: A = Q·T·Qᴴ reduces the function to the triangular
+// recurrence on T, the Schur-Parlett route. A real matrix answers in
+// real entries whenever the principal function is real, which for the
+// square root means no negative real eigenvalues and for the
+// logarithm none on the non-positive real axis; otherwise the call is
+// an error, honestly, rather than a complex answer in real clothing.
+func MatrixSqrt(a *core.Array) (*core.Array, error) {
+	n := a.Shape()[0]
+	if a.NDim() != 2 || n != a.Shape()[1] {
+		return nil, base.Errf("MatrixSqrt: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
+	}
+	if !isSymmetricMatrix(a) {
+		return matrixSqrtSchur(a)
+	}
+	vals, vecs, err := Eigen(a)
+	if err != nil {
+		return nil, base.Errf("MatrixSqrt: %w", err)
+	}
+	// Eigen returns dense float arrays, so the reconstruction reads the
+	// payloads directly rather than paying floatAt's dispatch per entry
+	// of this O(n³) triple product.
+	scale := 0.0
+	for _, lam := range vals.RawFloats() {
+		if v := math.Abs(lam); v > scale {
+			scale = v
+		}
+	}
+	tol := spectrumRelTol * scale
+	sqrtVals := make([]float64, n)
+	for i, lam := range vals.RawFloats() {
+		if lam < -tol {
+			return nil, base.Errf("MatrixSqrt: eigenvalue %v is negative; the matrix is not positive semi-definite", lam)
+		}
+		if lam < 0 {
+			lam = 0
+		}
+		sqrtVals[i] = math.Sqrt(lam)
+	}
+	// √A = V·diag(√λ)·Vᵀ.
+	out := core.New(core.Float, []int{n, n}...)
+	vMat := vecs.RawFloats()
+	for i := range n {
+		for j := range n {
+			s := 0.0
+			for k := range n {
+				s += vMat[i*n+k] * sqrtVals[k] * vMat[j*n+k]
+			}
+			out.RawFloats()[i*n+j] = s
+		}
+	}
+	return out, nil
+}
+
+// MatrixLog returns the principal matrix logarithm ln(A) of a real
+// symmetric positive-definite matrix, by the same diagonalisation
+// route as MatrixSqrt with ln applied to the eigenvalues. A negative
+// eigenvalue beyond the rounding floor n·eps·λmax is refused as
+// non-positive, and so is an exact zero: the principal logarithm of a
+// matrix with eigenvalues on the non-positive real axis leaves the
+// reals. A positive eigenvalue is logged whatever its size against
+// the spectrum, since ln 1e-300 is finite and exact to its own
+// digits.
+func MatrixLog(a *core.Array) (*core.Array, error) {
+	n := a.Shape()[0]
+	if a.NDim() != 2 || n != a.Shape()[1] {
+		return nil, base.Errf("MatrixLog: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
+	}
+	if !isSymmetricMatrix(a) {
+		return matrixLogSchur(a)
+	}
+	vals, vecs, err := Eigen(a)
+	if err != nil {
+		return nil, base.Errf("MatrixLog: %w", err)
+	}
+	// The negativity floor is the eigensolver's own rounding scale,
+	// n·eps·λmax, the floor the QR rank guard uses: a negative value
+	// beyond it is a genuine sign of a non-positive spectrum, one
+	// inside it is rounding of a zero, and a positive value is logged
+	// honestly however far below the spectrum's top it sits.
+	floor := float64(n) * base.EpsF * eigenMaxAbs(vals, n)
+	logVals := make([]float64, n)
+	for i, lam := range vals.RawFloats() {
+		if lam < -floor {
+			return nil, base.Errf("MatrixLog: eigenvalue %v is negative; the matrix is not positive definite", lam)
+		}
+		if lam <= 0 {
+			return nil, base.Errf("MatrixLog: eigenvalue %v is non-positive; the principal logarithm needs a positive-definite matrix", lam)
+		}
+		logVals[i] = math.Log(lam)
+	}
+	out := core.New(core.Float, []int{n, n}...)
+	vMat := vecs.RawFloats()
+	for i := range n {
+		for j := range n {
+			s := 0.0
+			for k := range n {
+				s += vMat[i*n+k] * logVals[k] * vMat[j*n+k]
+			}
+			out.RawFloats()[i*n+j] = s
+		}
+	}
+	return out, nil
+}
+
+// eigenMaxAbs returns the largest absolute eigenvalue.
+func eigenMaxAbs(vals *core.Array, n int) float64 {
+	m := 0.0
+	for _, v := range vals.RawFloats() {
+		if v := math.Abs(v); v > m {
+			m = v
+		}
+	}
+	return m
+}
+
+// isSymmetricMatrix reports whether a square real matrix is symmetric
+// to rounding, which routes the matrix functions to the cheaper
+// symmetric eigendecomposition. The tolerance is purely relative to
+// the matrix scale, deliberately without an absolute floor: a floor
+// would call a small-scale matrix symmetric on asymmetry it carries in
+// full relative measure, and the Eigen route then applies its own
+// relative check and refuses the matrix the routing just approved.
+// Relative-only keeps the two decisions in agreement, because any
+// asymmetry this check passes is within the fraction of the scale the
+// route's own guard tolerates.
+func isSymmetricMatrix(a *core.Array) bool {
+	if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] || a.Dtype() == core.Complex {
+		return false
+	}
+	n := a.Shape()[0]
+	scale := 0.0
+	for i := range n * n {
+		scale = math.Max(scale, math.Abs(a.FloatAt(i)))
+	}
+	tol := 1e-12 * scale
+	for i := range n {
+		for j := i + 1; j < n; j++ {
+			if math.Abs(a.FloatAt(i*n+j)-a.FloatAt(j*n+i)) > tol {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+// matrixSqrtSchur evaluates the principal square root through the
+// complex Schur decomposition and the triangular recurrence. A real
+// matrix with a negative real eigenvalue has a complex principal root
+// and is refused; a real answer is otherwise recovered to rounding.
+func matrixSqrtSchur(a *core.Array) (*core.Array, error) {
+	const name = "MatrixSqrt"
+	n := a.Shape()[0]
+	t, q, err := SchurComplex(a)
+	if err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	tm := make([]complex128, n*n)
+	copy(tm, t.RawComplexes())
+	if a.Dtype() != core.Complex {
+		tol := spectrumRelTol * cmplxAbsMax(tm, n)
+		for i := range n {
+			lam := tm[i*n+i]
+			if math.Abs(imag(lam)) <= tol && real(lam) < -tol {
+				return nil, base.Errf("%s: eigenvalue %v is negative real; the principal square root is complex", name, lam)
+			}
+		}
+	}
+	if err := triSqrt(tm, n); err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	return schurBackTransform(q, tm, n, a.Dtype() != core.Complex, name)
+}
+
+// matrixLogSchur evaluates the principal logarithm through the complex
+// Schur decomposition by inverse scaling and squaring: repeated
+// triangular square roots walk the matrix near the identity, where the
+// Mercator series finishes the job, and the squarings unwind. A real
+// matrix with an eigenvalue on the non-positive real axis has no real
+// principal logarithm and is refused.
+func matrixLogSchur(a *core.Array) (*core.Array, error) {
+	const name = "MatrixLog"
+	n := a.Shape()[0]
+	t, q, err := SchurComplex(a)
+	if err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	tm := make([]complex128, n*n)
+	copy(tm, t.RawComplexes())
+	logScale := 0.0
+	// Spectrum screen. Real input refuses eigenvalues on the
+	// non-positive real axis, where the principal logarithm leaves the
+	// reals; complex input refuses eigenvalues at or below the
+	// rounding floor of the spectrum, where the square-root walk below
+	// would diverge into garbage instead of converging.
+	tol := spectrumRelTol * cmplxAbsMax(tm, n)
+	for i := range n {
+		lam := tm[i*n+i]
+		if a.Dtype() != core.Complex {
+			if math.Abs(imag(lam)) <= tol && real(lam) <= tol {
+				return nil, base.Errf("%s: eigenvalue %v is on the non-positive real axis; the principal logarithm is complex", name, lam)
+			}
+		} else if cmplx.Abs(lam) <= tol {
+			return nil, base.Errf("%s: eigenvalue %v is at the rounding floor of the spectrum; the matrix is singular and has no principal logarithm", name, lam)
+		}
+	}
+	// Scale the spectrum towards one so the square-root walk converges
+	// in a handful of steps; the scalar shift unwinds afterwards.
+	c := 0.0
+	for i := range n {
+		c = math.Max(c, cmplx.Abs(tm[i*n+i]))
+	}
+	if c > 0 && (c < 0.5 || c > 2) {
+		logScale = math.Log(c)
+		for i := range n {
+			for j := range n {
+				tm[i*n+j] /= complex(c, 0)
+			}
+		}
+	}
+	// Square-root walk towards the identity.
+	squarings := 0
+	for range 60 {
+		dev := triDeviation(tm, n)
+		if dev <= 0.25 {
+			break
+		}
+		if err := triSqrt(tm, n); err != nil {
+			return nil, base.Errf("%s: %w", name, err)
+		}
+		squarings++
+	}
+	// A spectrum the walk cannot bring near the identity (typically a
+	// near-singular one the screen let through) would send the Mercator
+	// series below into open-ended divergence; refuse instead.
+	if dev := triDeviation(tm, n); dev > 0.25 {
+		return nil, base.Errf("%s: the square-root walk did not converge in 60 passes (deviation %.3g); the matrix is outside the domain of the principal logarithm", name, dev)
+	}
+	// Mercator series for log(I + E) on the near-identity triangle.
+	e := make([]complex128, n*n)
+	for i := range n {
+		for j := range n {
+			e[i*n+j] = tm[i*n+j]
+		}
+		e[i*n+i] -= 1
+	}
+	lg := make([]complex128, n*n)
+	// The series multiplies term by e each step, so the two buffers
+	// alternate instead of the product landing in a fresh matrix: every
+	// value read is the previous step's product exactly as it was when
+	// the product was copied back, because only the upper triangle is
+	// ever written and both buffers start with a zero lower triangle.
+	term, next := make([]complex128, n*n), make([]complex128, n*n)
+	copy(term, e)
+	for m := 1; m <= 40; m++ {
+		factor := complex(1/float64(m), 0)
+		if m%2 == 0 {
+			factor = complex(-1/float64(m), 0)
+		}
+		for i := range n * n {
+			lg[i] += factor * term[i]
+		}
+		if m == 40 {
+			break
+		}
+		for i := range n {
+			for j := i; j < n; j++ {
+				s := complex(0, 0)
+				for k := i; k <= j; k++ {
+					s += term[i*n+k] * e[k*n+j]
+				}
+				next[i*n+j] = s
+			}
+		}
+		tnorm := 0.0
+		for i := range n * n {
+			tnorm = math.Max(tnorm, cmplx.Abs(next[i]))
+		}
+		term, next = next, term
+		if tnorm <= 1e-17 {
+			break
+		}
+	}
+	// Unwind the walk: log(T) = 2^k·L + log(c)·I.
+	for i := range n * n {
+		lg[i] *= complex(math.Pow(2, float64(squarings)), 0)
+	}
+	if logScale != 0 {
+		for i := range n {
+			lg[i*n+i] += complex(logScale, 0)
+		}
+	}
+	return schurBackTransform(q, lg, n, a.Dtype() != core.Complex, name)
+}
+
+// schurBackTransform maps the triangular result back through the Schur
+// unitary, q·S·qᴴ, and returns real entries for real input when the
+// imaginary rounding cancels.
+func schurBackTransform(q *core.Array, s []complex128, n int, wantReal bool, name string) (*core.Array, error) {
+	qm := q.RawComplexes()
+	w := make([]complex128, n*n)
+	for i := range n {
+		for j := range n {
+			sum := complex(0, 0)
+			for k := i; k < n; k++ {
+				sum += s[i*n+k] * cmplx.Conj(qm[j*n+k])
+			}
+			w[i*n+j] = sum
+		}
+	}
+	out := make([]complex128, n*n)
+	for i := range n {
+		for j := range n {
+			sum := complex(0, 0)
+			for k := range n {
+				sum += qm[i*n+k] * w[k*n+j]
+			}
+			out[i*n+j] = sum
+		}
+	}
+	if !wantReal {
+		res, _ := core.FromComplexes(out, n, n)
+		return res, nil
+	}
+	worstIm, worstRe := 0.0, 0.0
+	for i := range n * n {
+		worstIm = math.Max(worstIm, math.Abs(imag(out[i])))
+		worstRe = math.Max(worstRe, math.Abs(real(out[i])))
+	}
+	if worstIm > 1e-8*math.Max(1, worstRe) {
+		return nil, base.Errf("%s: the result is complex (imaginary magnitude %.3g); the real answer does not exist", name, worstIm)
+	}
+	res := core.New(core.Float, []int{n, n}...)
+	for i := range n * n {
+		res.RawFloats()[i] = real(out[i])
+	}
+	return res, nil
+}
+
+// cmplxAbsMax returns the largest diagonal magnitude of a complex
+// triangular matrix.
+func cmplxAbsMax(t []complex128, n int) float64 {
+	m := 0.0
+	for i := range n {
+		m = math.Max(m, cmplx.Abs(t[i*n+i]))
+	}
+	return m
+}
+
+// triDeviation measures how far an upper triangular matrix still sits
+// from the identity: the largest magnitude among the diagonal offsets
+// and the strict upper part, the convergence gauge of the square-root
+// walk in matrixLogSchur.
+func triDeviation(t []complex128, n int) float64 {
+	dev := 0.0
+	for i := range n {
+		dev = math.Max(dev, cmplx.Abs(t[i*n+i]-1))
+	}
+	for i := range n {
+		for j := i + 1; j < n; j++ {
+			dev = math.Max(dev, cmplx.Abs(t[i*n+j]))
+		}
+	}
+	return dev
+}