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)
class Krylov:
 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.

@staticmethod
def arnoldi_mgs( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, k: int, tolerance=0.0001, **_):
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.

@staticmethod
def arnoldi_mgs_trunc( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, k: int, tolerance=0.0001, r=10):
 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.

@staticmethod
def arnoldi_mgs_2x( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, k: int, tolerance=0.0001, **_):
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.

def gmres( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, k: int, krylov='auto', tolerance=0.0001):
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$.

def sgmres( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, k: int, r: int, krylov='auto', tolerance=0.0001):
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$.

def invert(A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector):
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$.

def lstsq(A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector):
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$.

def lu(A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector):
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$.

def triangular( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, tolerance=1e-10):
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$.

def automatic( A: mathmat.mathmat.Matrix, b: mathmat.mathmat.Vector, get_method=False):
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$.