Coverage for /builds/debichem-team/python-ase/ase/optimize/bfgs.py: 77.91%

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 alpha: Optional[float] = None, 

24 **kwargs, 

25 ): 

26 """BFGS optimizer. 

27 

28 Parameters 

29 ---------- 

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

31 The Atoms object to relax. 

32 

33 restart: str 

34 JSON 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: str 

39 Trajectory file used to store optimisation path. 

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 alpha: float 

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

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

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

53 a lower value also means risk of instability. 

54 

55 kwargs : dict, optional 

56 Extra arguments passed to 

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

58 

59 """ 

60 if maxstep is None: 

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

62 else: 

63 self.maxstep = maxstep 

64 

65 if self.maxstep > 1.0: 

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

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

68 

69 self.alpha = alpha 

70 if self.alpha is None: 

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

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

73 logfile=logfile, trajectory=trajectory, 

74 append_trajectory=append_trajectory, 

75 **kwargs) 

76 

77 def initialize(self): 

78 # initial hessian 

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

80 

81 self.H = None 

82 self.pos0 = None 

83 self.forces0 = None 

84 

85 def read(self): 

86 file = self.load() 

87 if len(file) == 5: 

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

89 self.atoms.orig_cell) = file 

90 else: 

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

92 

93 def step(self, forces=None): 

94 optimizable = self.optimizable 

95 

96 if forces is None: 

97 forces = optimizable.get_forces() 

98 

99 pos = optimizable.get_positions() 

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

101 dpos = self.determine_step(dpos, steplengths) 

102 optimizable.set_positions(pos + dpos) 

103 if isinstance(self.atoms, UnitCellFilter): 

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

105 self.atoms.orig_cell)) 

106 else: 

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

108 

109 def prepare_step(self, pos, forces): 

110 forces = forces.reshape(-1) 

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

112 omega, V = eigh(self.H) 

113 

114 # FUTURE: Log this properly 

115 # # check for negative eigenvalues of the hessian 

116 # if any(omega < 0): 

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

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

119 # n_negative 

120 # ) 

121 # print(msg, flush=True) 

122 # if self.logfile is not None: 

123 # self.logfile.write(msg) 

124 # self.logfile.flush() 

125 

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

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

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

129 self.forces0 = forces.copy() 

130 return dpos, steplengths 

131 

132 def determine_step(self, dpos, steplengths): 

133 """Determine step to take according to maxstep 

134 

135 Normalize all steps as the largest step. This way 

136 we still move along the direction. 

137 """ 

138 maxsteplength = np.max(steplengths) 

139 if maxsteplength >= self.maxstep: 

140 scale = self.maxstep / maxsteplength 

141 # FUTURE: Log this properly 

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

143 # scale, self.maxstep 

144 # ) 

145 # print(msg, flush=True) 

146 

147 dpos *= scale 

148 return dpos 

149 

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

151 if self.H is None: 

152 self.H = self.H0 

153 return 

154 dpos = pos - pos0 

155 

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

157 # Same configuration again (maybe a restart): 

158 return 

159 

160 dforces = forces - forces0 

161 a = np.dot(dpos, dforces) 

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

163 b = np.dot(dpos, dg) 

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

165 

166 def replay_trajectory(self, traj): 

167 """Initialize hessian from old trajectory.""" 

168 if isinstance(traj, str): 

169 from ase.io.trajectory import Trajectory 

170 traj = Trajectory(traj, 'r') 

171 self.H = None 

172 atoms = traj[0] 

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

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

175 for atoms in traj: 

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

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

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

179 pos0 = pos 

180 forces0 = forces 

181 

182 self.pos0 = pos0 

183 self.forces0 = forces0 

184 

185 

186class oldBFGS(BFGS): 

187 def determine_step(self, dpos, steplengths): 

188 """Old BFGS behaviour for scaling step lengths 

189 

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

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

192 global minimum. 

193 """ 

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

195 return dpos