1import numpy as np 

2 

3from ase.optimize.optimize import Optimizer 

4from ase.constraints import UnitCellFilter 

5import time 

6 

7 

8class PreconFIRE(Optimizer): 

9 

10 def __init__(self, atoms, restart=None, logfile='-', trajectory=None, 

11 dt=0.1, maxmove=0.2, dtmax=1.0, Nmin=5, finc=1.1, fdec=0.5, 

12 astart=0.1, fa=0.99, a=0.1, theta=0.1, master=None, 

13 precon=None, use_armijo=True, variable_cell=False): 

14 """ 

15 Preconditioned version of the FIRE optimizer 

16 

17 Parameters: 

18 

19 atoms: Atoms object 

20 The Atoms object to relax. 

21 

22 restart: string 

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

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

25 be used, if the file exists. 

26 

27 trajectory: string 

28 Pickle file used to store trajectory of atomic movement. 

29 

30 logfile: file object or str 

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

32 Use '-' for stdout. 

33 

34 master: bool 

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

36 set to true, this rank will save files. 

37 

38 variable_cell: bool 

39 If True, wrap atoms in UnitCellFilter to relax cell and positions. 

40 

41 In time this implementation is expected to replace 

42 ase.optimize.fire.FIRE. 

43 """ 

44 if variable_cell: 

45 atoms = UnitCellFilter(atoms) 

46 Optimizer.__init__(self, atoms, restart, logfile, trajectory, master) 

47 

48 self.dt = dt 

49 self.Nsteps = 0 

50 self.maxmove = maxmove 

51 self.dtmax = dtmax 

52 self.Nmin = Nmin 

53 self.finc = finc 

54 self.fdec = fdec 

55 self.astart = astart 

56 self.fa = fa 

57 self.a = a 

58 self.theta = theta 

59 self.precon = precon 

60 self.use_armijo = use_armijo 

61 

62 def initialize(self): 

63 self.v = None 

64 self.skip_flag = False 

65 self.e1 = None 

66 

67 def read(self): 

68 self.v, self.dt = self.load() 

69 

70 def step(self, f=None): 

71 atoms = self.atoms 

72 

73 if f is None: 

74 f = atoms.get_forces() 

75 

76 r = atoms.get_positions() 

77 

78 if self.precon is not None: 

79 # Can this be moved out of the step method? 

80 self.precon.make_precon(atoms) 

81 invP_f = self.precon.solve(f.reshape(-1)).reshape(len(atoms), -1) 

82 

83 if self.v is None: 

84 self.v = np.zeros((len(self.atoms), 3)) 

85 else: 

86 if self.use_armijo: 

87 

88 if self.precon is None: 

89 v_test = self.v + self.dt * f 

90 else: 

91 v_test = self.v + self.dt * invP_f 

92 

93 r_test = r + self.dt * v_test 

94 

95 self.skip_flag = False 

96 func_val = self.func(r_test) 

97 self.e1 = func_val 

98 if (func_val > self.func(r) - 

99 self.theta * self.dt * np.vdot(v_test, f)): 

100 self.v[:] *= 0.0 

101 self.a = self.astart 

102 self.dt *= self.fdec 

103 self.Nsteps = 0 

104 self.skip_flag = True 

105 

106 if not self.skip_flag: 

107 

108 v_f = np.vdot(self.v, f) 

109 if v_f > 0.0: 

110 if self.precon is None: 

111 self.v = (1.0 - self.a) * self.v + self.a * f / \ 

112 np.sqrt(np.vdot(f, f)) * \ 

113 np.sqrt(np.vdot(self.v, self.v)) 

114 else: 

115 self.v = ( 

116 (1.0 - self.a) * self.v + 

117 self.a * 

118 (np.sqrt(self.precon.dot(self.v.reshape(-1), 

119 self.v.reshape(-1))) / 

120 np.sqrt(np.dot(f.reshape(-1), 

121 invP_f.reshape(-1))) * invP_f)) 

122 if self.Nsteps > self.Nmin: 

123 self.dt = min(self.dt * self.finc, self.dtmax) 

124 self.a *= self.fa 

125 self.Nsteps += 1 

126 else: 

127 self.v[:] *= 0.0 

128 self.a = self.astart 

129 self.dt *= self.fdec 

130 self.Nsteps = 0 

131 

132 if self.precon is None: 

133 self.v += self.dt * f 

134 else: 

135 self.v += self.dt * invP_f 

136 dr = self.dt * self.v 

137 normdr = np.sqrt(np.vdot(dr, dr)) 

138 if normdr > self.maxmove: 

139 dr = self.maxmove * dr / normdr 

140 atoms.set_positions(r + dr) 

141 self.dump((self.v, self.dt)) 

142 

143 def func(self, x): 

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

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

146 potl = self.atoms.get_potential_energy() 

147 return potl 

148 

149 def run(self, fmax=0.05, steps=100000000, smax=None): 

150 if smax is None: 

151 smax = fmax 

152 self.smax = smax 

153 return Optimizer.run(self, fmax, steps) 

154 

155 def converged(self, forces=None): 

156 """Did the optimization converge?""" 

157 if forces is None: 

158 forces = self.atoms.get_forces() 

159 if isinstance(self.atoms, UnitCellFilter): 

160 natoms = len(self.atoms.atoms) 

161 forces, stress = forces[:natoms], self.atoms.stress 

162 fmax_sq = (forces**2).sum(axis=1).max() 

163 smax_sq = (stress**2).max() 

164 return (fmax_sq < self.fmax**2 and smax_sq < self.smax**2) 

165 else: 

166 fmax_sq = (forces**2).sum(axis=1).max() 

167 return fmax_sq < self.fmax**2 

168 

169 def log(self, forces=None): 

170 if forces is None: 

171 forces = self.atoms.get_forces() 

172 if isinstance(self.atoms, UnitCellFilter): 

173 natoms = len(self.atoms.atoms) 

174 forces, stress = forces[:natoms], self.atoms.stress 

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

176 smax = np.sqrt((stress**2).max()) 

177 else: 

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

179 if self.e1 is not None: 

180 # reuse energy at end of line search to avoid extra call 

181 e = self.e1 

182 else: 

183 e = self.atoms.get_potential_energy() 

184 T = time.localtime() 

185 if self.logfile is not None: 

186 name = self.__class__.__name__ 

187 if isinstance(self.atoms, UnitCellFilter): 

188 self.logfile.write( 

189 '%s: %3d %02d:%02d:%02d %15.6f %12.4f %12.4f\n' % 

190 (name, self.nsteps, T[3], T[4], T[5], e, fmax, smax)) 

191 

192 else: 

193 self.logfile.write( 

194 '%s: %3d %02d:%02d:%02d %15.6f %12.4f\n' % 

195 (name, self.nsteps, T[3], T[4], T[5], e, fmax)) 

196 self.logfile.flush()