Coverage for gpaw/utilities/cg.py: 96%
26 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1import numpy as np
4def CG(A, X, B, maxiter=20, tolerance=1.0e-10, verbose=False):
5 """Solve X*A=B using conjugate gradient method.
7 ``X`` and ``B`` are ``ndarrays```of shape ``(m, nx, ny, nz)``
8 coresponding to matrices of size ``m*n`` (``n=nx*ny*nz``) and
9 ``A`` is a callable representing an ``n*n`` matrix::
11 A(X, Y)
13 will store ``X*A`` in the output array ``Y``.
15 On return ``X`` will be the solution to ``X*A=B`` within
16 ``tolerance``."""
18 m = len(X)
19 shape = (m, 1, 1, 1)
20 R = np.empty(X.shape, X.dtype.char)
21 Q = np.empty(X.shape, X.dtype.char)
22 A(X, R)
23 R -= B
24 P = R.copy()
25 c1 = A.sum(np.reshape([abs(np.vdot(r, r)) for r in R], shape))
26 for i in range(maxiter):
27 error = sum(c1.ravel())
28 if verbose:
29 print('CG-%d: %e' % (i, error))
30 if error < tolerance:
31 return i, error
32 A(P, Q)
33 alpha = c1 / A.sum(np.reshape([np.vdot(q, p)
34 for p, q in zip(P, Q)], shape))
35 X -= alpha * P
36 R -= alpha * Q
37 c0 = c1
38 c1 = A.sum(np.reshape([abs(np.vdot(r, r)) for r in R], shape))
39 beta = c1 / c0
40 P *= beta
41 P += R
43 raise ArithmeticError('Did not converge!')