Coverage for gpaw/new/pwfd/lbfgs.py: 13%

1import numpy as np 

2from gpaw.new import zips 

3 

4 

5class Vector: 

6 def __init__(self, x_unX): 

7 self.x_unX = x_unX 

8 

9 def copy(self): 

10 return Vector([x_nX.copy() for x_nX in self.x_unX]) 

11 

12 def __iter__(self): 

13 yield from self.x_unX 

14 

15 def zeros(self): 

16 y_unX = [] 

17 for x_nX in self: 

18 y_nX = x_nX.new() 

19 y_nX.data[:] = 0.0 

20 y_unX.append(y_nX) 

21 return Vector(y_unX) 

22 

23 def __mul__(self, other): 

24 y_unX = self.copy() 

25 for y_nX in y_unX: 

26 y_nX.data *= other 

27 return y_unX 

28 

29 def __matmul__(self, other): 

30 z = 0.0 

31 for x_nX, y_nX in zips(self, other): 

32 for x_X, y_X in zips(x_nX, y_nX): 

33 z += x_X.integrate(y_X) 

34 return 2 * z.real 

35 

36 def __sub__(self, other): 

37 x_unX = self.copy() 

38 for x_nX, y_nX in zips(x_unX, other): 

39 x_nX.data -= y_nX.data 

40 return x_unX 

41 

42 

43class LBFGS: 

44 def __init__(self, 

45 memory=3): 

46 self.memory = memory 

47 self.iters = 0 

48 self.kp = None 

49 self.p = None 

50 self.k = None 

51 self.s_k = {i: None for i in range(memory)} 

52 self.y_k = {i: None for i in range(memory)} 

53 self.rho_k = np.zeros(shape=memory) 

54 self.kp = {} 

55 self.p = 0 

56 self.k = 0 

57 self.stable = True 

58 

59 def update(self, 

60 x_k1, 

61 g_k1): 

62 x_k1 = Vector(x_k1) 

63 g_k1 = Vector(g_k1) 

64 self.iters += 1 

65 

66 if self.k == 0: 

67 self.kp[self.k] = self.p 

68 self.x_k = x_k1.copy() 

69 self.g_k = g_k1 

70 self.s_k[self.kp[self.k]] = g_k1.zeros() 

71 self.y_k[self.kp[self.k]] = g_k1.zeros() 

72 self.k += 1 

73 self.p += 1 

74 self.kp[self.k] = self.p 

75 return (g_k1 * -1.0).x_unX 

76 

77 if self.p == self.memory: 

78 self.p = 0 

79 self.kp[self.k] = self.p 

80 

81 s_k = self.s_k 

82 x_k = self.x_k 

83 y_k = self.y_k 

84 g_k = self.g_k 

85 

86 x_k1 = x_k1.copy() 

87 

88 rho_k = self.rho_k 

89 

90 kp = self.kp 

91 k = self.k 

92 m = self.memory 

93 

94 s_k[kp[k]] = x_k1 - x_k 

95 y_k[kp[k]] = g_k1 - g_k 

96 dot_ys = y_k[kp[k]] @ s_k[kp[k]] 

97 

98 if abs(dot_ys) > 1.0e-15: 

99 rho_k[kp[k]] = 1.0 / dot_ys 

100 else: 

101 rho_k[kp[k]] = 1.0e15 

102 

103 if dot_ys < 0.0: 

104 self.stable = False 

105 

106 q = g_k1.copy() 

107 

108 alpha = np.zeros(np.minimum(k + 1, m)) 

109 j = np.maximum(-1, k - m) 

110 

111 for i in range(k, j, -1): 

112 dot_sq = s_k[kp[i]] @ q 

113 alpha[kp[i]] = rho_k[kp[i]] * dot_sq 

114 q = q - y_k[kp[i]] * alpha[kp[i]] 

115 

116 t = k 

117 dot_yy = y_k[kp[t]] @ y_k[kp[t]] 

118 

119 if abs(dot_yy) > 1.0e-15: 

120 r = q * (1.0 / (rho_k[kp[t]] * dot_yy)) 

121 else: 

122 r = q * 1.0e15 

123 

124 for i in range(np.maximum(0, k - m + 1), k + 1): 

125 dot_yr = y_k[kp[i]] @ r 

126 beta = rho_k[kp[i]] * dot_yr 

127 r = r - s_k[kp[i]] * (beta - alpha[kp[i]]) 

128 

129 # save this step: 

130 self.x_k = x_k1.copy() 

131 self.g_k = g_k1.copy() 

132 self.k += 1 

133 self.p += 1 

134 self.kp[self.k] = self.p 

135 

136 return (r * -1.0).x_unX