Coverage for /builds/kinetik161/ase/ase/optimize/precon/lbfgs.py: 86.24%

1import time

2import warnings

3from math import sqrt

5import numpy as np

7from ase.filters import UnitCellFilter

8from ase.optimize.optimize import Optimizer

9from ase.optimize.precon.precon import make_precon

10from ase.utils.linesearch import LineSearch

11from ase.utils.linesearcharmijo import LineSearchArmijo

14class PreconLBFGS(Optimizer):

15 """Preconditioned version of the Limited memory BFGS optimizer, to

16 be used as a drop-in replacement for ase.optimize.lbfgs.LBFGS for systems

17 where a good preconditioner is available.

19 In the standard bfgs and lbfgs algorithms, the inverse of Hessian matrix

20 is a (usually fixed) diagonal matrix. By contrast, PreconLBFGS,

21 updates the hessian after each step with a general "preconditioner".

22 By default, the ase.optimize.precon.Exp preconditioner is applied.

23 This preconditioner is well-suited for large condensed phase structures,

24 in particular crystalline. For systems outside this category,

25 PreconLBFGS with Exp preconditioner may yield unpredictable results.

27 In time this implementation is expected to replace

28 ase.optimize.lbfgs.LBFGS.

30 See this article for full details: D. Packwood, J. R. Kermode, L. Mones,

31 N. Bernstein, J. Woolley, N. Gould, C. Ortner, and G. Csanyi, A universal

32 preconditioner for simulating condensed phase materials

33 J. Chem. Phys. 144, 164109 (2016), DOI: https://doi.org/10.1063/1.4947024

34 """

36 # CO : added parameters rigid_units and rotation_factors

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

38 maxstep=None, memory=100, damping=1.0, alpha=70.0,

39 master=None, precon='auto', variable_cell=False,

40 use_armijo=True, c1=0.23, c2=0.46, a_min=None,

41 rigid_units=None, rotation_factors=None, Hinv=None):

42 """Parameters:

44 atoms: Atoms object

45 The Atoms object to relax.

47 restart: string

48 Pickle file used to store vectors for updating the inverse of

49 Hessian matrix. If set, file with such a name will be searched

50 and information stored will be used, if the file exists.

52 logfile: file object or str

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

54 Use '-' for stdout.

56 trajectory: string

57 Pickle file used to store trajectory of atomic movement.

59 maxstep: float

60 How far is a single atom allowed to move. This is useful for DFT

61 calculations where wavefunctions can be reused if steps are small.

62 Default is 0.04 Angstrom.

64 memory: int

65 Number of steps to be stored. Default value is 100. Three numpy

66 arrays of this length containing floats are stored.

68 damping: float

69 The calculated step is multiplied with this number before added to

70 the positions.

72 alpha: float

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

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

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

76 a lower value also means risk of instability.

78 master: boolean

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

80 set to true, this rank will save files.

82 precon: ase.optimize.precon.Precon instance or compatible.

83 Apply the given preconditioner during optimization. Defaults to

84 'auto', which will choose the `Exp` preconditioner unless the system

85 is too small (< 100 atoms) in which case a standard LBFGS fallback

86 is used. To enforce use of the `Exp` preconditioner, use `precon =

87 'Exp'`. Other options include 'C1', 'Pfrommer' and 'FF' - see the

88 corresponding classes in the `ase.optimize.precon` module for more

89 details. Pass precon=None or precon='ID' to disable preconditioning.

91 use_armijo: boolean

92 Enforce only the Armijo condition of sufficient decrease of

93 of the energy, and not the second Wolff condition for the forces.

94 Often significantly faster than full Wolff linesearch.

95 Defaults to True.

97 c1: float

98 c1 parameter for the line search. Default is c1=0.23.

100 c2: float

101 c2 parameter for the line search. Default is c2=0.46.

102

103 a_min: float

104 minimal value for the line search step parameter. Default is

105 a_min=1e-8 (use_armijo=False) or 1e-10 (use_armijo=True).

106 Higher values can be useful to avoid performing many

107 line searches for comparatively small changes in geometry.

108

109 variable_cell: bool

110 If True, wrap atoms in UnitCellFilter to

111 relax both postions and cell. Default is False.

112

113 rigid_units: each I = rigid_units[i] is a list of indices, which

114 describes a subsystem of atoms that forms a (near-)rigid unit

115 If rigid_units is not None, then special search-paths are

116 are created to take the rigidness into account

117

118 rotation_factors: list of scalars; acceleration factors deteriming

119 the rate of rotation as opposed to the rate of stretch in the

120 rigid units

121 """

122 if variable_cell:

123 atoms = UnitCellFilter(atoms)

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

125

126 self._actual_atoms = atoms

127

128 # default preconditioner

129 # TODO: introduce a heuristic for different choices of preconditioners

130 if precon == 'auto':

131 if len(atoms) < 100:

132 precon = None

133 warnings.warn('The system is likely too small to benefit from '

134 'the standard preconditioner, hence it is '

135 'disabled. To re-enable preconditioning, call '

136 '`PreconLBFGS` by explicitly providing the '

137 'kwarg `precon`')

138 else:

139 precon = 'Exp'

140

141 if maxstep is not None:

142 if maxstep > 1.0:

143 raise ValueError('You are using a much too large value for ' +

144 'the maximum step size: %.1f Angstrom' %

145 maxstep)

146 self.maxstep = maxstep

147 else:

148 self.maxstep = 0.04

149

150 self.memory = memory

151 self.H0 = 1. / alpha # Initial approximation of inverse Hessian

152 # 1./70. is to emulate the behaviour of BFGS

153 # Note that this is never changed!

154 self.Hinv = Hinv

155 self.damping = damping

156 self.p = None

157

158 # construct preconditioner if passed as a string

159 self.precon = make_precon(precon)

160 self.use_armijo = use_armijo

161 self.c1 = c1

162 self.c2 = c2

163 self.a_min = a_min

164 if self.a_min is None:

165 self.a_min = 1e-10 if use_armijo else 1e-8

166

167 # CO

168 self.rigid_units = rigid_units

169 self.rotation_factors = rotation_factors

170

171 def reset_hessian(self):

172 """

173 Throw away history of the Hessian

174 """

175 self._just_reset_hessian = True

176 self.s = []

177 self.y = []

178 self.rho = [] # Store also rho, to avoid calculating the dot product

179 # again and again

180

181 def initialize(self):

182 """Initialize everything so no checks have to be done in step"""

183 self.iteration = 0

184 self.reset_hessian()

185 self.r0 = None

186 self.f0 = None

187 self.e0 = None

188 self.e1 = None

189 self.task = 'START'

190 self.load_restart = False

191

192 def read(self):

193 """Load saved arrays to reconstruct the Hessian"""

194 self.iteration, self.s, self.y, self.rho, \

195 self.r0, self.f0, self.e0, self.task = self.load()

196 self.load_restart = True

197

198 def step(self, f=None):

199 """Take a single step

200

201 Use the given forces, update the history and calculate the next step --

202 then take it"""

203 r = self._actual_atoms.get_positions()

204

205 if f is None:

206 f = self._actual_atoms.get_forces()

207

208 previously_reset_hessian = self._just_reset_hessian

209 self.update(r, f, self.r0, self.f0)

210

211 s = self.s

212 y = self.y

213 rho = self.rho

214 H0 = self.H0

215

216 loopmax = np.min([self.memory, len(self.y)])

217 a = np.empty((loopmax,), dtype=np.float64)

218

219 # The algorithm itself:

220 q = -f.reshape(-1)

221 for i in range(loopmax - 1, -1, -1):

222 a[i] = rho[i] * np.dot(s[i], q)

223 q -= a[i] * y[i]

224

225 if self.precon is None:

226 if self.Hinv is not None:

227 z = np.dot(self.Hinv, q)

228 else:

229 z = H0 * q

230 else:

231 self.precon.make_precon(self._actual_atoms)

232 z = self.precon.solve(q)

233

234 for i in range(loopmax):

235 b = rho[i] * np.dot(y[i], z)

236 z += s[i] * (a[i] - b)

237

238 self.p = - z.reshape((-1, 3))

239 ###

240

241 g = -f

242 if self.e1 is not None:

243 e = self.e1

244 else:

245 e = self.func(r)

246 self.line_search(r, g, e, previously_reset_hessian)

247 dr = (self.alpha_k * self.p).reshape(len(self._actual_atoms), -1)

248

249 if self.alpha_k != 0.0:

250 self._actual_atoms.set_positions(r + dr)

251

252 self.iteration += 1

253 self.r0 = r

254 self.f0 = -g

255 self.dump((self.iteration, self.s, self.y,

256 self.rho, self.r0, self.f0, self.e0, self.task))

257

258 def update(self, r, f, r0, f0):

259 """Update everything that is kept in memory

260

261 This function is mostly here to allow for replay_trajectory.

262 """

263 if not self._just_reset_hessian:

264 s0 = r.reshape(-1) - r0.reshape(-1)

265 self.s.append(s0)

266

267 # We use the gradient which is minus the force!

268 y0 = f0.reshape(-1) - f.reshape(-1)

269 self.y.append(y0)

270

271 rho0 = 1.0 / np.dot(y0, s0)

272 self.rho.append(rho0)

273 self._just_reset_hessian = False

274

275 if len(self.y) > self.memory:

276 self.s.pop(0)

277 self.y.pop(0)

278 self.rho.pop(0)

279

280 def replay_trajectory(self, traj):

281 """Initialize history from old trajectory."""

282 if isinstance(traj, str):

283 from ase.io.trajectory import Trajectory

284 traj = Trajectory(traj, 'r')

285 r0 = None

286 f0 = None

287 # The last element is not added, as we get that for free when taking

288 # the first qn-step after the replay

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

290 r = traj[i].get_positions()

291 f = traj[i].get_forces()

292 self.update(r, f, r0, f0)

293 r0 = r.copy()

294 f0 = f.copy()

295 self.iteration += 1

296 self.r0 = r0

297 self.f0 = f0

298

299 def func(self, x):

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

301 self._actual_atoms.set_positions(x.reshape(-1, 3))

302 potl = self._actual_atoms.get_potential_energy()

303 return potl

304

305 def fprime(self, x):

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

307 self._actual_atoms.set_positions(x.reshape(-1, 3))

308 # Remember that forces are minus the gradient!

309 return -self._actual_atoms.get_forces().reshape(-1)

310

311 def line_search(self, r, g, e, previously_reset_hessian):

312 self.p = self.p.ravel()

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

314 if p_size <= np.sqrt(len(self._actual_atoms) * 1e-10):

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

316 g = g.ravel()

317 r = r.ravel()

318

319 if self.use_armijo:

320 try:

321 # CO: modified call to ls.run

322 # TODO: pass also the old slope to the linesearch

323 # so that the RumPath can extract a better starting guess?

324 # alternatively: we can adjust the rotation_factors

325 # out using some extrapolation tricks?

326 ls = LineSearchArmijo(self.func, c1=self.c1, tol=1e-14)

327 step, func_val, no_update = ls.run(

328 r, self.p, a_min=self.a_min,

329 func_start=e,

330 func_prime_start=g,

331 func_old=self.e0,

332 rigid_units=self.rigid_units,

333 rotation_factors=self.rotation_factors,

334 maxstep=self.maxstep)

335 self.e0 = e

336 self.e1 = func_val

337 self.alpha_k = step

338 except (ValueError, RuntimeError):

339 if not previously_reset_hessian:

340 warnings.warn(

341 'Armijo linesearch failed, resetting Hessian and '

342 'trying again')

343 self.reset_hessian()

344 self.alpha_k = 0.0

345 else:

346 raise RuntimeError(

347 'Armijo linesearch failed after reset of Hessian, '

348 'aborting')

349

350 else:

351 ls = LineSearch()

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

353 ls._line_search(self.func, self.fprime, r, self.p, g,

354 e, self.e0, stpmin=self.a_min,

355 maxstep=self.maxstep, c1=self.c1,

356 c2=self.c2, stpmax=50.)

357 self.e1 = e

358 if self.alpha_k is None:

359 raise RuntimeError('Wolff lineSearch failed!')

360

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

362 if smax is None:

363 smax = fmax

364 self.smax = smax

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

366

367 def log(self, forces=None):

368 if forces is None:

369 forces = self._actual_atoms.get_forces()

370 if isinstance(self._actual_atoms, UnitCellFilter):

371 natoms = len(self._actual_atoms.atoms)

372 forces, stress = forces[:natoms], self._actual_atoms.stress

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

374 smax = sqrt((stress**2).max())

375 else:

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

377 if self.e1 is not None:

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

379 e = self.e1

380 else:

381 e = self._actual_atoms.get_potential_energy()

382 T = time.localtime()

383 if self.logfile is not None:

384 name = self.__class__.__name__

385 if isinstance(self._actual_atoms, UnitCellFilter):

386 self.logfile.write(

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

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

389

390 else:

391 self.logfile.write(

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

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

394 self.logfile.flush()

395

396 def converged(self, forces=None):

397 """Did the optimization converge?"""

398 if forces is None:

399 forces = self._actual_atoms.get_forces()

400 if isinstance(self._actual_atoms, UnitCellFilter):

401 natoms = len(self._actual_atoms.atoms)

402 forces, stress = forces[:natoms], self._actual_atoms.stress

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

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

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

406 else:

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

408 return fmax_sq < self.fmax**2