linalg: add polar decomposition

aneforge/linalg.py

@@ -529,11 +529,19 @@ def eigvals(A, iters: int = 60):
   "eigh", "eigvals", "generalized_eigh", "dominant_eig",
   "svd", "dominant_svd", "svdvals_topk", "randomized_svd", "pca",
   "kron",
+  "polar",
 ]
 
 
 # __main__ - self-test / validation vs numpy/scipy
 
+def _general_square(n, cond, seed):
+  """A general (non-symmetric) n x n matrix with the target condition number: U diag(s) V^T."""
+  rng = np.random.default_rng(seed)
+  U = np.linalg.qr(rng.standard_normal((n, n)))[0]; V = np.linalg.qr(rng.standard_normal((n, n)))[0]
+  return (U * np.geomspace(1.0, cond, n)) @ V.T
+
+
 def _make_spd(n, cond, seed):
   """SPD A with target condition number, fp16-stored (reference solves the fp16-rounded system)."""
   rng = np.random.default_rng(seed)
@@ -791,6 +799,26 @@ def matrix_power(A, n: int):
   return np.asarray(acc, np.float32)
 
 
+def polar(A):
+  """Polar decomposition A = U @ P, on the ANE.
+
+  U is orthogonal for square A and semi-orthogonal for either rectangular shape (U^T U = I for
+  tall A, U U^T = I for wide A); P is symmetric positive-semidefinite [n,n].
+  Composed from the on-ANE SVD: U_ S V^T -> U = U_ V^T, P = V diag(S) V^T.
+  Oracle: scipy.linalg.polar(A, side='right')."""
+  A16 = np.asarray(A, f16)
+  if A16.ndim != 2: raise ValueError(f"polar: expected 2-D; got shape {A16.shape}")
+  m, n = A16.shape
+  U_svd, S, Vt = randomized_svd(A16, k=min(m, n), oversample=5, power_iters=2)
+  V = Vt.T.astype(f16)
+  Umat = _ane_gemm(U_svd.astype(f16), Vt)                        # U_ @ V^T -> [m,n]
+  Sdiag = np.diag(S).astype(f16)
+  Pmat = _ane_gemm(_ane_gemm(V, Sdiag), Vt)                      # V diag(S) V^T -> [n,n]
+  P = np.asarray(Pmat, np.float32)
+  P = 0.5 * (P + P.T)                                            # fp16 GEMMs leave ~1e-4 asymmetry; sym kills it
+  return np.asarray(Umat, np.float32), P
+
+
 def main():
   print("=" * 90)
   print("aneforge.linalg - ITERATIVE linear algebra on the ANE  (matmuls=ANE, RNG/QR/SVD/loop=HOST)")
@@ -903,6 +931,23 @@ def main():
         f"(~{ane_flops/host_flops:.0f}x on ANE)")
   print()
 
+  # ---------------- polar decomposition ------------------------------- #
+  print("-" * 90)
+  print("POLAR DECOMPOSITION  (A = U P, U orthogonal, P SPD)  -  vs scipy.linalg.polar")
+  print("-" * 90)
+  import scipy.linalg
+  for n in (6, 8):
+    # a GENERAL matrix: for symmetric positive-definite input polar is the trivial case (U = I, P = A)
+    A_pol = f16(_general_square(n, 1e1, 80 + n))
+    U, P = polar(A_pol)
+    ref_P = scipy.linalg.polar(A_pol, side="right")[1]
+    recon_err = _relerr(U @ P, A_pol)
+    orth_err = _relerr(U.T @ U, np.eye(n))
+    spd_err = _relerr(P, ref_P)
+    print(f"  n={n}:  recon relerr={recon_err:.3e}  U^T U ~ I relerr={orth_err:.3e}  "
+          f"P vs scipy relerr={spd_err:.3e}")
+  print()
+
   # ---------------- verdict ---------------------------------------------- #
   print("#" * 90)
   print("# VERDICT - iterative linear algebra on the ANE")

tests/test_linalg.py

@@ -492,3 +492,52 @@ def test_kron():
     max_err = max(max_err, err)
     assert np.allclose(got16, ref16, atol=5e-4, rtol=0), f"kron({m},{n},{p},{q}) max abs err {err}"
   assert max_err <= 5e-4
+# ----------------------------- polar decomposition ----------------------------- #
+
+@pytest.mark.parametrize("n", [6, 8])
+def test_polar_reconstruction(n):
+  A = f16(_square(n, 1e1, 100 + n))
+  U, P = L.polar(A)
+  assert relerr(U @ P, A) <= 5e-2, f"polar recon: relerr {relerr(U @ P, A)}"
+
+
+@pytest.mark.parametrize("n", [6, 8])
+def test_polar_orthogonality(n):
+  A = f16(_square(n, 1e1, 110 + n))
+  U, _ = L.polar(A)
+  assert relerr(U.T @ U, np.eye(n)) <= 5e-2, f"U^T U ~ I: relerr {relerr(U.T @ U, np.eye(n))}"
+
+
+@pytest.mark.parametrize("n", [6, 8])
+def test_polar_psd(n):
+  A = f16(_square(n, 1e1, 120 + n))
+  _, P = L.polar(A)
+  # P should be symmetric (fp16 randomized SVD introduces ~1e-4 asymmetry)
+  assert relerr(P, P.T) <= 1e-3, f"P not symmetric: relerr {relerr(P, P.T)}"
+  # P should be PSD (all eigenvalues >= 0)
+  eig = np.linalg.eigvalsh(P)
+  assert eig.min() >= -1e-3, f"P has negative eigenvalue: {eig.min()}"
+
+
+def test_polar_matches_scipy():
+  scipy_linalg = pytest.importorskip("scipy.linalg")
+  A = f16(_square(8, 1e1, 130))
+  U, P = L.polar(A)
+  ref_U, ref_P = scipy_linalg.polar(A, side="right")
+  assert relerr(P, ref_P) <= 5e-2, f"P vs scipy: relerr {relerr(P, ref_P)}"
+
+
+@pytest.mark.parametrize("m,n", [(8, 5), (5, 8)])
+def test_polar_rectangular(m, n):
+  """Both rectangular shapes: U is [m,n] semi-orthogonal, P is [n,n], and U P reconstructs A."""
+  A = np.asarray(np.random.default_rng(3).standard_normal((m, n)), f16)
+  U, P = L.polar(A)
+  assert U.shape == (m, n) and P.shape == (n, n)
+  assert relerr(U @ P, A) <= 5e-2, f"polar recon [{m},{n}]: relerr {relerr(U @ P, A)}"
+  I = U.T @ U if m >= n else U @ U.T
+  assert relerr(I, np.eye(I.shape[0])) <= 5e-2, f"U not semi-orthogonal: relerr {relerr(I, np.eye(I.shape[0]))}"
+
+
+def test_polar_rejects():
+  with pytest.raises(ValueError): L.polar(np.zeros(4, f16))              # not 2-D
+  with pytest.raises(ValueError): L.polar(np.zeros((2, 2, 2), f16))