Coverage for /builds/kinetik161/ase/ase/optimize/bfgs.py: 75.58%

1import warnings 

2from typing import IO, Optional, Union 

3 

4import numpy as np 

5from numpy.linalg import eigh 

6 

7from ase import Atoms 

8from ase.optimize.optimize import Optimizer, UnitCellFilter 

9 

10 

11class BFGS(Optimizer): 

12 # default parameters 

13 defaults = {**Optimizer.defaults, 'alpha': 70.0} 

14 

15 def __init__( 

16 self, 

17 atoms: Atoms, 

18 restart: Optional[str] = None, 

19 logfile: Optional[Union[IO, str]] = '-', 

20 trajectory: Optional[str] = None, 

21 append_trajectory: bool = False, 

22 maxstep: Optional[float] = None, 

23 master: Optional[bool] = None, 

24 alpha: Optional[float] = None, 

25 ): 

26 """BFGS optimizer. 

27 

28 Parameters: 

29 

30 atoms: Atoms object 

31 The Atoms object to relax. 

32 

33 restart: string 

34 Pickle file used to store hessian matrix. If set, file with 

35 such a name will be searched and hessian matrix stored will 

36 be used, if the file exists. 

37 

38 trajectory: string 

39 Pickle file used to store trajectory of atomic movement. 

40 

41 logfile: file object or str 

42 If *logfile* is a string, a file with that name will be opened. 

43 Use '-' for stdout. 

44 

45 maxstep: float 

46 Used to set the maximum distance an atom can move per 

47 iteration (default value is 0.2 Å). 

48 

49 master: boolean 

50 Defaults to None, which causes only rank 0 to save files. If 

51 set to true, this rank will save files. 

52 

53 alpha: float 

54 Initial guess for the Hessian (curvature of energy surface). A 

55 conservative value of 70.0 is the default, but number of needed 

56 steps to converge might be less if a lower value is used. However, 

57 a lower value also means risk of instability. 

58 """ 

59 if maxstep is None: 

60 self.maxstep = self.defaults['maxstep'] 

61 else: 

62 self.maxstep = maxstep 

63 

64 if self.maxstep > 1.0: 

65 warnings.warn('You are using a *very* large value for ' 

66 'the maximum step size: %.1f Å' % self.maxstep) 

67 

68 self.alpha = alpha 

69 if self.alpha is None: 

70 self.alpha = self.defaults['alpha'] 

71 Optimizer.__init__(self, atoms=atoms, restart=restart, 

72 logfile=logfile, trajectory=trajectory, 

73 master=master, append_trajectory=append_trajectory) 

74 

75 def initialize(self): 

76 # initial hessian 

77 self.H0 = np.eye(3 * len(self.optimizable)) * self.alpha 

78 

79 self.H = None 

80 self.pos0 = None 

81 self.forces0 = None 

82 

83 def read(self): 

84 file = self.load() 

85 if len(file) == 5: 

86 (self.H, self.pos0, self.forces0, self.maxstep, 

87 self.atoms.orig_cell) = file 

88 else: 

89 self.H, self.pos0, self.forces0, self.maxstep = file 

90 

91 def step(self, forces=None): 

92 optimizable = self.optimizable 

93 

94 if forces is None: 

95 forces = optimizable.get_forces() 

96 

97 pos = optimizable.get_positions() 

98 dpos, steplengths = self.prepare_step(pos, forces) 

99 dpos = self.determine_step(dpos, steplengths) 

100 optimizable.set_positions(pos + dpos) 

101 if isinstance(self.atoms, UnitCellFilter): 

102 self.dump((self.H, self.pos0, self.forces0, self.maxstep, 

103 self.atoms.orig_cell)) 

104 else: 

105 self.dump((self.H, self.pos0, self.forces0, self.maxstep)) 

106 

107 def prepare_step(self, pos, forces): 

108 forces = forces.reshape(-1) 

109 self.update(pos.flat, forces, self.pos0, self.forces0) 

110 omega, V = eigh(self.H) 

111 

112 # FUTURE: Log this properly 

113 # # check for negative eigenvalues of the hessian 

114 # if any(omega < 0): 

115 # n_negative = len(omega[omega < 0]) 

116 # msg = '\n** BFGS Hessian has {} negative eigenvalues.'.format( 

117 # n_negative 

118 # ) 

119 # print(msg, flush=True) 

120 # if self.logfile is not None: 

121 # self.logfile.write(msg) 

122 # self.logfile.flush() 

123 

124 dpos = np.dot(V, np.dot(forces, V) / np.fabs(omega)).reshape((-1, 3)) 

125 steplengths = (dpos**2).sum(1)**0.5 

126 self.pos0 = pos.flat.copy() 

127 self.forces0 = forces.copy() 

128 return dpos, steplengths 

129 

130 def determine_step(self, dpos, steplengths): 

131 """Determine step to take according to maxstep 

132 

133 Normalize all steps as the largest step. This way 

134 we still move along the eigendirection. 

135 """ 

136 maxsteplength = np.max(steplengths) 

137 if maxsteplength >= self.maxstep: 

138 scale = self.maxstep / maxsteplength 

139 # FUTURE: Log this properly 

140 # msg = '\n** scale step by {:.3f} to be shorter than {}'.format( 

141 # scale, self.maxstep 

142 # ) 

143 # print(msg, flush=True) 

144 

145 dpos *= scale 

146 return dpos 

147 

148 def update(self, pos, forces, pos0, forces0): 

149 if self.H is None: 

150 self.H = self.H0 

151 return 

152 dpos = pos - pos0 

153 

154 if np.abs(dpos).max() < 1e-7: 

155 # Same configuration again (maybe a restart): 

156 return 

157 

158 dforces = forces - forces0 

159 a = np.dot(dpos, dforces) 

160 dg = np.dot(self.H, dpos) 

161 b = np.dot(dpos, dg) 

162 self.H -= np.outer(dforces, dforces) / a + np.outer(dg, dg) / b 

163 

164 def replay_trajectory(self, traj): 

165 """Initialize hessian from old trajectory.""" 

166 if isinstance(traj, str): 

167 from ase.io.trajectory import Trajectory 

168 traj = Trajectory(traj, 'r') 

169 self.H = None 

170 atoms = traj[0] 

171 pos0 = atoms.get_positions().ravel() 

172 forces0 = atoms.get_forces().ravel() 

173 for atoms in traj: 

174 pos = atoms.get_positions().ravel() 

175 forces = atoms.get_forces().ravel() 

176 self.update(pos, forces, pos0, forces0) 

177 pos0 = pos 

178 forces0 = forces 

179 

180 self.pos0 = pos0 

181 self.forces0 = forces0 

182 

183 

184class oldBFGS(BFGS): 

185 def determine_step(self, dpos, steplengths): 

186 """Old BFGS behaviour for scaling step lengths 

187 

188 This keeps the behaviour of truncating individual steps. Some might 

189 depend of this as some absurd kind of stimulated annealing to find the 

190 global minimum. 

191 """ 

192 dpos /= np.maximum(steplengths / self.maxstep, 1.0).reshape(-1, 1) 

193 return dpos