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

1import numpy as np 

2 

3 

4def CG(A, X, B, maxiter=20, tolerance=1.0e-10, verbose=False): 

5 """Solve X*A=B using conjugate gradient method. 

6 

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:: 

10 

11 A(X, Y) 

12 

13 will store ``X*A`` in the output array ``Y``. 

14 

15 On return ``X`` will be the solution to ``X*A=B`` within 

16 ``tolerance``.""" 

17 

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 

42 

43 raise ArithmeticError('Did not converge!')