Coverage for gpaw/preconditioner.py: 95%

1import gpaw.cgpaw as cgpaw 

2import numpy as np 

3 

4from gpaw import debug 

5from gpaw.fd_operators import Laplace 

6from gpaw.transformers import Transformer 

7 

8 

9class Preconditioner: 

10 def __init__(self, gd0, kin0, dtype=float, block=1, xp=np): 

11 gd1 = gd0.coarsen() 

12 gd2 = gd1.coarsen() 

13 self.kin0 = kin0 

14 self.kin1 = Laplace(gd1, -0.5, 1, dtype, xp=xp) 

15 self.kin2 = Laplace(gd2, -0.5, 1, dtype, xp=xp) 

16 self.scratch0 = gd0.zeros((2, block), dtype, False, xp=xp) 

17 self.scratch1 = gd1.zeros((3, block), dtype, False, xp=xp) 

18 self.scratch2 = gd2.zeros((3, block), dtype, False, xp=xp) 

19 self.step = 0.66666666 / kin0.get_diagonal_element() 

20 

21 self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype, xp=xp) 

22 self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype, xp=xp) 

23 self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype, xp=xp) 

24 self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype, xp=xp) 

25 self.restrictor0 = self.restrictor_object0.apply 

26 self.restrictor1 = self.restrictor_object1.apply 

27 self.interpolator2 = self.interpolator_object2.apply 

28 self.interpolator1 = self.interpolator_object1.apply 

29 self.use_c_precond = True 

30 self.xp = xp 

31 

32 def calculate_kinetic_energy(self, psit_xG, kpt): 

33 return None 

34 

35 def __call__(self, residuals, kpt, ekin=None, out=None): 

36 if residuals.ndim == 3: 

37 if out is None: 

38 return self.__call__(residuals[np.newaxis], kpt)[0] 

39 return self.__call__(residuals[np.newaxis], kpt, 

40 out=out[np.newaxis])[0] 

41 

42 nb = len(residuals) # number of bands 

43 nb0 = self.scratch0.shape[1] 

44 if nb > nb0: 

45 assert out is not None 

46 for n1 in range(0, nb, nb0): 

47 self(residuals[n1:n1 + nb0], kpt, out=out[n1:n1 + nb0]) 

48 return out 

49 

50 phases = kpt.phase_cd 

51 step = self.step 

52 if out is None: 

53 d0, q0 = self.scratch0[:, :nb] 

54 else: 

55 d0 = out 

56 q0 = self.scratch0[0, :nb] 

57 r1, d1, q1 = self.scratch1[:, :nb] 

58 r2, d2, q2 = self.scratch2[:, :nb] 

59 if self.use_c_precond and self.xp is np: 

60 transformers = [self.restrictor_object0.transformer, 

61 self.restrictor_object1.transformer, 

62 self.interpolator_object1.transformer, 

63 self.interpolator_object2.transformer] 

64 if debug: 

65 # Unwrap wrapper: 

66 transformers = [getattr(t, 'transformer', t) 

67 for t in transformers] 

68 cgpaw.fd_precond(*transformers, 

69 self.kin0.operator, self.kin1.operator, 

70 self.kin2.operator, 

71 d0, q0, r1, d1, q1, r2, d2, q2, 

72 residuals, -residuals, step, phases) 

73 return d0 

74 self.restrictor0(-residuals, r1, phases) 

75 d1[:] = 4 * step * r1 

76 self.kin1.apply(d1, q1, phases) 

77 q1 -= r1 

78 self.restrictor1(q1, r2, phases) 

79 d2 = 16 * step * r2 

80 self.kin2.apply(d2, q2, phases) 

81 q2 -= r2 

82 d2 -= 16 * step * q2 

83 self.interpolator2(d2, q1, phases) 

84 d1 -= q1 

85 self.kin1.apply(d1, q1, phases) 

86 q1 -= r1 

87 d1 -= 4 * step * q1 

88 self.interpolator1(-d1, d0, phases) 

89 self.kin0.apply(d0, q0, phases) 

90 q0 -= residuals 

91 d0 -= step * q0 

92 d0 *= -1.0 

93 return d0