Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# ******NOTICE*************** 

2# optimize.py module by Travis E. Oliphant 

3# 

4# You may copy and use this module as you see fit with no 

5# guarantee implied provided you keep this notice in all copies. 

6# *****END NOTICE************ 

7 

8import time 

9import numpy as np 

10from numpy import eye, absolute, sqrt, isinf 

11from ase.utils.linesearch import LineSearch 

12from ase.optimize.optimize import Optimizer 

13 

14# These have been copied from Numeric's MLab.py 

15# I don't think they made the transition to scipy_core 

16 

17# Modified from scipy_optimize 

18abs = absolute 

19pymin = min 

20pymax = max 

21__version__ = '0.1' 

22 

23 

24class BFGSLineSearch(Optimizer): 

25 def __init__(self, atoms, restart=None, logfile='-', maxstep=None, 

26 trajectory=None, c1=0.23, c2=0.46, alpha=10.0, stpmax=50.0, 

27 master=None, force_consistent=None): 

28 """Optimize atomic positions in the BFGSLineSearch algorithm, which 

29 uses both forces and potential energy information. 

30 

31 Parameters: 

32 

33 atoms: Atoms object 

34 The Atoms object to relax. 

35 

36 restart: string 

37 Pickle 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: string 

42 Pickle file used to store trajectory of atomic movement. 

43 

44 maxstep: float 

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

46 iteration (default value is 0.2 Angstroms). 

47 

48 logfile: file object or str 

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

50 Use '-' for stdout. 

51 

52 master: boolean 

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

54 set to true, this rank will save files. 

55 

56 force_consistent: boolean or None 

57 Use force-consistent energy calls (as opposed to the energy 

58 extrapolated to 0 K). By default (force_consistent=None) uses 

59 force-consistent energies if available in the calculator, but 

60 falls back to force_consistent=False if not. 

61 """ 

62 if maxstep is None: 

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

64 else: 

65 self.maxstep = maxstep 

66 self.stpmax = stpmax 

67 self.alpha = alpha 

68 self.H = None 

69 self.c1 = c1 

70 self.c2 = c2 

71 self.force_calls = 0 

72 self.function_calls = 0 

73 self.r0 = None 

74 self.g0 = None 

75 self.e0 = None 

76 self.load_restart = False 

77 self.task = 'START' 

78 self.rep_count = 0 

79 self.p = None 

80 self.alpha_k = None 

81 self.no_update = False 

82 self.replay = False 

83 

84 Optimizer.__init__(self, atoms, restart, logfile, trajectory, 

85 master, force_consistent=force_consistent) 

86 

87 def read(self): 

88 self.r0, self.g0, self.e0, self.task, self.H = self.load() 

89 self.load_restart = True 

90 

91 def reset(self): 

92 self.H = None 

93 self.r0 = None 

94 self.g0 = None 

95 self.e0 = None 

96 self.rep_count = 0 

97 

98 def step(self, forces=None): 

99 atoms = self.atoms 

100 

101 if forces is None: 

102 forces = atoms.get_forces() 

103 

104 from ase.neb import NEB 

105 if isinstance(atoms, NEB): 

106 raise TypeError('NEB calculations cannot use the BFGSLineSearch' 

107 ' optimizer. Use BFGS or another optimizer.') 

108 r = atoms.get_positions() 

109 r = r.reshape(-1) 

110 g = -forces.reshape(-1) / self.alpha 

111 p0 = self.p 

112 self.update(r, g, self.r0, self.g0, p0) 

113 # o,v = np.linalg.eigh(self.B) 

114 e = self.func(r) 

115 

116 self.p = -np.dot(self.H, g) 

117 p_size = np.sqrt((self.p**2).sum()) 

118 if p_size <= np.sqrt(len(atoms) * 1e-10): 

119 self.p /= (p_size / np.sqrt(len(atoms) * 1e-10)) 

120 ls = LineSearch() 

121 self.alpha_k, e, self.e0, self.no_update = \ 

122 ls._line_search(self.func, self.fprime, r, self.p, g, e, self.e0, 

123 maxstep=self.maxstep, c1=self.c1, 

124 c2=self.c2, stpmax=self.stpmax) 

125 if self.alpha_k is None: 

126 raise RuntimeError("LineSearch failed!") 

127 

128 dr = self.alpha_k * self.p 

129 atoms.set_positions((r + dr).reshape(len(atoms), -1)) 

130 self.r0 = r 

131 self.g0 = g 

132 self.dump((self.r0, self.g0, self.e0, self.task, self.H)) 

133 

134 def update(self, r, g, r0, g0, p0): 

135 self.I = eye(len(self.atoms) * 3, dtype=int) 

136 if self.H is None: 

137 self.H = eye(3 * len(self.atoms)) 

138 # self.B = np.linalg.inv(self.H) 

139 return 

140 else: 

141 dr = r - r0 

142 dg = g - g0 

143 # self.alpha_k can be None!!! 

144 if not (((self.alpha_k or 0) > 0 and 

145 abs(np.dot(g, p0)) - abs(np.dot(g0, p0)) < 0) or 

146 self.replay): 

147 return 

148 if self.no_update is True: 

149 print('skip update') 

150 return 

151 

152 try: # this was handled in numeric, let it remain for more safety 

153 rhok = 1.0 / (np.dot(dg, dr)) 

154 except ZeroDivisionError: 

155 rhok = 1000.0 

156 print("Divide-by-zero encountered: rhok assumed large") 

157 if isinf(rhok): # this is patch for np 

158 rhok = 1000.0 

159 print("Divide-by-zero encountered: rhok assumed large") 

160 A1 = self.I - dr[:, np.newaxis] * dg[np.newaxis, :] * rhok 

161 A2 = self.I - dg[:, np.newaxis] * dr[np.newaxis, :] * rhok 

162 self.H = (np.dot(A1, np.dot(self.H, A2)) + 

163 rhok * dr[:, np.newaxis] * dr[np.newaxis, :]) 

164 # self.B = np.linalg.inv(self.H) 

165 

166 def func(self, x): 

167 """Objective function for use of the optimizers""" 

168 self.atoms.set_positions(x.reshape(-1, 3)) 

169 self.function_calls += 1 

170 # Scale the problem as SciPy uses I as initial Hessian. 

171 return (self.atoms.get_potential_energy( 

172 force_consistent=self.force_consistent) / self.alpha) 

173 

174 def fprime(self, x): 

175 """Gradient of the objective function for use of the optimizers""" 

176 self.atoms.set_positions(x.reshape(-1, 3)) 

177 self.force_calls += 1 

178 # Remember that forces are minus the gradient! 

179 # Scale the problem as SciPy uses I as initial Hessian. 

180 forces = self.atoms.get_forces().reshape(-1) 

181 return - forces / self.alpha 

182 

183 def replay_trajectory(self, traj): 

184 """Initialize hessian from old trajectory.""" 

185 self.replay = True 

186 from ase.utils import IOContext 

187 

188 with IOContext() as files: 

189 if isinstance(traj, str): 

190 from ase.io.trajectory import Trajectory 

191 traj = files.closelater(Trajectory(traj, mode='r')) 

192 

193 r0 = None 

194 g0 = None 

195 for i in range(0, len(traj) - 1): 

196 r = traj[i].get_positions().ravel() 

197 g = - traj[i].get_forces().ravel() / self.alpha 

198 self.update(r, g, r0, g0, self.p) 

199 self.p = -np.dot(self.H, g) 

200 r0 = r.copy() 

201 g0 = g.copy() 

202 self.r0 = r0 

203 self.g0 = g0 

204 

205 def log(self, forces=None): 

206 if self.logfile is None: 

207 return 

208 if forces is None: 

209 forces = self.atoms.get_forces() 

210 fmax = sqrt((forces**2).sum(axis=1).max()) 

211 e = self.atoms.get_potential_energy( 

212 force_consistent=self.force_consistent) 

213 T = time.localtime() 

214 name = self.__class__.__name__ 

215 w = self.logfile.write 

216 if self.nsteps == 0: 

217 w('%s %4s[%3s] %8s %15s %12s\n' % 

218 (' ' * len(name), 'Step', 'FC', 'Time', 'Energy', 'fmax')) 

219 if self.force_consistent: 

220 w('*Force-consistent energies used in optimization.\n') 

221 w('%s: %3d[%3d] %02d:%02d:%02d %15.6f%1s %12.4f\n' 

222 % (name, self.nsteps, self.force_calls, T[3], T[4], T[5], e, 

223 {1: '*', 0: ''}[self.force_consistent], fmax)) 

224 self.logfile.flush() 

225 

226 

227def wrap_function(function, args): 

228 ncalls = [0] 

229 

230 def function_wrapper(x): 

231 ncalls[0] += 1 

232 return function(x, *args) 

233 return ncalls, function_wrapper