Spaces:
Runtime error
Runtime error
File size: 2,934 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import numpy as np
def cg(
A,
b,
x0=None,
atol=0.0,
rtol=1e-7,
maxiter=10000,
callback=None,
M=None,
reorthogonalize=False,
):
"""Solves a system of linear equations :math:`Ax=b` using conjugate gradient descent :cite:`hestenes1952methods`
Parameters
----------
A: scipy.sparse.csr_matrix
Square matrix
b: numpy.ndarray
Vector describing the right-hand side of the system
x0: numpy.ndarray
Initialization, if `None` then :code:`x=np.zeros_like(b)`
atol: float
Absolute tolerance. The loop terminates if the :math:`||r||` is smaller than `atol`, where :math:`r` denotes the residual of the current iterate.
rtol: float
Relative tolerance. The loop terminates if :math:`{||r||}/{||b||}` is smaller than `rtol`, where :math:`r` denotes the residual of the current iterate.
callback: function
Function :code:`callback(A, x, b, norm_b, r, norm_r)` called after each iteration, defaults to `None`
M: function or scipy.sparse.csr_matrix
Function that applies the preconditioner to a vector. Alternatively, `M` can be a matrix describing the precondioner.
reorthogonalize: boolean
Wether to apply reorthogonalization of the residuals after each update, defaults to `False`
Returns
-------
x: numpy.ndarray
Solution of the system
Example
-------
>>> from pymatting import *
>>> import numpy as np
>>> A = np.array([[3.0, 1.0], [1.0, 2.0]])
>>> M = jacobi(A)
>>> b = np.array([4.0, 3.0])
>>> cg(A, b, M=M)
array([1., 1.])
"""
if M is None:
def precondition(x):
return x
elif callable(M):
precondition = M
else:
def precondition(x):
return M.dot(x)
x = np.zeros_like(b) if x0 is None else x0.copy()
norm_b = np.linalg.norm(b)
if callable(A):
r = b - A(x)
else:
r = b - A.dot(x)
norm_r = np.linalg.norm(r)
if norm_r < atol or norm_r < rtol * norm_b:
return x
z = precondition(r)
p = z.copy()
rz = np.inner(r, z)
for _ in range(maxiter):
r_old = r.copy()
if callable(A):
Ap = A(p)
else:
Ap = A.dot(p)
alpha = rz / np.inner(p, Ap)
x += alpha * p
r -= alpha * Ap
norm_r = np.linalg.norm(r)
if callback is not None:
callback(A, x, b, norm_b, r, norm_r)
if norm_r < atol or norm_r < rtol * norm_b:
return x
z = precondition(r)
if reorthogonalize:
beta = np.inner(r - r_old, z) / rz
rz = np.inner(r, z)
else:
beta = 1.0 / rz
rz = np.inner(r, z)
beta *= rz
p *= beta
p += z
raise ValueError(
"Conjugate gradient descent did not converge within %d iterations" % maxiter
)
|