mathmat.solve
MathMat Linear System Solving Toolbox
1"""MathMat Linear System Solving Toolbox""" 2 3import lazy_import 4 5from .mathmat import Matrix, Vector, MatrixSizeException, MatrixTypeException 6from .factor import lu as lu_factor 7import numpy as np 8from scipy.linalg import solve_triangular 9 10mathmat_random = lazy_import.lazy_module("mathmat.random") 11 12 13class Krylov: 14 """A collection of methods for generating Krylov subspaces.""" 15 16 @staticmethod 17 def arnoldi_mgs(A: Matrix, b: Vector, k: int, 18 tolerance=1e-4, **_): 19 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 20 Uses Modified Gram-Schmidt (non-repeated) full orthogonalization. 21 Optionally configure the tolerance for stopping early.""" 22 # If rank is known, restrict max iteration 23 k = min(k, A._computed.get("rank", k)) 24 # Check if we need to use conjugates 25 use_conj = A.is_complex() 26 27 def dot(q, v): return np.dot(np.conjugate( 28 q), v) if use_conj else np.dot(q, v) 29 # The Krylov basis vectors. 30 K_k = np.zeros((len(b), k), 31 dtype=np._complex if use_conj else np.double) 32 # Set the first vector to normalized b 33 K_k[:, 0] = b.unit().entries.ravel() 34 # The upper Hessenberg matrix 35 H = np.zeros((k+1, k), 36 dtype=np._complex if use_conj else np.double) 37 38 # Outer iteration 39 for j in range(k-1): 40 # Arnoldi step 41 v = A.entries @ K_k[:, j] 42 43 # Orthogonalization 44 for i in range(j+1): 45 # Get the Krylov vector to orthogonalize against 46 q_i = K_k[:, i] 47 # Compute the Hessenberg term 48 H[i, j] = dot(q_i, v) 49 # Orthogonalize 50 v -= H[i, j] * q_i 51 52 # Fill in the sub-diagonal entry 53 H[j+1, j] = np.linalg.norm(v, ord=2) 54 55 if j > 0 and H[j+1, j] <= tolerance: 56 # Exact solution reached, trim the generated subspace 57 K_k = K_k[:, :j-1] 58 H = H[:j, :j-1] 59 # We have learned the rank of A 60 A.set("rank", j-1) 61 break 62 63 # Normalize and store the vector 64 v /= H[j+1, j] 65 K_k[:, j+1] = v 66 67 K_k = Matrix(K_k) 68 K_k.set("rank", K_k.entries.shape[1]) 69 K_k.set("unitary" if use_conj else "orthogonal", True) 70 71 H = Matrix(H) 72 H.set("upper_hessenberg", True) 73 74 return K_k, H 75 76 @staticmethod 77 def arnoldi_mgs_trunc(A: Matrix, b: Vector, k: int, 78 tolerance=1e-4, r=10): 79 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 80 Uses Modified Gram-Schmidt (non-repeated) `r`-truncated 81 orthogonalization (see Nakatuskasa and Tropp 2022). 82 Optionally configure the tolerance for stopping early.""" 83 # If rank is known, restrict max iteration 84 k = min(k, A._computed.get("rank", k)) 85 # Check if we need to use conjugates 86 use_conj = A.is_complex() 87 88 def dot(q, v): return np.dot(np.conjugate( 89 q), v) if use_conj else np.dot(q, v) 90 # The Krylov basis vectors. 91 K_k = np.zeros((len(b), k), 92 dtype=np._complex if use_conj else np.double) 93 # Set the first vector to normalized b 94 K_k[:, 0] = b.unit().entries.ravel() 95 # The upper Hessenberg matrix 96 H = np.zeros((k+1, k), 97 dtype=np._complex if use_conj else np.double) 98 99 # Outer iteration 100 for j in range(k-1): 101 # Arnoldi step 102 v = A.entries @ K_k[:, j] 103 104 # Truncated Orthogonalization 105 for i in range(max(0, j-r), j+1): 106 # Get the Krylov vector to orthogonalize against 107 q_i = K_k[:, i] 108 # Compute the Hessenberg term 109 H[i, j] = dot(q_i, v) 110 # Orthogonalize 111 v -= H[i, j] * q_i 112 113 # Fill in the sub-diagonal entry 114 H[j+1, j] = np.linalg.norm(v, ord=2) 115 116 if j > 0 and H[j+1, j] <= tolerance: 117 # Exact solution reached, trim the generated subspace 118 K_k = K_k[:, :j-1] 119 H = H[:j, :j-1] 120 # We have learned the rank of A 121 A.set("rank", j-1) 122 break 123 124 # Normalize and store the vector 125 v /= H[j+1, j] 126 K_k[:, j+1] = v 127 128 K_k = Matrix(K_k) 129 H = Matrix(H) 130 H.set("upper_hessenberg", True) 131 132 return K_k, H 133 134 @staticmethod 135 def arnoldi_mgs_2x(A: Matrix, b: Vector, k: int, 136 tolerance=1e-4, **_): 137 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 138 Uses Modified Gram-Schmidt (twice) full orthogonalization. 139 Optionally configure the tolerance for stopping early.""" 140 # If rank is known, restrict max iteration 141 k = min(k, A._computed.get("rank", k)) 142 # Check if we need to use conjugates 143 use_conj = A.is_complex() 144 145 def dot(q, v): return np.dot(np.conjugate( 146 q), v) if use_conj else np.dot(q, v) 147 # The Krylov basis vectors. 148 K_k = np.zeros((len(b), k), 149 dtype=np._complex if use_conj else np.double) 150 # Set the first vector to normalized b 151 K_k[:, 0] = b.unit().entries.ravel() 152 # The upper Hessenberg matrix 153 H = np.zeros((k+1, k), 154 dtype=np._complex if use_conj else np.double) 155 156 # Outer iteration 157 for j in range(k-1): 158 # Arnoldi step 159 v = A.entries @ K_k[:, j] 160 161 # Orthogonalization 162 for i in range(j+1): 163 # Get the Krylov vector to orthogonalize against 164 q_i = K_k[:, i] 165 # Compute the Hessenberg term 166 H[i, j] = dot(q_i, v) 167 # Orthogonalize once 168 v -= H[i, j] * q_i 169 # Orthogonalize again 170 h_tmp = dot(q_i, v) 171 H[i, j] += h_tmp 172 v -= h_tmp * q_i 173 174 # Fill in the sub-diagonal entry 175 H[j+1, j] = np.linalg.norm(v, ord=2) 176 177 if j > 0 and H[j+1, j] <= tolerance: 178 # Exact solution reached, trim the generated subspace 179 K_k = K_k[:, :j-1] 180 H = H[:j, :j-1] 181 # We have learned the rank of A 182 A.set("rank", j-1) 183 break 184 185 # Normalize and store the vector 186 v /= H[j+1, j] 187 K_k[:, j+1] = v 188 189 K_k = Matrix(K_k) 190 K_k.set("rank", K_k.entries.shape[1]) 191 K_k.set("unitary" if use_conj else "orthogonal", True) 192 193 H = Matrix(H) 194 H.set("upper_hessenberg", True) 195 196 return K_k, H 197 198 199def gmres(A: Matrix, b: Vector, k: int, 200 krylov="auto", tolerance=1e-4): 201 """Solve Ax=b for x iteratively using the GMRES algorithm. 202 Stops after `k` iterations, or when the sub-diagonal entry 203 is below the given `tolerance`. 204 205 Optionally specify the method of generating the Krylov 206 subspace by passing a callable to `krylov`. Otherwise, 207 this is derermined automatically from A. The callable 208 should accept the arguments `(A, b, k, tol)` and return the 209 Matrices $K_k$ and $H$. 210 """ 211 if krylov == "auto": 212 krylov = Krylov.arnoldi_mgs 213 214 # Use the specified method to compute the basis 215 K_k, H = krylov(A, b, k, tolerance) 216 217 # Solve the least-squares problem 218 rhs = np.zeros((H.nr, 1)) 219 rhs[0] = b.norm_2() 220 221 y = lstsq(H, Vector(rhs)) 222 x = K_k @ y 223 224 return x, K_k, H 225 226 227def sgmres(A: Matrix, b: Vector, k: int, r: int, 228 krylov="auto", tolerance=1e-4): 229 """Solve Ax=b for x iteratively using the sGMRES algorithm. 230 Generates a Krylov subspace of dimension `k` via a truncated 231 Arnoldi process that orthogonalizes against `r` previous vectors. 232 The solution is then produced by sketching the matrix using a FFT. 233 234 Optionally specify the method of generating the Krylov 235 subspace by passing a callable to `krylov`. Otherwise, 236 truncated MGS Arnoldi will be used. The callable 237 should accept the arguments `(A, b, k, tol, r)` and return the 238 Matrices $K_k$ and $H$. 239 """ 240 if krylov == "auto": 241 krylov = Krylov.arnoldi_mgs_trunc 242 243 # Use the specified method to compute the basis 244 K_k, H = krylov(A, b, k, tolerance, r) 245 246 # Sketch 247 SA = mathmat_random.FFTSketchedMatrix(A @ K_k) 248 Sb = mathmat_random.FFTSketchedMatrix(b) 249 rows = np.random.choice(len(b), min(len(b), 2*K_k.nc)) 250 SA = SA.entries[rows, :] 251 Sb = Sb.entries[rows, :] 252 253 x = np.linalg.lstsq(SA, Sb, rcond=None)[0] 254 x = K_k.entries @ x 255 256 return Vector(x), K_k, H 257 258 259def invert(A: Matrix, b: Vector): 260 """Solve $Ax=b$ for $x$ by inverting the matrix $A$.""" 261 if A.nr != b.nr: 262 raise MatrixSizeException(A, b, "Linear System Solve") 263 return A.inverse() @ b 264 265 266def lstsq(A: Matrix, b: Vector): 267 """Find the least-squares, min. norm solution to $Ax=b$.""" 268 if A.nr != b.nr: 269 raise MatrixSizeException(A, b, "Least Squares Solve") 270 if "qr" in A._computed: 271 # Use the QR of A if it exists 272 Q, R = A.qr() 273 return R.lin_solver()(Q.transpose() @ b) 274 # Computing lstsq also reveals some properties of $A$. 275 x, _, rank, sigmas = np.linalg.lstsq(A.entries, b.entries, rcond=None) 276 if "rank" not in A._computed: 277 A.set("rank", rank) 278 if "sigmas" not in A._computed: 279 A.set("sigmas", np.flip(np.sort(sigmas))) 280 return Vector(x) 281 282 283def lu(A: Matrix, b: Vector): 284 """Solve $Ax=b$ for $x% by taking an $LU$ factorization of $A$.""" 285 if not A.is_square(): 286 raise MatrixTypeException("LU factorization", "square") 287 if A.is_sparse(): 288 # We obtain the sparse LU solver 289 A.set("lin_solver", lambda v: Vector(lu_factor(A).solve(v.entries))) 290 return Vector(A.lin_solver()(b)) 291 else: 292 P, L, U = lu_factor(A) 293 y = L.lin_solver()(P.inverse() @ b) 294 return U.lin_solver()(y) 295 296 297def triangular(A: Matrix, b: Vector, tolerance=1e-10): 298 """Solve $Ax=b$ for $x$ for a triangular $A$.""" 299 if not (A.is_triangular_U(tolerance) or A.is_triangular_L(tolerance)): 300 raise MatrixTypeException("triangular solve", "triangular (any)") 301 return Vector(solve_triangular(A.entries, b.entries, check_finite=False)) 302 303 304def automatic(A: Matrix, b: Vector, get_method=False): 305 """Try to pick an optimal solving strategy based on the properties of $A$. 306 307 Note that computing some properties may be expensive, and therefore using 308 `automatic` may be sub-optimal if used only once. 309 Optionally, if `get_method` is `True`, return the solver method instead of 310 actually solving for $x$. 311 """ 312 if "lin_solver" not in A._computed: 313 # We need to find a suitable linear system solver 314 if "inverse" in A._computed: 315 # An inverse exists and thus a direct solution is available. 316 A.set("lin_solver", lambda v: invert(A, v)) 317 A.lin_solver().name = "auto: invert" 318 elif A.is_triangular_L() or A.is_triangular_U(): 319 # A triangular matrix can be solved using a faster algorithm. 320 A.set("lin_solver", lambda v: triangular(A, v)) 321 A.lin_solver().name = "auto: triangular" 322 elif "qr" in A._computed: 323 # If a QR exists then lstsq will use it 324 A.set("lin_solver", lambda v: lstsq(A, v)) 325 A.lin_solver().name = "auto: qr" 326 elif "plu" in A._computed or A.is_square(): 327 # If a LU (or Cholesky) exists or can be found, use the LU solver 328 A.set("lin_solver", lambda v: lu(A, v)) 329 A.lin_solver().name = "auto: lu" 330 else: 331 # The matrix is nonsquare and has no available factorizations. 332 A.set("lin_solver", lambda v: gmres(A, v, A.nc)[0]) 333 A.lin_solver().name = "auto: gmres" 334 335 if get_method: 336 return A._computed["lin_solver"] 337 return A.lin_solver(b)
14class Krylov: 15 """A collection of methods for generating Krylov subspaces.""" 16 17 @staticmethod 18 def arnoldi_mgs(A: Matrix, b: Vector, k: int, 19 tolerance=1e-4, **_): 20 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 21 Uses Modified Gram-Schmidt (non-repeated) full orthogonalization. 22 Optionally configure the tolerance for stopping early.""" 23 # If rank is known, restrict max iteration 24 k = min(k, A._computed.get("rank", k)) 25 # Check if we need to use conjugates 26 use_conj = A.is_complex() 27 28 def dot(q, v): return np.dot(np.conjugate( 29 q), v) if use_conj else np.dot(q, v) 30 # The Krylov basis vectors. 31 K_k = np.zeros((len(b), k), 32 dtype=np._complex if use_conj else np.double) 33 # Set the first vector to normalized b 34 K_k[:, 0] = b.unit().entries.ravel() 35 # The upper Hessenberg matrix 36 H = np.zeros((k+1, k), 37 dtype=np._complex if use_conj else np.double) 38 39 # Outer iteration 40 for j in range(k-1): 41 # Arnoldi step 42 v = A.entries @ K_k[:, j] 43 44 # Orthogonalization 45 for i in range(j+1): 46 # Get the Krylov vector to orthogonalize against 47 q_i = K_k[:, i] 48 # Compute the Hessenberg term 49 H[i, j] = dot(q_i, v) 50 # Orthogonalize 51 v -= H[i, j] * q_i 52 53 # Fill in the sub-diagonal entry 54 H[j+1, j] = np.linalg.norm(v, ord=2) 55 56 if j > 0 and H[j+1, j] <= tolerance: 57 # Exact solution reached, trim the generated subspace 58 K_k = K_k[:, :j-1] 59 H = H[:j, :j-1] 60 # We have learned the rank of A 61 A.set("rank", j-1) 62 break 63 64 # Normalize and store the vector 65 v /= H[j+1, j] 66 K_k[:, j+1] = v 67 68 K_k = Matrix(K_k) 69 K_k.set("rank", K_k.entries.shape[1]) 70 K_k.set("unitary" if use_conj else "orthogonal", True) 71 72 H = Matrix(H) 73 H.set("upper_hessenberg", True) 74 75 return K_k, H 76 77 @staticmethod 78 def arnoldi_mgs_trunc(A: Matrix, b: Vector, k: int, 79 tolerance=1e-4, r=10): 80 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 81 Uses Modified Gram-Schmidt (non-repeated) `r`-truncated 82 orthogonalization (see Nakatuskasa and Tropp 2022). 83 Optionally configure the tolerance for stopping early.""" 84 # If rank is known, restrict max iteration 85 k = min(k, A._computed.get("rank", k)) 86 # Check if we need to use conjugates 87 use_conj = A.is_complex() 88 89 def dot(q, v): return np.dot(np.conjugate( 90 q), v) if use_conj else np.dot(q, v) 91 # The Krylov basis vectors. 92 K_k = np.zeros((len(b), k), 93 dtype=np._complex if use_conj else np.double) 94 # Set the first vector to normalized b 95 K_k[:, 0] = b.unit().entries.ravel() 96 # The upper Hessenberg matrix 97 H = np.zeros((k+1, k), 98 dtype=np._complex if use_conj else np.double) 99 100 # Outer iteration 101 for j in range(k-1): 102 # Arnoldi step 103 v = A.entries @ K_k[:, j] 104 105 # Truncated Orthogonalization 106 for i in range(max(0, j-r), j+1): 107 # Get the Krylov vector to orthogonalize against 108 q_i = K_k[:, i] 109 # Compute the Hessenberg term 110 H[i, j] = dot(q_i, v) 111 # Orthogonalize 112 v -= H[i, j] * q_i 113 114 # Fill in the sub-diagonal entry 115 H[j+1, j] = np.linalg.norm(v, ord=2) 116 117 if j > 0 and H[j+1, j] <= tolerance: 118 # Exact solution reached, trim the generated subspace 119 K_k = K_k[:, :j-1] 120 H = H[:j, :j-1] 121 # We have learned the rank of A 122 A.set("rank", j-1) 123 break 124 125 # Normalize and store the vector 126 v /= H[j+1, j] 127 K_k[:, j+1] = v 128 129 K_k = Matrix(K_k) 130 H = Matrix(H) 131 H.set("upper_hessenberg", True) 132 133 return K_k, H 134 135 @staticmethod 136 def arnoldi_mgs_2x(A: Matrix, b: Vector, k: int, 137 tolerance=1e-4, **_): 138 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 139 Uses Modified Gram-Schmidt (twice) full orthogonalization. 140 Optionally configure the tolerance for stopping early.""" 141 # If rank is known, restrict max iteration 142 k = min(k, A._computed.get("rank", k)) 143 # Check if we need to use conjugates 144 use_conj = A.is_complex() 145 146 def dot(q, v): return np.dot(np.conjugate( 147 q), v) if use_conj else np.dot(q, v) 148 # The Krylov basis vectors. 149 K_k = np.zeros((len(b), k), 150 dtype=np._complex if use_conj else np.double) 151 # Set the first vector to normalized b 152 K_k[:, 0] = b.unit().entries.ravel() 153 # The upper Hessenberg matrix 154 H = np.zeros((k+1, k), 155 dtype=np._complex if use_conj else np.double) 156 157 # Outer iteration 158 for j in range(k-1): 159 # Arnoldi step 160 v = A.entries @ K_k[:, j] 161 162 # Orthogonalization 163 for i in range(j+1): 164 # Get the Krylov vector to orthogonalize against 165 q_i = K_k[:, i] 166 # Compute the Hessenberg term 167 H[i, j] = dot(q_i, v) 168 # Orthogonalize once 169 v -= H[i, j] * q_i 170 # Orthogonalize again 171 h_tmp = dot(q_i, v) 172 H[i, j] += h_tmp 173 v -= h_tmp * q_i 174 175 # Fill in the sub-diagonal entry 176 H[j+1, j] = np.linalg.norm(v, ord=2) 177 178 if j > 0 and H[j+1, j] <= tolerance: 179 # Exact solution reached, trim the generated subspace 180 K_k = K_k[:, :j-1] 181 H = H[:j, :j-1] 182 # We have learned the rank of A 183 A.set("rank", j-1) 184 break 185 186 # Normalize and store the vector 187 v /= H[j+1, j] 188 K_k[:, j+1] = v 189 190 K_k = Matrix(K_k) 191 K_k.set("rank", K_k.entries.shape[1]) 192 K_k.set("unitary" if use_conj else "orthogonal", True) 193 194 H = Matrix(H) 195 H.set("upper_hessenberg", True) 196 197 return K_k, H
A collection of methods for generating Krylov subspaces.
17 @staticmethod 18 def arnoldi_mgs(A: Matrix, b: Vector, k: int, 19 tolerance=1e-4, **_): 20 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 21 Uses Modified Gram-Schmidt (non-repeated) full orthogonalization. 22 Optionally configure the tolerance for stopping early.""" 23 # If rank is known, restrict max iteration 24 k = min(k, A._computed.get("rank", k)) 25 # Check if we need to use conjugates 26 use_conj = A.is_complex() 27 28 def dot(q, v): return np.dot(np.conjugate( 29 q), v) if use_conj else np.dot(q, v) 30 # The Krylov basis vectors. 31 K_k = np.zeros((len(b), k), 32 dtype=np._complex if use_conj else np.double) 33 # Set the first vector to normalized b 34 K_k[:, 0] = b.unit().entries.ravel() 35 # The upper Hessenberg matrix 36 H = np.zeros((k+1, k), 37 dtype=np._complex if use_conj else np.double) 38 39 # Outer iteration 40 for j in range(k-1): 41 # Arnoldi step 42 v = A.entries @ K_k[:, j] 43 44 # Orthogonalization 45 for i in range(j+1): 46 # Get the Krylov vector to orthogonalize against 47 q_i = K_k[:, i] 48 # Compute the Hessenberg term 49 H[i, j] = dot(q_i, v) 50 # Orthogonalize 51 v -= H[i, j] * q_i 52 53 # Fill in the sub-diagonal entry 54 H[j+1, j] = np.linalg.norm(v, ord=2) 55 56 if j > 0 and H[j+1, j] <= tolerance: 57 # Exact solution reached, trim the generated subspace 58 K_k = K_k[:, :j-1] 59 H = H[:j, :j-1] 60 # We have learned the rank of A 61 A.set("rank", j-1) 62 break 63 64 # Normalize and store the vector 65 v /= H[j+1, j] 66 K_k[:, j+1] = v 67 68 K_k = Matrix(K_k) 69 K_k.set("rank", K_k.entries.shape[1]) 70 K_k.set("unitary" if use_conj else "orthogonal", True) 71 72 H = Matrix(H) 73 H.set("upper_hessenberg", True) 74 75 return K_k, H
Generate $K_k = \text{span} ( b, Ab, ... A^{k-1}b )$. Uses Modified Gram-Schmidt (non-repeated) full orthogonalization. Optionally configure the tolerance for stopping early.
77 @staticmethod 78 def arnoldi_mgs_trunc(A: Matrix, b: Vector, k: int, 79 tolerance=1e-4, r=10): 80 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 81 Uses Modified Gram-Schmidt (non-repeated) `r`-truncated 82 orthogonalization (see Nakatuskasa and Tropp 2022). 83 Optionally configure the tolerance for stopping early.""" 84 # If rank is known, restrict max iteration 85 k = min(k, A._computed.get("rank", k)) 86 # Check if we need to use conjugates 87 use_conj = A.is_complex() 88 89 def dot(q, v): return np.dot(np.conjugate( 90 q), v) if use_conj else np.dot(q, v) 91 # The Krylov basis vectors. 92 K_k = np.zeros((len(b), k), 93 dtype=np._complex if use_conj else np.double) 94 # Set the first vector to normalized b 95 K_k[:, 0] = b.unit().entries.ravel() 96 # The upper Hessenberg matrix 97 H = np.zeros((k+1, k), 98 dtype=np._complex if use_conj else np.double) 99 100 # Outer iteration 101 for j in range(k-1): 102 # Arnoldi step 103 v = A.entries @ K_k[:, j] 104 105 # Truncated Orthogonalization 106 for i in range(max(0, j-r), j+1): 107 # Get the Krylov vector to orthogonalize against 108 q_i = K_k[:, i] 109 # Compute the Hessenberg term 110 H[i, j] = dot(q_i, v) 111 # Orthogonalize 112 v -= H[i, j] * q_i 113 114 # Fill in the sub-diagonal entry 115 H[j+1, j] = np.linalg.norm(v, ord=2) 116 117 if j > 0 and H[j+1, j] <= tolerance: 118 # Exact solution reached, trim the generated subspace 119 K_k = K_k[:, :j-1] 120 H = H[:j, :j-1] 121 # We have learned the rank of A 122 A.set("rank", j-1) 123 break 124 125 # Normalize and store the vector 126 v /= H[j+1, j] 127 K_k[:, j+1] = v 128 129 K_k = Matrix(K_k) 130 H = Matrix(H) 131 H.set("upper_hessenberg", True) 132 133 return K_k, H
Generate $K_k = \text{span} ( b, Ab, ... A^{k-1}b )$.
Uses Modified Gram-Schmidt (non-repeated) r
-truncated
orthogonalization (see Nakatuskasa and Tropp 2022).
Optionally configure the tolerance for stopping early.
135 @staticmethod 136 def arnoldi_mgs_2x(A: Matrix, b: Vector, k: int, 137 tolerance=1e-4, **_): 138 """Generate $K_k = \\text{span} ( b, Ab, ... A^{k-1}b )$. 139 Uses Modified Gram-Schmidt (twice) full orthogonalization. 140 Optionally configure the tolerance for stopping early.""" 141 # If rank is known, restrict max iteration 142 k = min(k, A._computed.get("rank", k)) 143 # Check if we need to use conjugates 144 use_conj = A.is_complex() 145 146 def dot(q, v): return np.dot(np.conjugate( 147 q), v) if use_conj else np.dot(q, v) 148 # The Krylov basis vectors. 149 K_k = np.zeros((len(b), k), 150 dtype=np._complex if use_conj else np.double) 151 # Set the first vector to normalized b 152 K_k[:, 0] = b.unit().entries.ravel() 153 # The upper Hessenberg matrix 154 H = np.zeros((k+1, k), 155 dtype=np._complex if use_conj else np.double) 156 157 # Outer iteration 158 for j in range(k-1): 159 # Arnoldi step 160 v = A.entries @ K_k[:, j] 161 162 # Orthogonalization 163 for i in range(j+1): 164 # Get the Krylov vector to orthogonalize against 165 q_i = K_k[:, i] 166 # Compute the Hessenberg term 167 H[i, j] = dot(q_i, v) 168 # Orthogonalize once 169 v -= H[i, j] * q_i 170 # Orthogonalize again 171 h_tmp = dot(q_i, v) 172 H[i, j] += h_tmp 173 v -= h_tmp * q_i 174 175 # Fill in the sub-diagonal entry 176 H[j+1, j] = np.linalg.norm(v, ord=2) 177 178 if j > 0 and H[j+1, j] <= tolerance: 179 # Exact solution reached, trim the generated subspace 180 K_k = K_k[:, :j-1] 181 H = H[:j, :j-1] 182 # We have learned the rank of A 183 A.set("rank", j-1) 184 break 185 186 # Normalize and store the vector 187 v /= H[j+1, j] 188 K_k[:, j+1] = v 189 190 K_k = Matrix(K_k) 191 K_k.set("rank", K_k.entries.shape[1]) 192 K_k.set("unitary" if use_conj else "orthogonal", True) 193 194 H = Matrix(H) 195 H.set("upper_hessenberg", True) 196 197 return K_k, H
Generate $K_k = \text{span} ( b, Ab, ... A^{k-1}b )$. Uses Modified Gram-Schmidt (twice) full orthogonalization. Optionally configure the tolerance for stopping early.
200def gmres(A: Matrix, b: Vector, k: int, 201 krylov="auto", tolerance=1e-4): 202 """Solve Ax=b for x iteratively using the GMRES algorithm. 203 Stops after `k` iterations, or when the sub-diagonal entry 204 is below the given `tolerance`. 205 206 Optionally specify the method of generating the Krylov 207 subspace by passing a callable to `krylov`. Otherwise, 208 this is derermined automatically from A. The callable 209 should accept the arguments `(A, b, k, tol)` and return the 210 Matrices $K_k$ and $H$. 211 """ 212 if krylov == "auto": 213 krylov = Krylov.arnoldi_mgs 214 215 # Use the specified method to compute the basis 216 K_k, H = krylov(A, b, k, tolerance) 217 218 # Solve the least-squares problem 219 rhs = np.zeros((H.nr, 1)) 220 rhs[0] = b.norm_2() 221 222 y = lstsq(H, Vector(rhs)) 223 x = K_k @ y 224 225 return x, K_k, H
Solve Ax=b for x iteratively using the GMRES algorithm.
Stops after k
iterations, or when the sub-diagonal entry
is below the given tolerance
.
Optionally specify the method of generating the Krylov
subspace by passing a callable to krylov
. Otherwise,
this is derermined automatically from A. The callable
should accept the arguments (A, b, k, tol)
and return the
Matrices $K_k$ and $H$.
228def sgmres(A: Matrix, b: Vector, k: int, r: int, 229 krylov="auto", tolerance=1e-4): 230 """Solve Ax=b for x iteratively using the sGMRES algorithm. 231 Generates a Krylov subspace of dimension `k` via a truncated 232 Arnoldi process that orthogonalizes against `r` previous vectors. 233 The solution is then produced by sketching the matrix using a FFT. 234 235 Optionally specify the method of generating the Krylov 236 subspace by passing a callable to `krylov`. Otherwise, 237 truncated MGS Arnoldi will be used. The callable 238 should accept the arguments `(A, b, k, tol, r)` and return the 239 Matrices $K_k$ and $H$. 240 """ 241 if krylov == "auto": 242 krylov = Krylov.arnoldi_mgs_trunc 243 244 # Use the specified method to compute the basis 245 K_k, H = krylov(A, b, k, tolerance, r) 246 247 # Sketch 248 SA = mathmat_random.FFTSketchedMatrix(A @ K_k) 249 Sb = mathmat_random.FFTSketchedMatrix(b) 250 rows = np.random.choice(len(b), min(len(b), 2*K_k.nc)) 251 SA = SA.entries[rows, :] 252 Sb = Sb.entries[rows, :] 253 254 x = np.linalg.lstsq(SA, Sb, rcond=None)[0] 255 x = K_k.entries @ x 256 257 return Vector(x), K_k, H
Solve Ax=b for x iteratively using the sGMRES algorithm.
Generates a Krylov subspace of dimension k
via a truncated
Arnoldi process that orthogonalizes against r
previous vectors.
The solution is then produced by sketching the matrix using a FFT.
Optionally specify the method of generating the Krylov
subspace by passing a callable to krylov
. Otherwise,
truncated MGS Arnoldi will be used. The callable
should accept the arguments (A, b, k, tol, r)
and return the
Matrices $K_k$ and $H$.
260def invert(A: Matrix, b: Vector): 261 """Solve $Ax=b$ for $x$ by inverting the matrix $A$.""" 262 if A.nr != b.nr: 263 raise MatrixSizeException(A, b, "Linear System Solve") 264 return A.inverse() @ b
Solve $Ax=b$ for $x$ by inverting the matrix $A$.
267def lstsq(A: Matrix, b: Vector): 268 """Find the least-squares, min. norm solution to $Ax=b$.""" 269 if A.nr != b.nr: 270 raise MatrixSizeException(A, b, "Least Squares Solve") 271 if "qr" in A._computed: 272 # Use the QR of A if it exists 273 Q, R = A.qr() 274 return R.lin_solver()(Q.transpose() @ b) 275 # Computing lstsq also reveals some properties of $A$. 276 x, _, rank, sigmas = np.linalg.lstsq(A.entries, b.entries, rcond=None) 277 if "rank" not in A._computed: 278 A.set("rank", rank) 279 if "sigmas" not in A._computed: 280 A.set("sigmas", np.flip(np.sort(sigmas))) 281 return Vector(x)
Find the least-squares, min. norm solution to $Ax=b$.
284def lu(A: Matrix, b: Vector): 285 """Solve $Ax=b$ for $x% by taking an $LU$ factorization of $A$.""" 286 if not A.is_square(): 287 raise MatrixTypeException("LU factorization", "square") 288 if A.is_sparse(): 289 # We obtain the sparse LU solver 290 A.set("lin_solver", lambda v: Vector(lu_factor(A).solve(v.entries))) 291 return Vector(A.lin_solver()(b)) 292 else: 293 P, L, U = lu_factor(A) 294 y = L.lin_solver()(P.inverse() @ b) 295 return U.lin_solver()(y)
Solve $Ax=b$ for $x% by taking an $LU$ factorization of $A$.
298def triangular(A: Matrix, b: Vector, tolerance=1e-10): 299 """Solve $Ax=b$ for $x$ for a triangular $A$.""" 300 if not (A.is_triangular_U(tolerance) or A.is_triangular_L(tolerance)): 301 raise MatrixTypeException("triangular solve", "triangular (any)") 302 return Vector(solve_triangular(A.entries, b.entries, check_finite=False))
Solve $Ax=b$ for $x$ for a triangular $A$.
305def automatic(A: Matrix, b: Vector, get_method=False): 306 """Try to pick an optimal solving strategy based on the properties of $A$. 307 308 Note that computing some properties may be expensive, and therefore using 309 `automatic` may be sub-optimal if used only once. 310 Optionally, if `get_method` is `True`, return the solver method instead of 311 actually solving for $x$. 312 """ 313 if "lin_solver" not in A._computed: 314 # We need to find a suitable linear system solver 315 if "inverse" in A._computed: 316 # An inverse exists and thus a direct solution is available. 317 A.set("lin_solver", lambda v: invert(A, v)) 318 A.lin_solver().name = "auto: invert" 319 elif A.is_triangular_L() or A.is_triangular_U(): 320 # A triangular matrix can be solved using a faster algorithm. 321 A.set("lin_solver", lambda v: triangular(A, v)) 322 A.lin_solver().name = "auto: triangular" 323 elif "qr" in A._computed: 324 # If a QR exists then lstsq will use it 325 A.set("lin_solver", lambda v: lstsq(A, v)) 326 A.lin_solver().name = "auto: qr" 327 elif "plu" in A._computed or A.is_square(): 328 # If a LU (or Cholesky) exists or can be found, use the LU solver 329 A.set("lin_solver", lambda v: lu(A, v)) 330 A.lin_solver().name = "auto: lu" 331 else: 332 # The matrix is nonsquare and has no available factorizations. 333 A.set("lin_solver", lambda v: gmres(A, v, A.nc)[0]) 334 A.lin_solver().name = "auto: gmres" 335 336 if get_method: 337 return A._computed["lin_solver"] 338 return A.lin_solver(b)
Try to pick an optimal solving strategy based on the properties of $A$.
Note that computing some properties may be expensive, and therefore using
automatic
may be sub-optimal if used only once.
Optionally, if get_method
is True
, return the solver method instead of
actually solving for $x$.