Coverage for /builds/ase/ase/ase/optimize/minimahopping.py: 57.57%

1# fmt: off

3import os

5import numpy as np

7from ase import io, units

8from ase.md import MDLogger, VelocityVerlet

9from ase.md.velocitydistribution import MaxwellBoltzmannDistribution

10from ase.optimize import QuasiNewton

11from ase.parallel import paropen, world

14class MinimaHopping:

15 """Implements the minima hopping method of global optimization outlined

16 by S. Goedecker, J. Chem. Phys. 120: 9911 (2004). Initialize with an

17 ASE atoms object. Optional parameters are fed through keywords.

18 To run multiple searches in parallel, specify the minima_traj keyword,

19 and have each run point to the same path.

20 """

22 _default_settings = {

23 'T0': 1000., # K, initial MD 'temperature'

24 'beta1': 1.1, # temperature adjustment parameter

25 'beta2': 1.1, # temperature adjustment parameter

26 'beta3': 1. / 1.1, # temperature adjustment parameter

27 'Ediff0': 0.5, # eV, initial energy acceptance threshold

28 'alpha1': 0.98, # energy threshold adjustment parameter

29 'alpha2': 1. / 0.98, # energy threshold adjustment parameter

30 'mdmin': 2, # criteria to stop MD simulation (no. of minima)

31 'logfile': 'hop.log', # text log

32 'minima_threshold': 0.5, # A, threshold for identical configs

33 'timestep': 1.0, # fs, timestep for MD simulations

34 'optimizer': QuasiNewton, # local optimizer to use

35 'minima_traj': 'minima.traj', # storage file for minima list

36 'fmax': 0.05} # eV/A, max force for optimizations

38 def __init__(self, atoms, **kwargs):

39 """Initialize with an ASE atoms object and keyword arguments."""

40 self._atoms = atoms

41 for key in kwargs:

42 if key not in self._default_settings:

43 raise RuntimeError(f'Unknown keyword: {key}')

44 for k, v in self._default_settings.items():

45 setattr(self, f'_{k}', kwargs.pop(k, v))

47 # when a MD sim. has passed a local minimum:

48 self._passedminimum = PassedMinimum()

50 # Misc storage.

51 self._previous_optimum = None

52 self._previous_energy = None

53 self._temperature = self._T0

54 self._Ediff = self._Ediff0

56 def __call__(self, totalsteps=None, maxtemp=None):

57 """Run the minima hopping algorithm. Can specify stopping criteria

58 with total steps allowed or maximum searching temperature allowed.

59 If neither is specified, runs indefinitely (or until stopped by

60 batching software)."""

61 self._startup()

62 while True:

63 if (totalsteps and self._counter >= totalsteps):

64 self._log('msg', 'Run terminated. Step #%i reached of '

65 '%i allowed. Increase totalsteps if resuming.'

66 % (self._counter, totalsteps))

67 return

68 if (maxtemp and self._temperature >= maxtemp):

69 self._log('msg', 'Run terminated. Temperature is %.2f K;'

70 ' max temperature allowed %.2f K.'

71 % (self._temperature, maxtemp))

72 return

74 self._previous_optimum = self._atoms.copy()

75 self._previous_energy = self._atoms.get_potential_energy()

76 self._molecular_dynamics()

77 self._optimize()

78 self._counter += 1

79 self._check_results()

81 def _startup(self):

82 """Initiates a run, and determines if running from previous data or

83 a fresh run."""

85 status = np.array(-1.)

86 exists = self._read_minima()

87 if world.rank == 0:

88 if not exists:

89 # Fresh run with new minima file.

90 status = np.array(0.)

91 elif not os.path.exists(self._logfile):

92 # Fresh run with existing or shared minima file.

93 status = np.array(1.)

94 else:

95 # Must be resuming from within a working directory.

96 status = np.array(2.)

97 world.barrier()

98 world.broadcast(status, 0)

100 if status == 2.:

101 self._resume()

102 else:

103 self._counter = 0

104 self._log('init')

105 self._log('msg', 'Performing initial optimization.')

106 if status == 1.:

107 self._log('msg', 'Using existing minima file with %i prior '

108 'minima: %s' % (len(self._minima),

109 self._minima_traj))

110 self._optimize()

111 self._check_results()

112 self._counter += 1

113

114 def _resume(self):

115 """Attempt to resume a run, based on information in the log

116 file. Note it will almost always be interrupted in the middle of

117 either a qn or md run or when exceeding totalsteps, so it only has

118 been tested in those cases currently."""

119 f = paropen(self._logfile, 'r')

120 lines = f.read().splitlines()

121 f.close()

122 self._log('msg', 'Attempting to resume stopped run.')

123 self._log('msg', 'Using existing minima file with %i prior '

124 'minima: %s' % (len(self._minima), self._minima_traj))

125 mdcount, qncount = 0, 0

126 for line in lines:

127 if (line[:4] == 'par:') and ('Ediff' not in line):

128 self._temperature = float(line.split()[1])

129 self._Ediff = float(line.split()[2])

130 elif line[:18] == 'msg: Optimization:':

131 qncount = int(line[19:].split('qn')[1])

132 elif line[:24] == 'msg: Molecular dynamics:':

133 mdcount = int(line[25:].split('md')[1])

134 self._counter = max((mdcount, qncount))

135 if qncount == mdcount:

136 # Either stopped during local optimization or terminated due to

137 # max steps.

138 self._log('msg', 'Attempting to resume at qn%05i' % qncount)

139 if qncount > 0:

140 atoms = io.read('qn%05i.traj' % (qncount - 1), index=-1)

141 self._previous_optimum = atoms.copy()

142 self._previous_energy = atoms.get_potential_energy()

143 if os.path.getsize('qn%05i.traj' % qncount) > 0:

144 atoms = io.read('qn%05i.traj' % qncount, index=-1)

145 else:

146 atoms = io.read('md%05i.traj' % qncount, index=-3)

147 self._atoms.positions = atoms.get_positions()

148 fmax = np.sqrt((atoms.get_forces() ** 2).sum(axis=1).max())

149 if fmax < self._fmax:

150 # Stopped after a qn finished.

151 self._log('msg', 'qn%05i fmax already less than fmax=%.3f'

152 % (qncount, self._fmax))

153 self._counter += 1

154 return

155 self._optimize()

156 self._counter += 1

157 if qncount > 0:

158 self._check_results()

159 else:

160 self._record_minimum()

161 self._log('msg', 'Found a new minimum.')

162 self._log('msg', 'Accepted new minimum.')

163 self._log('par')

164 elif qncount < mdcount:

165 # Probably stopped during molecular dynamics.

166 self._log('msg', 'Attempting to resume at md%05i.' % mdcount)

167 atoms = io.read('qn%05i.traj' % qncount, index=-1)

168 self._previous_optimum = atoms.copy()

169 self._previous_energy = atoms.get_potential_energy()

170 self._molecular_dynamics(resume=mdcount)

171 self._optimize()

172 self._counter += 1

173 self._check_results()

174

175 def _check_results(self):

176 """Adjusts parameters and positions based on outputs."""

177

178 # No prior minima found?

179 self._read_minima()

180 if len(self._minima) == 0:

181 self._log('msg', 'Found a new minimum.')

182 self._log('msg', 'Accepted new minimum.')

183 self._record_minimum()

184 self._log('par')

185 return

186 # Returned to starting position?

187 if self._previous_optimum:

188 compare = ComparePositions(translate=False)

189 dmax = compare(self._atoms, self._previous_optimum)

190 self._log('msg', 'Max distance to last minimum: %.3f A' % dmax)

191 if dmax < self._minima_threshold:

192 self._log('msg', 'Re-found last minimum.')

193 self._temperature *= self._beta1

194 self._log('par')

195 return

196 # In a previously found position?

197 unique, dmax_closest = self._unique_minimum_position()

198 self._log('msg', 'Max distance to closest minimum: %.3f A' %

199 dmax_closest)

200 if not unique:

201 self._temperature *= self._beta2

202 self._log('msg', 'Found previously found minimum.')

203 self._log('par')

204 if self._previous_optimum:

205 self._log('msg', 'Restoring last minimum.')

206 self._atoms.positions = self._previous_optimum.positions

207 return

208 # Must have found a unique minimum.

209 self._temperature *= self._beta3

210 self._log('msg', 'Found a new minimum.')

211 self._log('par')

212 if (self._previous_energy is None or

213 (self._atoms.get_potential_energy() <

214 self._previous_energy + self._Ediff)):

215 self._log('msg', 'Accepted new minimum.')

216 self._Ediff *= self._alpha1

217 self._log('par')

218 self._record_minimum()

219 else:

220 self._log('msg', 'Rejected new minimum due to energy. '

221 'Restoring last minimum.')

222 self._atoms.positions = self._previous_optimum.positions

223 self._Ediff *= self._alpha2

224 self._log('par')

225

226 def _log(self, cat='msg', message=None):

227 """Records the message as a line in the log file."""

228 if cat == 'init':

229 if world.rank == 0:

230 if os.path.exists(self._logfile):

231 raise RuntimeError(f'File exists: {self._logfile}')

232 fd = paropen(self._logfile, 'w')

233 fd.write('par: %12s %12s %12s\n' % ('T (K)', 'Ediff (eV)',

234 'mdmin'))

235 fd.write('ene: %12s %12s %12s\n' % ('E_current', 'E_previous',

236 'Difference'))

237 fd.close()

238 return

239 fd = paropen(self._logfile, 'a')

240 if cat == 'msg':

241 line = f'msg: {message}'

242 elif cat == 'par':

243 line = ('par: %12.4f %12.4f %12i' %

244 (self._temperature, self._Ediff, self._mdmin))

245 elif cat == 'ene':

246 current = self._atoms.get_potential_energy()

247 if self._previous_optimum:

248 previous = self._previous_energy

249 line = ('ene: %12.5f %12.5f %12.5f' %

250 (current, previous, current - previous))

251 else:

252 line = ('ene: %12.5f' % current)

253 fd.write(line + '\n')

254 fd.close()

255

256 def _optimize(self):

257 """Perform an optimization."""

258 self._atoms.set_momenta(np.zeros(self._atoms.get_momenta().shape))

259 with self._optimizer(self._atoms,

260 trajectory='qn%05i.traj' % self._counter,

261 logfile='qn%05i.log' % self._counter) as opt:

262 self._log('msg', 'Optimization: qn%05i' % self._counter)

263 opt.run(fmax=self._fmax)

264 self._log('ene')

265

266 def _record_minimum(self):

267 """Adds the current atoms configuration to the minima list."""

268 with io.Trajectory(self._minima_traj, 'a') as traj:

269 traj.write(self._atoms)

270 self._read_minima()

271 self._log('msg', 'Recorded minima #%i.' % (len(self._minima) - 1))

272

273 def _read_minima(self):

274 """Reads in the list of minima from the minima file."""

275 exists = os.path.exists(self._minima_traj)

276 if exists:

277 empty = os.path.getsize(self._minima_traj) == 0

278 if not empty:

279 with io.Trajectory(self._minima_traj, 'r') as traj:

280 self._minima = [atoms for atoms in traj]

281 else:

282 self._minima = []

283 return True

284 else:

285 self._minima = []

286 return False

287

288 def _molecular_dynamics(self, resume=None):

289 """Performs a molecular dynamics simulation, until mdmin is

290 exceeded. If resuming, the file number (md%05i) is expected."""

291 self._log('msg', 'Molecular dynamics: md%05i' % self._counter)

292 mincount = 0

293 energies, oldpositions = [], []

294 thermalized = False

295 if resume:

296 self._log('msg', 'Resuming MD from md%05i.traj' % resume)

297 if os.path.getsize('md%05i.traj' % resume) == 0:

298 self._log('msg', 'md%05i.traj is empty. Resuming from '

299 'qn%05i.traj.' % (resume, resume - 1))

300 atoms = io.read('qn%05i.traj' % (resume - 1), index=-1)

301 else:

302 with io.Trajectory('md%05i.traj' % resume, 'r') as images:

303 for atoms in images:

304 energies.append(atoms.get_potential_energy())

305 oldpositions.append(atoms.positions.copy())

306 passedmin = self._passedminimum(energies)

307 if passedmin:

308 mincount += 1

309 self._atoms.set_momenta(atoms.get_momenta())

310 thermalized = True

311 self._atoms.positions = atoms.get_positions()

312 self._log('msg', 'Starting MD with %i existing energies.' %

313 len(energies))

314 if not thermalized:

315 MaxwellBoltzmannDistribution(self._atoms,

316 temperature_K=self._temperature,

317 force_temp=True)

318 traj = io.Trajectory('md%05i.traj' % self._counter, 'a',

319 self._atoms)

320 dyn = VelocityVerlet(self._atoms, timestep=self._timestep * units.fs)

321 log = MDLogger(dyn, self._atoms, 'md%05i.log' % self._counter,

322 header=True, stress=False, peratom=False)

323

324 with traj, dyn, log:

325 dyn.attach(log, interval=1)

326 dyn.attach(traj, interval=1)

327 while mincount < self._mdmin:

328 dyn.run(1)

329 energies.append(self._atoms.get_potential_energy())

330 passedmin = self._passedminimum(energies)

331 if passedmin:

332 mincount += 1

333 oldpositions.append(self._atoms.positions.copy())

334 # Reset atoms to minimum point.

335 self._atoms.positions = oldpositions[passedmin[0]]

336

337 def _unique_minimum_position(self):

338 """Identifies if the current position of the atoms, which should be

339 a local minima, has been found before."""

340 unique = True

341 dmax_closest = 99999.

342 compare = ComparePositions(translate=True)

343 self._read_minima()

344 for minimum in self._minima:

345 dmax = compare(minimum, self._atoms)

346 if dmax < self._minima_threshold:

347 unique = False

348 if dmax < dmax_closest:

349 dmax_closest = dmax

350 return unique, dmax_closest

351

352

353class ComparePositions:

354 """Class that compares the atomic positions between two ASE atoms

355 objects. Returns the maximum distance that any atom has moved, assuming

356 all atoms of the same element are indistinguishable. If translate is

357 set to True, allows for arbitrary translations within the unit cell,

358 as well as translations across any periodic boundary conditions. When

359 called, returns the maximum displacement of any one atom."""

360

361 def __init__(self, translate=True):

362 self._translate = translate

363

364 def __call__(self, atoms1, atoms2):

365 atoms1 = atoms1.copy()

366 atoms2 = atoms2.copy()

367 if not self._translate:

368 dmax = self. _indistinguishable_compare(atoms1, atoms2)

369 else:

370 dmax = self._translated_compare(atoms1, atoms2)

371 return dmax

372

373 def _translated_compare(self, atoms1, atoms2):

374 """Moves the atoms around and tries to pair up atoms, assuming any

375 atoms with the same symbol are indistinguishable, and honors

376 periodic boundary conditions (for example, so that an atom at

377 (0.1, 0., 0.) correctly is found to be close to an atom at

378 (7.9, 0., 0.) if the atoms are in an orthorhombic cell with

379 x-dimension of 8. Returns dmax, the maximum distance between any

380 two atoms in the optimal configuration."""

381 atoms1.set_constraint()

382 atoms2.set_constraint()

383 for index in range(3):

384 assert atoms1.pbc[index] == atoms2.pbc[index]

385 least = self._get_least_common(atoms1)

386 indices1 = [atom.index for atom in atoms1 if atom.symbol == least[0]]

387 indices2 = [atom.index for atom in atoms2 if atom.symbol == least[0]]

388 # Make comparison sets from atoms2, which contain repeated atoms in

389 # all pbc's and bring the atom listed in indices2 to (0,0,0)

390 comparisons = []

391 repeat = []

392 for bc in atoms2.pbc:

393 if bc:

394 repeat.append(3)

395 else:

396 repeat.append(1)

397 repeated = atoms2.repeat(repeat)

398 moved_cell = atoms2.cell * atoms2.pbc

399 for moved in moved_cell:

400 repeated.translate(-moved)

401 repeated.set_cell(atoms2.cell)

402 for index in indices2:

403 comparison = repeated.copy()

404 comparison.translate(-atoms2[index].position)

405 comparisons.append(comparison)

406 # Bring the atom listed in indices1 to (0,0,0) [not whole list]

407 standard = atoms1.copy()

408 standard.translate(-atoms1[indices1[0]].position)

409 # Compare the standard to the comparison sets.

410 dmaxes = []

411 for comparison in comparisons:

412 dmax = self._indistinguishable_compare(standard, comparison)

413 dmaxes.append(dmax)

414 return min(dmaxes)

415

416 def _get_least_common(self, atoms):

417 """Returns the least common element in atoms. If more than one,

418 returns the first encountered."""

419 symbols = [atom.symbol for atom in atoms]

420 least = ['', np.inf]

421 for element in set(symbols):

422 count = symbols.count(element)

423 if count < least[1]:

424 least = [element, count]

425 return least

426

427 def _indistinguishable_compare(self, atoms1, atoms2):

428 """Finds each atom in atoms1's nearest neighbor with the same

429 chemical symbol in atoms2. Return dmax, the farthest distance an

430 individual atom differs by."""

431 atoms2 = atoms2.copy() # allow deletion

432 atoms2.set_constraint()

433 dmax = 0.

434 for atom1 in atoms1:

435 closest = [np.nan, np.inf]

436 for index, atom2 in enumerate(atoms2):

437 if atom2.symbol == atom1.symbol:

438 d = np.linalg.norm(atom1.position - atom2.position)

439 if d < closest[1]:

440 closest = [index, d]

441 if closest[1] > dmax:

442 dmax = closest[1]

443 del atoms2[closest[0]]

444 return dmax

445

446

447class PassedMinimum:

448 """Simple routine to find if a minimum in the potential energy surface

449 has been passed. In its default settings, a minimum is found if the

450 sequence ends with two downward points followed by two upward points.

451 Initialize with n_down and n_up, integer values of the number of up and

452 down points. If it has successfully determined it passed a minimum, it

453 returns the value (energy) of that minimum and the number of positions

454 back it occurred, otherwise returns None."""

455

456 def __init__(self, n_down=2, n_up=2):

457 self._ndown = n_down

458 self._nup = n_up

459

460 def __call__(self, energies):

461 if len(energies) < (self._nup + self._ndown + 1):

462 return None

463 status = True

464 index = -1

465 for _ in range(self._nup):

466 if energies[index] < energies[index - 1]:

467 status = False

468 index -= 1

469 for _ in range(self._ndown):

470 if energies[index] > energies[index - 1]:

471 status = False

472 index -= 1

473 if status:

474 return (-self._nup - 1), energies[-self._nup - 1]

475

476

477class MHPlot:

478 """Makes a plot summarizing the output of the MH algorithm from the

479 specified rundirectory. If no rundirectory is supplied, uses the

480 current directory."""

481

482 def __init__(self, rundirectory=None, logname='hop.log'):

483 if not rundirectory:

484 rundirectory = os.getcwd()

485 self._rundirectory = rundirectory

486 self._logname = logname

487 self._read_log()

488 self._fig, self._ax = self._makecanvas()

489 self._plot_data()

490

491 def get_figure(self):

492 """Returns the matplotlib figure object."""

493 return self._fig

494

495 def save_figure(self, filename):

496 """Saves the file to the specified path, with any allowed

497 matplotlib extension (e.g., .pdf, .png, etc.)."""

498 self._fig.savefig(filename)

499

500 def _read_log(self):

501 """Reads relevant parts of the log file."""

502 data = [] # format: [energy, status, temperature, ediff]

503

504 with open(os.path.join(self._rundirectory, self._logname)) as fd:

505 lines = fd.read().splitlines()

506

507 step_almost_over = False

508 step_over = False

509 for line in lines:

510 if line.startswith('msg: Molecular dynamics:'):

511 status = 'performing MD'

512 elif line.startswith('msg: Optimization:'):

513 status = 'performing QN'

514 elif line.startswith('ene:'):

515 status = 'local optimum reached'

516 energy = floatornan(line.split()[1])

517 elif line.startswith('msg: Accepted new minimum.'):

518 status = 'accepted'

519 step_almost_over = True

520 elif line.startswith('msg: Found previously found minimum.'):

521 status = 'previously found minimum'

522 step_almost_over = True

523 elif line.startswith('msg: Re-found last minimum.'):

524 status = 'previous minimum'

525 step_almost_over = True

526 elif line.startswith('msg: Rejected new minimum'):

527 status = 'rejected'

528 step_almost_over = True

529 elif line.startswith('par: '):

530 temperature = floatornan(line.split()[1])

531 ediff = floatornan(line.split()[2])

532 if step_almost_over:

533 step_over = True

534 step_almost_over = False

535 if step_over:

536 data.append([energy, status, temperature, ediff])

537 step_over = False

538 if data[-1][1] != status:

539 data.append([np.nan, status, temperature, ediff])

540 self._data = data

541

542 def _makecanvas(self):

543 from matplotlib import pyplot

544 from matplotlib.ticker import ScalarFormatter

545 fig = pyplot.figure(figsize=(6., 8.))

546 lm, rm, bm, tm = 0.22, 0.02, 0.05, 0.04

547 vg1 = 0.01 # between adjacent energy plots

548 vg2 = 0.03 # between different types of plots

549 ratio = 2. # size of an energy plot to a parameter plot

550 figwidth = 1. - lm - rm

551 totalfigheight = 1. - bm - tm - vg1 - 2. * vg2

552 parfigheight = totalfigheight / (2. * ratio + 2)

553 epotheight = ratio * parfigheight

554 ax1 = fig.add_axes((lm, bm, figwidth, epotheight))

555 ax2 = fig.add_axes((lm, bm + epotheight + vg1,

556 figwidth, epotheight))

557 for ax in [ax1, ax2]:

558 ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))

559 ediffax = fig.add_axes((lm, bm + 2. * epotheight + vg1 + vg2,

560 figwidth, parfigheight))

561 tempax = fig.add_axes((lm, (bm + 2 * epotheight + vg1 + 2 * vg2 +

562 parfigheight), figwidth, parfigheight))

563 for ax in [ax2, tempax, ediffax]:

564 ax.set_xticklabels([])

565 ax1.set_xlabel('step')

566 tempax.set_ylabel('$T$, K')

567 ediffax.set_ylabel(r'$E_\mathrm{diff}$, eV')

568 for ax in [ax1, ax2]:

569 ax.set_ylabel(r'$E_\mathrm{pot}$, eV')

570 ax = CombinedAxis(ax1, ax2, tempax, ediffax)

571 self._set_zoomed_range(ax)

572 ax1.spines['top'].set_visible(False)

573 ax2.spines['bottom'].set_visible(False)

574 return fig, ax

575

576 def _set_zoomed_range(self, ax):

577 """Try to intelligently set the range for the zoomed-in part of the

578 graph."""

579 energies = [line[0] for line in self._data

580 if not np.isnan(line[0])]

581 dr = max(energies) - min(energies)

582 if dr == 0.:

583 dr = 1.

584 ax.set_ax1_range((min(energies) - 0.2 * dr,

585 max(energies) + 0.2 * dr))

586

587 def _plot_data(self):

588 for step, line in enumerate(self._data):

589 self._plot_energy(step, line)

590 self._plot_qn(step, line)

591 self._plot_md(step, line)

592 self._plot_parameters()

593 self._ax.set_xlim(self._ax.ax1.get_xlim())

594

595 def _plot_energy(self, step, line):

596 """Plots energy and annotation for acceptance."""

597 energy, status = line[0], line[1]

598 if np.isnan(energy):

599 return

600 self._ax.plot([step, step + 0.5], [energy] * 2, '-',

601 color='k', linewidth=2.)

602 if status == 'accepted':

603 self._ax.text(step + 0.51, energy, r'$\checkmark$')

604 elif status == 'rejected':

605 self._ax.text(step + 0.51, energy, r'$\Uparrow$', color='red')

606 elif status == 'previously found minimum':

607 self._ax.text(step + 0.51, energy, r'$\hookleftarrow$',

608 color='red', va='center')

609 elif status == 'previous minimum':

610 self._ax.text(step + 0.51, energy, r'$\leftarrow$',

611 color='red', va='center')

612

613 def _plot_md(self, step, line):

614 """Adds a curved plot of molecular dynamics trajectory."""

615 if step == 0:

616 return

617 energies = [self._data[step - 1][0]]

618 file = os.path.join(self._rundirectory, 'md%05i.traj' % step)

619 with io.Trajectory(file, 'r') as traj:

620 for atoms in traj:

621 energies.append(atoms.get_potential_energy())

622 xi = step - 1 + .5

623 if len(energies) > 2:

624 xf = xi + (step + 0.25 - xi) * len(energies) / (len(energies) - 2.)

625 else:

626 xf = step

627 if xf > (step + .75):

628 xf = step

629 self._ax.plot(np.linspace(xi, xf, num=len(energies)), energies,

630 '-k')

631

632 def _plot_qn(self, index, line):

633 """Plots a dashed vertical line for the optimization."""

634 if line[1] == 'performing MD':

635 return

636 file = os.path.join(self._rundirectory, 'qn%05i.traj' % index)

637 if os.path.getsize(file) == 0:

638 return

639 with io.Trajectory(file, 'r') as traj:

640 energies = [traj[0].get_potential_energy(),

641 traj[-1].get_potential_energy()]

642 if index > 0:

643 file = os.path.join(self._rundirectory, 'md%05i.traj' % index)

644 atoms = io.read(file, index=-3)

645 energies[0] = atoms.get_potential_energy()

646 self._ax.plot([index + 0.25] * 2, energies, ':k')

647

648 def _plot_parameters(self):

649 """Adds a plot of temperature and Ediff to the plot."""

650 steps, Ts, ediffs = [], [], []

651 for step, line in enumerate(self._data):

652 steps.extend([step + 0.5, step + 1.5])

653 Ts.extend([line[2]] * 2)

654 ediffs.extend([line[3]] * 2)

655 self._ax.tempax.plot(steps, Ts)

656 self._ax.ediffax.plot(steps, ediffs)

657

658 for ax in [self._ax.tempax, self._ax.ediffax]:

659 ylim = ax.get_ylim()

660 yrange = ylim[1] - ylim[0]

661 ax.set_ylim((ylim[0] - 0.1 * yrange, ylim[1] + 0.1 * yrange))

662

663

664def floatornan(value):

665 """Converts the argument into a float if possible, np.nan if not."""

666 try:

667 output = float(value)

668 except ValueError:

669 output = np.nan

670 return output

671

672

673class CombinedAxis:

674 """Helper class for MHPlot to plot on split y axis and adjust limits

675 simultaneously."""

676

677 def __init__(self, ax1, ax2, tempax, ediffax):

678 self.ax1 = ax1

679 self.ax2 = ax2

680 self.tempax = tempax

681 self.ediffax = ediffax

682 self._ymax = -np.inf

683

684 def set_ax1_range(self, ylim):

685 self._ax1_ylim = ylim

686 self.ax1.set_ylim(ylim)

687

688 def plot(self, *args, **kwargs):

689 self.ax1.plot(*args, **kwargs)

690 self.ax2.plot(*args, **kwargs)

691 # Re-adjust yrange

692 for yvalue in args[1]:

693 if yvalue > self._ymax:

694 self._ymax = yvalue

695 self.ax1.set_ylim(self._ax1_ylim)

696 self.ax2.set_ylim((self._ax1_ylim[1], self._ymax))

697

698 def set_xlim(self, *args):

699 self.ax1.set_xlim(*args)

700 self.ax2.set_xlim(*args)

701 self.tempax.set_xlim(*args)

702 self.ediffax.set_xlim(*args)

703

704 def text(self, *args, **kwargs):

705 y = args[1]

706 if y < self._ax1_ylim[1]:

707 ax = self.ax1

708 else:

709 ax = self.ax2

710 ax.text(*args, **kwargs)