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

1# fmt: off

3import time

4import warnings

5from math import sqrt

7import numpy as np

9from ase.filters import UnitCellFilter

10from ase.optimize.optimize import Optimizer

11from ase.optimize.precon.precon import make_precon

12from ase.utils.linesearch import LineSearch

13from ase.utils.linesearcharmijo import LineSearchArmijo

16class PreconLBFGS(Optimizer):

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

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

19 where a good preconditioner is available.

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

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

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

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

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

26 in particular crystalline. For systems outside this category,

27 PreconLBFGS with Exp preconditioner may yield unpredictable results.

29 In time this implementation is expected to replace

30 ase.optimize.lbfgs.LBFGS.

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

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

34 preconditioner for simulating condensed phase materials

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

36 """

38 # CO : added parameters rigid_units and rotation_factors

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

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

41 precon='auto', variable_cell=False,

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

43 rigid_units=None, rotation_factors=None, Hinv=None, **kwargs):

44 """

46 Parameters

47 ----------

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

49 The Atoms object to relax.

51 restart: str

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

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

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

56 logfile: file object or str

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

58 Use '-' for stdout.

60 trajectory: str

61 Trajectory file used to store optimisation path.

63 maxstep: float

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

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

66 Default is 0.04 Angstrom.

68 memory: int

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

70 arrays of this length containing floats are stored.

72 damping: float

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

74 the positions.

76 alpha: float

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

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

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

80 a lower value also means risk of instability.

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 kwargs : dict, optional

123 Extra arguments passed to

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

125

126 """

127 if variable_cell:

128 atoms = UnitCellFilter(atoms)

129 Optimizer.__init__(self, atoms, restart, logfile, trajectory, **kwargs)

130

131 self._actual_atoms = atoms

132

133 # default preconditioner

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

135 if precon == 'auto':

136 if len(atoms) < 100:

137 precon = None

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

139 'the standard preconditioner, hence it is '

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

141 '`PreconLBFGS` by explicitly providing the '

142 'kwarg `precon`')

143 else:

144 precon = 'Exp'

145

146 if maxstep is not None:

147 if maxstep > 1.0:

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

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

150 maxstep)

151 self.maxstep = maxstep

152 else:

153 self.maxstep = 0.04

154

155 self.memory = memory

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

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

158 # Note that this is never changed!

159 self.Hinv = Hinv

160 self.damping = damping

161 self.p = None

162

163 # construct preconditioner if passed as a string

164 self.precon = make_precon(precon)

165 self.use_armijo = use_armijo

166 self.c1 = c1

167 self.c2 = c2

168 self.a_min = a_min

169 if self.a_min is None:

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

171

172 # CO

173 self.rigid_units = rigid_units

174 self.rotation_factors = rotation_factors

175

176 def reset_hessian(self):

177 """

178 Throw away history of the Hessian

179 """

180 self._just_reset_hessian = True

181 self.s = []

182 self.y = []

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

184 # again and again

185

186 def initialize(self):

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

188 self.iteration = 0

189 self.reset_hessian()

190 self.r0 = None

191 self.f0 = None

192 self.e0 = None

193 self.e1 = None

194 self.task = 'START'

195 self.load_restart = False

196

197 def read(self):

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

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

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

201 self.load_restart = True

202

203 def step(self, f=None):

204 """Take a single step

205

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

207 then take it"""

208 r = self._actual_atoms.get_positions()

209

210 if f is None:

211 f = self._actual_atoms.get_forces()

212

213 previously_reset_hessian = self._just_reset_hessian

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

215

216 s = self.s

217 y = self.y

218 rho = self.rho

219 H0 = self.H0

220

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

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

223

224 # The algorithm itself:

225 q = -f.reshape(-1)

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

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

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

229

230 if self.precon is None:

231 if self.Hinv is not None:

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

233 else:

234 z = H0 * q

235 else:

236 self.precon.make_precon(self._actual_atoms)

237 z = self.precon.solve(q)

238

239 for i in range(loopmax):

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

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

242

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

244 ###

245

246 g = -f

247 if self.e1 is not None:

248 e = self.e1

249 else:

250 e = self.func(r)

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

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

253

254 if self.alpha_k != 0.0:

255 self._actual_atoms.set_positions(r + dr)

256

257 self.iteration += 1

258 self.r0 = r

259 self.f0 = -g

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

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

262

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

264 """Update everything that is kept in memory

265

266 This function is mostly here to allow for replay_trajectory.

267 """

268 if not self._just_reset_hessian:

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

270 self.s.append(s0)

271

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

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

274 self.y.append(y0)

275

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

277 self.rho.append(rho0)

278 self._just_reset_hessian = False

279

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

281 self.s.pop(0)

282 self.y.pop(0)

283 self.rho.pop(0)

284

285 def replay_trajectory(self, traj):

286 """Initialize history from old trajectory."""

287 if isinstance(traj, str):

288 from ase.io.trajectory import Trajectory

289 traj = Trajectory(traj, 'r')

290 r0 = None

291 f0 = None

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

293 # the first qn-step after the replay

294 for i in range(len(traj) - 1):

295 r = traj[i].get_positions()

296 f = traj[i].get_forces()

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

298 r0 = r.copy()

299 f0 = f.copy()

300 self.iteration += 1

301 self.r0 = r0

302 self.f0 = f0

303

304 def func(self, x):

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

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

307 potl = self._actual_atoms.get_potential_energy()

308 return potl

309

310 def fprime(self, x):

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

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

313 # Remember that forces are minus the gradient!

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

315

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

317 self.p = self.p.ravel()

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

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

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

321 g = g.ravel()

322 r = r.ravel()

323

324 if self.use_armijo:

325 try:

326 # CO: modified call to ls.run

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

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

329 # alternatively: we can adjust the rotation_factors

330 # out using some extrapolation tricks?

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

332 step, func_val, _no_update = ls.run(

333 r, self.p, a_min=self.a_min,

334 func_start=e,

335 func_prime_start=g,

336 func_old=self.e0,

337 rigid_units=self.rigid_units,

338 rotation_factors=self.rotation_factors,

339 maxstep=self.maxstep)

340 self.e0 = e

341 self.e1 = func_val

342 self.alpha_k = step

343 except (ValueError, RuntimeError):

344 if not previously_reset_hessian:

345 warnings.warn(

346 'Armijo linesearch failed, resetting Hessian and '

347 'trying again')

348 self.reset_hessian()

349 self.alpha_k = 0.0

350 else:

351 raise RuntimeError(

352 'Armijo linesearch failed after reset of Hessian, '

353 'aborting')

354

355 else:

356 ls = LineSearch()

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

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

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

360 maxstep=self.maxstep, c1=self.c1,

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

362 self.e1 = e

363 if self.alpha_k is None:

364 raise RuntimeError('Wolff lineSearch failed!')

365

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

367 if smax is None:

368 smax = fmax

369 self.smax = smax

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

371

372 def log(self, gradient):

373 forces = self._actual_atoms.get_forces()

374 if isinstance(self._actual_atoms, UnitCellFilter):

375 natoms = len(self._actual_atoms.atoms)

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

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

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

379 else:

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

381 if self.e1 is not None:

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

383 e = self.e1

384 else:

385 e = self._actual_atoms.get_potential_energy()

386 T = time.localtime()

387 if self.logfile is not None:

388 name = self.__class__.__name__

389 if isinstance(self._actual_atoms, UnitCellFilter):

390 self.logfile.write(

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

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

393

394 else:

395 self.logfile.write(

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

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

398 self.logfile.flush()

399

400 def converged(self, gradient):

401 """Did the optimization converge?"""

402 # XXX ignoring gradient

403 forces = self._actual_atoms.get_forces()

404 if isinstance(self._actual_atoms, UnitCellFilter):

405 natoms = len(self._actual_atoms.atoms)

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

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

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

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

410 else:

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

412 return fmax_sq < self.fmax**2