trace and history added

numethods/solvers.py

@@ -1,9 +1,17 @@
 from __future__ import annotations
 from .linalg import Matrix, Vector, forward_substitution, backward_substitution
-from .exceptions import NonSquareMatrixError, SingularMatrixError, NotSymmetricError, NotPositiveDefiniteError, ConvergenceError
+from .exceptions import (
+    NonSquareMatrixError,
+    SingularMatrixError,
+    NotSymmetricError,
+    NotPositiveDefiniteError,
+    ConvergenceError,
+)
+
 
 class LUDecomposition:
-    """LU with partial pivoting: PA = LU"""
+    """LU decomposition with partial pivoting: PA = LU"""
+
     def __init__(self, A: Matrix):
         if not A.is_square():
             raise NonSquareMatrixError("A must be square")
@@ -11,6 +19,7 @@ class LUDecomposition:
         self.L = Matrix.identity(self.n)
         self.U = A.copy()
         self.P = Matrix.identity(self.n)
+        self.steps: list[tuple[int, Matrix, Matrix, Matrix]] = []  # store pivot steps
         self._decompose()
 
     def _decompose(self) -> None:
@@ -22,26 +31,47 @@ class LUDecomposition:
             self.U.swap_rows(k, pivot_row)
             self.P.swap_rows(k, pivot_row)
             if k > 0:
-                self.L.data[k][:k], self.L.data[pivot_row][:k] = self.L.data[pivot_row][:k], self.L.data[k][:k]
+                self.L.data[k][:k], self.L.data[pivot_row][:k] = (
+                    self.L.data[pivot_row][:k],
+                    self.L.data[k][:k],
+                )
             for i in range(k + 1, n):
                 m = self.U.data[i][k] / self.U.data[k][k]
                 self.L.data[i][k] = m
                 for j in range(k, n):
                     self.U.data[i][j] -= m * self.U.data[k][j]
+            # record step
+            self.steps.append((k, self.L.copy(), self.U.copy(), self.P.copy()))
 
     def solve(self, b: Vector) -> Vector:
-        Pb = Vector([sum(self.P.data[i][j]*b[j] for j in range(self.n)) for i in range(self.n)])
+        Pb = Vector(
+            [
+                sum(self.P.data[i][j] * b[j] for j in range(self.n))
+                for i in range(self.n)
+            ]
+        )
         y = forward_substitution(self.L, Pb)
         x = backward_substitution(self.U, y)
         return x
 
+    def trace(self):
+        print("LU Decomposition Trace (steps of elimination)")
+        for k, L, U, P in self.steps:
+            print(f"\nStep {k}:")
+            print(f"L = {L}")
+            print(f"U = {U}")
+            print(f"P = {P}")
+
 
 class GaussJordan:
+    """Gauss–Jordan elimination."""
+
     def __init__(self, A: Matrix):
         if not A.is_square():
             raise NonSquareMatrixError("A must be square")
         self.n = A.n
         self.A = A.copy()
+        self.steps: list[tuple[int, Matrix]] = []
 
     def solve(self, b: Vector) -> Vector:
         n = self.n
@@ -57,12 +87,26 @@ class GaussJordan:
                 if r == col:
                     continue
                 factor = Ab.data[r][col]
-                Ab.data[r] = [rv - factor*cv for rv, cv in zip(Ab.data[r], Ab.data[col])]
+                Ab.data[r] = [
+                    rv - factor * cv for rv, cv in zip(Ab.data[r], Ab.data[col])
+                ]
+            # record step
+            self.steps.append((col, Ab.copy()))
         return Vector(row[-1] for row in Ab.data)
 
+    def trace(self):
+        print("Gauss–Jordan Trace (row reduction steps)")
+        for step, Ab in self.steps:
+            print(f"\nColumn {step}:")
+            print(f"Augmented matrix = {Ab}")
+
 
 class Jacobi:
-    def __init__(self, A: Matrix, b: Vector, tol: float = 1e-10, max_iter: int = 10_000):
+    """Jacobi iterative method for Ax = b."""
+
+    def __init__(
+        self, A: Matrix, b: Vector, tol: float = 1e-10, max_iter: int = 10_000
+    ):
         if not A.is_square():
             raise NonSquareMatrixError("A must be square")
         if A.n != len(b):
@@ -71,6 +115,7 @@ class Jacobi:
         self.b = b.copy()
         self.tol = tol
         self.max_iter = max_iter
+        self.history: list[float] = []
 
     def solve(self, x0: Vector | None = None) -> Vector:
         n = self.A.n
@@ -84,14 +129,28 @@ class Jacobi:
                 s = sum(self.A.data[i][j] * x[j] for j in range(n) if j != i)
                 x_new[i] = (self.b[i] - s) / diag
             x_new = Vector(x_new)
-            if (x_new - x).norm_inf() <= self.tol * (1.0 + x_new.norm_inf()):
+            r = (self.A @ x_new) - self.b
+            res_norm = r.norm2()
+            self.history.append(res_norm)
+            if res_norm <= self.tol * (1.0 + x_new.norm2()):
                 return x_new
             x = x_new
         raise ConvergenceError("Jacobi did not converge within max_iter")
 
+    def trace(self):
+        print("Jacobi Iteration Trace")
+        print(f"{'iter':>6} | {'residual norm':>14}")
+        print("-" * 26)
+        for k, res in enumerate(self.history):
+            print(f"{k:6d} | {res:14.6e}")
+
 
 class GaussSeidel:
-    def __init__(self, A: Matrix, b: Vector, tol: float = 1e-10, max_iter: int = 10_000):
+    """Gauss–Seidel iterative method for Ax = b."""
+
+    def __init__(
+        self, A: Matrix, b: Vector, tol: float = 1e-10, max_iter: int = 10_000
+    ):
         if not A.is_square():
             raise NonSquareMatrixError("A must be square")
         if A.n != len(b):
@@ -100,6 +159,7 @@ class GaussSeidel:
         self.b = b.copy()
         self.tol = tol
         self.max_iter = max_iter
+        self.history: list[float] = []
 
     def solve(self, x0: Vector | None = None) -> Vector:
         n = self.A.n
@@ -113,13 +173,24 @@ class GaussSeidel:
                 s1 = sum(self.A.data[i][j] * x[j] for j in range(i))
                 s2 = sum(self.A.data[i][j] * x_old[j] for j in range(i + 1, n))
                 x[i] = (self.b[i] - s1 - s2) / diag
-            if (x - x_old).norm_inf() <= self.tol * (1.0 + x.norm_inf()):
+            r = (self.A @ x) - self.b
+            res_norm = r.norm2()
+            self.history.append(res_norm)
+            if res_norm <= self.tol * (1.0 + x.norm2()):
                 return x
         raise ConvergenceError("Gauss-Seidel did not converge within max_iter")
 
+    def trace(self):
+        print("Gauss–Seidel Iteration Trace")
+        print(f"{'iter':>6} | {'residual norm':>14}")
+        print("-" * 26)
+        for k, res in enumerate(self.history):
+            print(f"{k:6d} | {res:14.6e}")
+
 
 class Cholesky:
-    """A = L L^T for SPD matrices."""
+    """Cholesky factorization A = L L^T for SPD matrices."""
+
     def __init__(self, A: Matrix):
         if not A.is_square():
             raise NonSquareMatrixError("A must be square")
@@ -130,6 +201,7 @@ class Cholesky:
                     raise NotSymmetricError("Matrix is not symmetric")
         self.n = n
         self.L = Matrix.zeros(n, n)
+        self.steps: list[tuple[int, Matrix]] = []
         self._decompose(A.copy())
 
     def _decompose(self, A: Matrix) -> None:
@@ -140,12 +212,22 @@ class Cholesky:
                 if i == j:
                     val = A.data[i][i] - s
                     if val <= 0.0:
-                        raise NotPositiveDefiniteError("Matrix is not positive definite")
+                        raise NotPositiveDefiniteError(
+                            "Matrix is not positive definite"
+                        )
                     self.L.data[i][j] = val**0.5
                 else:
                     self.L.data[i][j] = (A.data[i][j] - s) / self.L.data[j][j]
+            # record after each row i
+            self.steps.append((i, self.L.copy()))
 
     def solve(self, b: Vector) -> Vector:
         y = forward_substitution(self.L, b)
         x = backward_substitution(self.L.transpose(), y)
         return x
+
+    def trace(self):
+        print("Cholesky Decomposition Trace")
+        for i, L in self.steps:
+            print(f"\nRow {i}:")
+            print(f"L = {L}")