Coverage for /builds/alexhroom/ase/ase/optimize/bfgs.py: 78.16%

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.parallel import world 

9from ase.optimize.optimize import Optimizer, UnitCellFilter 

10 

11 

12class BFGS(Optimizer): 

13 # default parameters 

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

15 

16 def __init__( 

17 self, 

18 atoms: Atoms, 

19 restart: Optional[str] = None, 

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

21 trajectory: Optional[str] = None, 

22 append_trajectory: bool = False, 

23 maxstep: Optional[float] = None, 

24 master: Optional[bool] = None, 

25 alpha: Optional[float] = None, 

26 comm=world 

27 ): 

28 """BFGS optimizer. 

29 

30 Parameters: 

31 

32 atoms: Atoms object 

33 The Atoms object to relax. 

34 

35 restart: string 

36 JSON file used to store hessian matrix. If set, file with 

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

38 be used, if the file exists. 

39 

40 trajectory: string 

41 Trajectory file used to store optimisation path. 

42 

43 logfile: file object or str 

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

45 Use '-' for stdout. 

46 

47 maxstep: float 

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

49 iteration (default value is 0.2 Å). 

50 

51 master: boolean 

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

53 set to true, this rank will save files. 

54 

55 comm: Communicator object 

56 Defaults to world. Communicator to handle parallel file reading 

57 and writing, set by ase.parallel.world. 

58 

59 alpha: float 

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

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

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

63 a lower value also means risk of instability. 

64 """ 

65 if maxstep is None: 

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

67 else: 

68 self.maxstep = maxstep 

69 

70 if self.maxstep > 1.0: 

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

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

73 

74 self.alpha = alpha 

75 if self.alpha is None: 

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

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

78 logfile=logfile, trajectory=trajectory, 

79 master=master, append_trajectory=append_trajectory, 

80 comm=comm) 

81 

82 def initialize(self): 

83 # initial hessian 

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

85 

86 self.H = None 

87 self.pos0 = None 

88 self.forces0 = None 

89 

90 def read(self): 

91 file = self.load() 

92 if len(file) == 5: 

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

94 self.atoms.orig_cell) = file 

95 else: 

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

97 

98 def step(self, forces=None): 

99 optimizable = self.optimizable 

100 

101 if forces is None: 

102 forces = optimizable.get_forces() 

103 

104 pos = optimizable.get_positions() 

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

106 dpos = self.determine_step(dpos, steplengths) 

107 optimizable.set_positions(pos + dpos) 

108 if isinstance(self.atoms, UnitCellFilter): 

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

110 self.atoms.orig_cell)) 

111 else: 

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

113 

114 def prepare_step(self, pos, forces): 

115 forces = forces.reshape(-1) 

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

117 omega, V = eigh(self.H) 

118 

119 # FUTURE: Log this properly 

120 # # check for negative eigenvalues of the hessian 

121 # if any(omega < 0): 

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

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

124 # n_negative 

125 # ) 

126 # print(msg, flush=True) 

127 # if self.logfile is not None: 

128 # self.logfile.write(msg) 

129 # self.logfile.flush() 

130 

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

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

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

134 self.forces0 = forces.copy() 

135 return dpos, steplengths 

136 

137 def determine_step(self, dpos, steplengths): 

138 """Determine step to take according to maxstep 

139 

140 Normalize all steps as the largest step. This way 

141 we still move along the direction. 

142 """ 

143 maxsteplength = np.max(steplengths) 

144 if maxsteplength >= self.maxstep: 

145 scale = self.maxstep / maxsteplength 

146 # FUTURE: Log this properly 

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

148 # scale, self.maxstep 

149 # ) 

150 # print(msg, flush=True) 

151 

152 dpos *= scale 

153 return dpos 

154 

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

156 if self.H is None: 

157 self.H = self.H0 

158 return 

159 dpos = pos - pos0 

160 

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

162 # Same configuration again (maybe a restart): 

163 return 

164 

165 dforces = forces - forces0 

166 a = np.dot(dpos, dforces) 

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

168 b = np.dot(dpos, dg) 

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

170 

171 def replay_trajectory(self, traj): 

172 """Initialize hessian from old trajectory.""" 

173 if isinstance(traj, str): 

174 from ase.io.trajectory import Trajectory 

175 traj = Trajectory(traj, 'r') 

176 self.H = None 

177 atoms = traj[0] 

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

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

180 for atoms in traj: 

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

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

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

184 pos0 = pos 

185 forces0 = forces 

186 

187 self.pos0 = pos0 

188 self.forces0 = forces0 

189 

190 

191class oldBFGS(BFGS): 

192 def determine_step(self, dpos, steplengths): 

193 """Old BFGS behaviour for scaling step lengths 

194 

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

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

197 global minimum. 

198 """ 

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

200 return dpos