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

1# fmt: off 

2 

3import warnings 

4from pathlib import Path 

5from typing import IO, Optional, Union 

6 

7import numpy as np 

8from numpy.linalg import eigh 

9 

10from ase import Atoms 

11from ase.optimize.optimize import Optimizer, UnitCellFilter 

12 

13 

14class BFGS(Optimizer): 

15 # default parameters 

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

17 

18 def __init__( 

19 self, 

20 atoms: Atoms, 

21 restart: Optional[str] = None, 

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

23 trajectory: Optional[Union[str, Path]] = None, 

24 append_trajectory: bool = False, 

25 maxstep: Optional[float] = None, 

26 alpha: Optional[float] = None, 

27 **kwargs, 

28 ): 

29 """BFGS optimizer. 

30 

31 Parameters 

32 ---------- 

33 atoms: :class:`~ase.Atoms` 

34 The Atoms object to relax. 

35 

36 restart: str 

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

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

39 be used, if the file exists. 

40 

41 trajectory: str or Path 

42 Trajectory file used to store optimisation path. 

43 

44 logfile: file object, Path, or str 

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

46 Use '-' for stdout. 

47 

48 maxstep: float 

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

50 iteration (default value is 0.2 Å). 

51 

52 alpha: float 

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

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

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

56 a lower value also means risk of instability. 

57 

58 kwargs : dict, optional 

59 Extra arguments passed to 

60 :class:`~ase.optimize.optimize.Optimizer`. 

61 

62 """ 

63 if maxstep is None: 

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

65 else: 

66 self.maxstep = maxstep 

67 

68 if self.maxstep > 1.0: 

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

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

71 

72 self.alpha = alpha 

73 if self.alpha is None: 

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

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

76 logfile=logfile, trajectory=trajectory, 

77 append_trajectory=append_trajectory, 

78 **kwargs) 

79 

80 def initialize(self): 

81 # initial hessian 

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

83 

84 self.H = None 

85 self.pos0 = None 

86 self.forces0 = None 

87 

88 def read(self): 

89 file = self.load() 

90 if len(file) == 5: 

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

92 self.atoms.orig_cell) = file 

93 else: 

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

95 

96 def step(self, forces=None): 

97 optimizable = self.optimizable 

98 

99 if forces is None: 

100 forces = optimizable.get_forces() 

101 

102 pos = optimizable.get_positions() 

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

104 dpos = self.determine_step(dpos, steplengths) 

105 optimizable.set_positions(pos + dpos) 

106 if isinstance(self.atoms, UnitCellFilter): 

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

108 self.atoms.orig_cell)) 

109 else: 

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

111 

112 def prepare_step(self, pos, forces): 

113 forces = forces.reshape(-1) 

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

115 omega, V = eigh(self.H) 

116 

117 # FUTURE: Log this properly 

118 # # check for negative eigenvalues of the hessian 

119 # if any(omega < 0): 

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

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

122 # n_negative 

123 # ) 

124 # print(msg, flush=True) 

125 # if self.logfile is not None: 

126 # self.logfile.write(msg) 

127 # self.logfile.flush() 

128 

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

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

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

132 self.forces0 = forces.copy() 

133 return dpos, steplengths 

134 

135 def determine_step(self, dpos, steplengths): 

136 """Determine step to take according to maxstep 

137 

138 Normalize all steps as the largest step. This way 

139 we still move along the direction. 

140 """ 

141 maxsteplength = np.max(steplengths) 

142 if maxsteplength >= self.maxstep: 

143 scale = self.maxstep / maxsteplength 

144 # FUTURE: Log this properly 

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

146 # scale, self.maxstep 

147 # ) 

148 # print(msg, flush=True) 

149 

150 dpos *= scale 

151 return dpos 

152 

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

154 if self.H is None: 

155 self.H = self.H0 

156 return 

157 dpos = pos - pos0 

158 

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

160 # Same configuration again (maybe a restart): 

161 return 

162 

163 dforces = forces - forces0 

164 a = np.dot(dpos, dforces) 

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

166 b = np.dot(dpos, dg) 

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

168 

169 def replay_trajectory(self, traj): 

170 """Initialize hessian from old trajectory.""" 

171 if isinstance(traj, str): 

172 from ase.io.trajectory import Trajectory 

173 traj = Trajectory(traj, 'r') 

174 self.H = None 

175 atoms = traj[0] 

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

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

178 for atoms in traj: 

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

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

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

182 pos0 = pos 

183 forces0 = forces 

184 

185 self.pos0 = pos0 

186 self.forces0 = forces0 

187 

188 

189class oldBFGS(BFGS): 

190 def determine_step(self, dpos, steplengths): 

191 """Old BFGS behaviour for scaling step lengths 

192 

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

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

195 global minimum. 

196 """ 

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

198 return dpos