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

1"""Structure optimization. """ 

2import time 

3import warnings 

4from collections.abc import Callable 

5from math import sqrt 

6from os.path import isfile 

7from typing import IO, Any, Dict, List, Optional, Tuple, Union 

8 

9from ase import Atoms 

10from ase.calculators.calculator import PropertyNotImplementedError 

11from ase.filters import UnitCellFilter 

12from ase.parallel import world 

13from ase.utils import IOContext, lazyproperty 

14from ase.utils.abc import Optimizable 

15 

16DEFAULT_MAX_STEPS = 100_000_000 

17 

18 

19class RestartError(RuntimeError): 

20 pass 

21 

22 

23class OptimizableAtoms(Optimizable): 

24 def __init__(self, atoms): 

25 self.atoms = atoms 

26 

27 def get_positions(self): 

28 return self.atoms.get_positions() 

29 

30 def set_positions(self, positions): 

31 self.atoms.set_positions(positions) 

32 

33 def get_forces(self): 

34 return self.atoms.get_forces() 

35 

36 @lazyproperty 

37 def _use_force_consistent_energy(self): 

38 # This boolean is in principle invalidated if the 

39 # calculator changes. This can lead to weird things 

40 # in multi-step optimizations. 

41 try: 

42 self.atoms.get_potential_energy(force_consistent=True) 

43 except PropertyNotImplementedError: 

44 # warnings.warn( 

45 # 'Could not get force consistent energy (\'free_energy\'). ' 

46 # 'Please make sure calculator provides \'free_energy\', even ' 

47 # 'if equal to the ordinary energy. ' 

48 # 'This will raise an error in future versions of ASE.', 

49 # FutureWarning) 

50 return False 

51 else: 

52 return True 

53 

54 def get_potential_energy(self): 

55 force_consistent = self._use_force_consistent_energy 

56 return self.atoms.get_potential_energy( 

57 force_consistent=force_consistent) 

58 

59 def iterimages(self): 

60 # XXX document purpose of iterimages 

61 return self.atoms.iterimages() 

62 

63 def __len__(self): 

64 # TODO: return 3 * len(self.atoms), because we want the length 

65 # of this to be the number of DOFs 

66 return len(self.atoms) 

67 

68 

69class Dynamics(IOContext): 

70 """Base-class for all MD and structure optimization classes.""" 

71 

72 def __init__( 

73 self, 

74 atoms: Atoms, 

75 logfile: Optional[Union[IO, str]] = None, 

76 trajectory: Optional[str] = None, 

77 append_trajectory: bool = False, 

78 master: Optional[bool] = None, 

79 comm=world, 

80 *, 

81 loginterval: int = 1, 

82 ): 

83 """Dynamics object. 

84 

85 Parameters 

86 ---------- 

87 atoms : Atoms object 

88 The Atoms object to operate on. 

89 

90 logfile : file object or str 

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

92 Use '-' for stdout. 

93 

94 trajectory : Trajectory object or str 

95 Attach trajectory object. If *trajectory* is a string a 

96 Trajectory will be constructed. Use *None* for no 

97 trajectory. 

98 

99 append_trajectory : bool 

100 Defaults to False, which causes the trajectory file to be 

101 overwriten each time the dynamics is restarted from scratch. 

102 If True, the new structures are appended to the trajectory 

103 file instead. 

104 

105 master : bool 

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

107 true, this rank will save files. 

108 

109 comm : Communicator object 

110 Communicator to handle parallel file reading and writing. 

111 

112 loginterval : int, default: 1 

113 Only write a log line for every *loginterval* time steps. 

114 """ 

115 self.atoms = atoms 

116 self.optimizable = atoms.__ase_optimizable__() 

117 self.logfile = self.openfile(file=logfile, comm=comm, mode='a') 

118 self.observers: List[Tuple[Callable, int, Tuple, Dict[str, Any]]] = [] 

119 self.nsteps = 0 

120 self.max_steps = 0 # to be updated in run or irun 

121 self.comm = comm 

122 

123 if trajectory is not None: 

124 if isinstance(trajectory, str): 

125 from ase.io.trajectory import Trajectory 

126 mode = "a" if append_trajectory else "w" 

127 trajectory = self.closelater(Trajectory( 

128 trajectory, mode=mode, master=master, comm=comm 

129 )) 

130 self.attach( 

131 trajectory, 

132 interval=loginterval, 

133 atoms=self.optimizable, 

134 ) 

135 

136 self.trajectory = trajectory 

137 

138 def todict(self) -> Dict[str, Any]: 

139 raise NotImplementedError 

140 

141 def get_number_of_steps(self): 

142 return self.nsteps 

143 

144 def insert_observer( 

145 self, function, position=0, interval=1, *args, **kwargs 

146 ): 

147 """Insert an observer. 

148 

149 This can be used for pre-processing before logging and dumping. 

150 

151 Examples 

152 -------- 

153 >>> from ase.build import bulk 

154 >>> from ase.calculators.emt import EMT 

155 >>> from ase.optimize import BFGS 

156 ... 

157 ... 

158 >>> def update_info(atoms, opt): 

159 ... atoms.info["nsteps"] = opt.nsteps 

160 ... 

161 ... 

162 >>> atoms = bulk("Cu", cubic=True) * 2 

163 >>> atoms.rattle() 

164 >>> atoms.calc = EMT() 

165 >>> with BFGS(atoms, logfile=None, trajectory="opt.traj") as opt: 

166 ... opt.insert_observer(update_info, atoms=atoms, opt=opt) 

167 ... opt.run(fmax=0.05, steps=10) 

168 True 

169 """ 

170 if not isinstance(function, Callable): 

171 function = function.write 

172 self.observers.insert(position, (function, interval, args, kwargs)) 

173 

174 def attach(self, function, interval=1, *args, **kwargs): 

175 """Attach callback function. 

176 

177 If *interval > 0*, at every *interval* steps, call *function* with 

178 arguments *args* and keyword arguments *kwargs*. 

179 

180 If *interval <= 0*, after step *interval*, call *function* with 

181 arguments *args* and keyword arguments *kwargs*. This is 

182 currently zero indexed.""" 

183 

184 if hasattr(function, "set_description"): 

185 d = self.todict() 

186 d.update(interval=interval) 

187 function.set_description(d) 

188 if not isinstance(function, Callable): 

189 function = function.write 

190 self.observers.append((function, interval, args, kwargs)) 

191 

192 def call_observers(self): 

193 for function, interval, args, kwargs in self.observers: 

194 call = False 

195 # Call every interval iterations 

196 if interval > 0: 

197 if (self.nsteps % interval) == 0: 

198 call = True 

199 # Call only on iteration interval 

200 elif interval <= 0: 

201 if self.nsteps == abs(interval): 

202 call = True 

203 if call: 

204 function(*args, **kwargs) 

205 

206 def irun(self, steps=DEFAULT_MAX_STEPS): 

207 """Run dynamics algorithm as generator. 

208 

209 Parameters 

210 ---------- 

211 steps : int, default=DEFAULT_MAX_STEPS 

212 Number of dynamics steps to be run. 

213 

214 Yields 

215 ------ 

216 converged : bool 

217 True if the forces on atoms are converged. 

218 

219 Examples 

220 -------- 

221 This method allows, e.g., to run two optimizers or MD thermostats at 

222 the same time. 

223 >>> opt1 = BFGS(atoms) 

224 >>> opt2 = BFGS(StrainFilter(atoms)).irun() 

225 >>> for _ in opt2: 

226 ... opt1.run() 

227 """ 

228 

229 # update the maximum number of steps 

230 self.max_steps = self.nsteps + steps 

231 

232 # compute the initial step 

233 self.optimizable.get_forces() 

234 

235 # log the initial step 

236 if self.nsteps == 0: 

237 self.log() 

238 

239 # we write a trajectory file if it is None 

240 if self.trajectory is None: 

241 self.call_observers() 

242 # We do not write on restart w/ an existing trajectory file 

243 # present. This duplicates the same entry twice 

244 elif len(self.trajectory) == 0: 

245 self.call_observers() 

246 

247 # check convergence 

248 is_converged = self.converged() 

249 yield is_converged 

250 

251 # run the algorithm until converged or max_steps reached 

252 while not is_converged and self.nsteps < self.max_steps: 

253 # compute the next step 

254 self.step() 

255 self.nsteps += 1 

256 

257 # log the step 

258 self.log() 

259 self.call_observers() 

260 

261 # check convergence 

262 is_converged = self.converged() 

263 yield is_converged 

264 

265 def run(self, steps=DEFAULT_MAX_STEPS): 

266 """Run dynamics algorithm. 

267 

268 This method will return when the forces on all individual 

269 atoms are less than *fmax* or when the number of steps exceeds 

270 *steps*. 

271 

272 Parameters 

273 ---------- 

274 steps : int, default=DEFAULT_MAX_STEPS 

275 Number of dynamics steps to be run. 

276 

277 Returns 

278 ------- 

279 converged : bool 

280 True if the forces on atoms are converged. 

281 """ 

282 

283 for converged in Dynamics.irun(self, steps=steps): 

284 pass 

285 return converged 

286 

287 def converged(self): 

288 """" a dummy function as placeholder for a real criterion, e.g. in 

289 Optimizer """ 

290 return False 

291 

292 def log(self, *args): 

293 """ a dummy function as placeholder for a real logger, e.g. in 

294 Optimizer """ 

295 return True 

296 

297 def step(self): 

298 """this needs to be implemented by subclasses""" 

299 raise RuntimeError("step not implemented.") 

300 

301 

302class Optimizer(Dynamics): 

303 """Base-class for all structure optimization classes.""" 

304 

305 # default maxstep for all optimizers 

306 defaults = {'maxstep': 0.2} 

307 _deprecated = object() 

308 

309 def __init__( 

310 self, 

311 atoms: Atoms, 

312 restart: Optional[str] = None, 

313 logfile: Optional[Union[IO, str]] = None, 

314 trajectory: Optional[str] = None, 

315 append_trajectory: bool = False, 

316 **kwargs, 

317 ): 

318 """ 

319 

320 Parameters 

321 ---------- 

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

323 The Atoms object to relax. 

324 

325 restart: str 

326 Filename for restart file. Default value is *None*. 

327 

328 logfile: file object or str 

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

330 Use '-' for stdout. 

331 

332 trajectory: Trajectory object or str 

333 Attach trajectory object. If *trajectory* is a string a 

334 Trajectory will be constructed. Use *None* for no 

335 trajectory. 

336 

337 append_trajectory: bool 

338 Appended to the trajectory file instead of overwriting it. 

339 

340 kwargs : dict, optional 

341 Extra arguments passed to :class:`~ase.optimize.optimize.Dynamics`. 

342 

343 """ 

344 super().__init__( 

345 atoms=atoms, 

346 logfile=logfile, 

347 trajectory=trajectory, 

348 append_trajectory=append_trajectory, 

349 **kwargs, 

350 ) 

351 

352 self.restart = restart 

353 

354 self.fmax = None 

355 

356 if restart is None or not isfile(restart): 

357 self.initialize() 

358 else: 

359 self.read() 

360 self.comm.barrier() 

361 

362 def read(self): 

363 raise NotImplementedError 

364 

365 def todict(self): 

366 description = { 

367 "type": "optimization", 

368 "optimizer": self.__class__.__name__, 

369 } 

370 # add custom attributes from subclasses 

371 for attr in ('maxstep', 'alpha', 'max_steps', 'restart', 

372 'fmax'): 

373 if hasattr(self, attr): 

374 description.update({attr: getattr(self, attr)}) 

375 return description 

376 

377 def initialize(self): 

378 pass 

379 

380 def irun(self, fmax=0.05, steps=DEFAULT_MAX_STEPS): 

381 """Run optimizer as generator. 

382 

383 Parameters 

384 ---------- 

385 fmax : float 

386 Convergence criterion of the forces on atoms. 

387 steps : int, default=DEFAULT_MAX_STEPS 

388 Number of optimizer steps to be run. 

389 

390 Yields 

391 ------ 

392 converged : bool 

393 True if the forces on atoms are converged. 

394 """ 

395 self.fmax = fmax 

396 return Dynamics.irun(self, steps=steps) 

397 

398 def run(self, fmax=0.05, steps=DEFAULT_MAX_STEPS): 

399 """Run optimizer. 

400 

401 Parameters 

402 ---------- 

403 fmax : float 

404 Convergence criterion of the forces on atoms. 

405 steps : int, default=DEFAULT_MAX_STEPS 

406 Number of optimizer steps to be run. 

407 

408 Returns 

409 ------- 

410 converged : bool 

411 True if the forces on atoms are converged. 

412 """ 

413 self.fmax = fmax 

414 return Dynamics.run(self, steps=steps) 

415 

416 def converged(self, forces=None): 

417 """Did the optimization converge?""" 

418 if forces is None: 

419 forces = self.optimizable.get_forces() 

420 return self.optimizable.converged(forces, self.fmax) 

421 

422 def log(self, forces=None): 

423 if forces is None: 

424 forces = self.optimizable.get_forces() 

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

426 e = self.optimizable.get_potential_energy() 

427 T = time.localtime() 

428 if self.logfile is not None: 

429 name = self.__class__.__name__ 

430 if self.nsteps == 0: 

431 args = (" " * len(name), "Step", "Time", "Energy", "fmax") 

432 msg = "%s %4s %8s %15s %12s\n" % args 

433 self.logfile.write(msg) 

434 

435 args = (name, self.nsteps, T[3], T[4], T[5], e, fmax) 

436 msg = "%s: %3d %02d:%02d:%02d %15.6f %15.6f\n" % args 

437 self.logfile.write(msg) 

438 self.logfile.flush() 

439 

440 def dump(self, data): 

441 from ase.io.jsonio import write_json 

442 if self.comm.rank == 0 and self.restart is not None: 

443 with open(self.restart, 'w') as fd: 

444 write_json(fd, data) 

445 

446 def load(self): 

447 from ase.io.jsonio import read_json 

448 with open(self.restart) as fd: 

449 try: 

450 from ase.optimize import BFGS 

451 if not isinstance(self, BFGS) and isinstance( 

452 self.atoms, UnitCellFilter 

453 ): 

454 warnings.warn( 

455 "WARNING: restart function is untested and may result " 

456 "in unintended behavior. Namely orig_cell is not " 

457 "loaded in the UnitCellFilter. Please test on your own" 

458 " to ensure consistent results." 

459 ) 

460 return read_json(fd, always_array=False) 

461 except Exception as ex: 

462 msg = ('Could not decode restart file as JSON. ' 

463 'You may need to delete the restart file ' 

464 f'{self.restart}') 

465 raise RestartError(msg) from ex