Coverage for /builds/ase/ase/ase/autoneb.py : 75.00%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import numpy as np
2import shutil
3import os
4import types
5from math import log
6from math import exp
7from contextlib import ExitStack
8from pathlib import Path
9from warnings import warn
11from ase.io import Trajectory
12from ase.io import read
13from ase.neb import NEB
14from ase.optimize import BFGS
15from ase.optimize import FIRE
16from ase.calculators.singlepoint import SinglePointCalculator
17import ase.parallel as mpi
20class AutoNEB:
21 """AutoNEB object.
23 The AutoNEB algorithm streamlines the execution of NEB and CI-NEB
24 calculations following the algorithm described in:
26 E. L. Kolsbjerg, M. N. Groves, and B. Hammer, J. Chem. Phys,
27 145, 094107, 2016. (doi: 10.1063/1.4961868)
29 The user supplies at minimum the two end-points and possibly also some
30 intermediate images.
32 The stages are:
33 1) Define a set of images and name them sequentially.
34 Must at least have a relaxed starting and ending image
35 User can supply intermediate guesses which do not need to
36 have previously determined energies (probably from another
37 NEB calculation with a lower level of theory)
38 2) AutoNEB will first evaluate the user provided intermediate images
39 3) AutoNEB will then add additional images dynamically until n_max
40 is reached
41 4) A climbing image will attempt to locate the saddle point
42 5) All the images between the highest point and the starting point
43 are further relaxed to smooth the path
44 6) All the images between the highest point and the ending point are
45 further relaxed to smooth the path
47 Step 4 and 5-6 are optional steps!
49 Parameters:
51 attach_calculators:
52 Function which adds valid calculators to the list of images supplied.
53 prefix: string or path
54 All files that the AutoNEB method reads and writes are prefixed with
55 prefix
56 n_simul: int
57 The number of relaxations run in parallel.
58 n_max: int
59 The number of images along the NEB path when done.
60 This number includes the two end-points.
61 Important: due to the dynamic adding of images around the peak n_max
62 must be updated if the NEB is restarted.
63 climb: boolean
64 Should a CI-NEB calculation be done at the top-point
65 fmax: float or list of floats
66 The maximum force along the NEB path
67 maxsteps: int
68 The maximum number of steps in each NEB relaxation.
69 If a list is given the first number of steps is used in the build-up
70 and final scan phase;
71 the second number of steps is used in the CI step after all images
72 have been inserted.
73 k: float
74 The spring constant along the NEB path
75 method: str (see neb.py)
76 Choice betweeen three method:
77 'aseneb', standard ase NEB implementation
78 'improvedtangent', published NEB implementation
79 'eb', full spring force implementation (default)
80 optimizer: object
81 Optimizer object, defaults to FIRE
82 Use of the valid strings 'BFGS' and 'FIRE' is deprecated.
83 space_energy_ratio: float
84 The preference for new images to be added in a big energy gab
85 with a preference around the peak or in the biggest geometric gab.
86 A space_energy_ratio set to 1 will only considder geometric gabs
87 while one set to 0 will result in only images for energy
88 resolution.
90 The AutoNEB method uses a fixed file-naming convention.
91 The initial images should have the naming prefix000.traj, prefix001.traj,
92 ... up until the final image in prefix00N.traj
93 Images are dynamically added in between the first and last image until
94 n_max images have been reached.
95 When doing the i'th NEB optimization a set of files
96 prefixXXXiter00i.traj exists with XXX ranging from 000 to the N images
97 currently in the NEB.
99 The most recent NEB path can always be monitored by:
100 $ ase-gui -n -1 neb???.traj
101 """
103 def __init__(self, attach_calculators, prefix, n_simul, n_max,
104 iter_folder='AutoNEB_iter',
105 fmax=0.025, maxsteps=10000, k=0.1, climb=True, method='eb',
106 optimizer=FIRE,
107 remove_rotation_and_translation=False, space_energy_ratio=0.5,
108 world=None,
109 parallel=True, smooth_curve=False, interpolate_method='idpp'):
110 self.attach_calculators = attach_calculators
111 self.prefix = Path(prefix)
112 self.n_simul = n_simul
113 self.n_max = n_max
114 self.climb = climb
115 self.all_images = []
117 self.parallel = parallel
118 self.maxsteps = maxsteps
119 self.fmax = fmax
120 self.k = k
121 self.method = method
122 self.remove_rotation_and_translation = remove_rotation_and_translation
123 self.space_energy_ratio = space_energy_ratio
124 if interpolate_method not in ['idpp', 'linear']:
125 self.interpolate_method = 'idpp'
126 print('Interpolation method not implementet.',
127 'Using the IDPP method.')
128 else:
129 self.interpolate_method = interpolate_method
130 if world is None:
131 world = mpi.world
132 self.world = world
133 self.smooth_curve = smooth_curve
135 if isinstance(optimizer, str):
136 warn('Please set optimizer as an object and not as string',
137 FutureWarning)
138 try:
139 self.optimizer = {
140 'BFGS': BFGS, 'FIRE': FIRE}[optimizer]
141 except KeyError:
142 raise Exception('Optimizer needs to be BFGS or FIRE')
143 else:
144 self.optimizer = optimizer
146 self.iter_folder = Path(self.prefix.parent) / iter_folder
147 self.iter_folder.mkdir(exist_ok=True)
149 def execute_one_neb(self, n_cur, to_run, climb=False, many_steps=False):
150 with ExitStack() as exitstack:
151 self._execute_one_neb(exitstack, n_cur, to_run,
152 climb=climb, many_steps=many_steps)
154 def iter_trajpath(self, i, iiter):
155 """When doing the i'th NEB optimization a set of files
156 prefixXXXiter00i.traj exists with XXX ranging from 000 to the N images
157 currently in the NEB."""
158 return (self.iter_folder /
159 (self.prefix.name + f'{i:03d}iter{iiter:03d}.traj'))
161 def _execute_one_neb(self, exitstack, n_cur, to_run,
162 climb=False, many_steps=False):
163 '''Internal method which executes one NEB optimization.'''
165 closelater = exitstack.enter_context
167 self.iteration += 1
168 # First we copy around all the images we are not using in this
169 # neb (for reproducability purposes)
170 if self.world.rank == 0:
171 for i in range(n_cur):
172 if i not in to_run[1: -1]:
173 filename = '%s%03d.traj' % (self.prefix, i)
174 with Trajectory(filename, mode='w',
175 atoms=self.all_images[i]) as traj:
176 traj.write()
177 filename_ref = self.iter_trajpath(i, self.iteration)
178 if os.path.isfile(filename):
179 shutil.copy2(filename, filename_ref)
180 if self.world.rank == 0:
181 print('Now starting iteration %d on ' % self.iteration, to_run)
182 # Attach calculators to all the images we will include in the NEB
183 self.attach_calculators([self.all_images[i] for i in to_run[1: -1]])
184 neb = NEB([self.all_images[i] for i in to_run],
185 k=[self.k[i] for i in to_run[0:-1]],
186 method=self.method,
187 parallel=self.parallel,
188 remove_rotation_and_translation=self
189 .remove_rotation_and_translation,
190 climb=climb)
192 # Do the actual NEB calculation
193 logpath = (self.iter_folder
194 / f'{self.prefix.name}_log_iter{self.iteration:03d}.log')
195 qn = closelater(self.optimizer(neb, logfile=logpath))
197 # Find the ranks which are masters for each their calculation
198 if self.parallel:
199 nneb = to_run[0]
200 nim = len(to_run) - 2
201 n = self.world.size // nim # number of cpu's per image
202 j = 1 + self.world.rank // n # my image number
203 assert nim * n == self.world.size
204 traj = closelater(Trajectory(
205 '%s%03d.traj' % (self.prefix, j + nneb), 'w',
206 self.all_images[j + nneb],
207 master=(self.world.rank % n == 0)
208 ))
209 filename_ref = self.iter_trajpath(j + nneb, self.iteration)
210 trajhist = closelater(Trajectory(
211 filename_ref, 'w',
212 self.all_images[j + nneb],
213 master=(self.world.rank % n == 0)
214 ))
215 qn.attach(traj)
216 qn.attach(trajhist)
217 else:
218 num = 1
219 for i, j in enumerate(to_run[1: -1]):
220 filename_ref = self.iter_trajpath(j, self.iteration)
221 trajhist = closelater(Trajectory(
222 filename_ref, 'w', self.all_images[j]
223 ))
224 qn.attach(seriel_writer(trajhist, i, num).write)
226 traj = closelater(Trajectory(
227 '%s%03d.traj' % (self.prefix, j), 'w',
228 self.all_images[j]
229 ))
230 qn.attach(seriel_writer(traj, i, num).write)
231 num += 1
233 if isinstance(self.maxsteps, (list, tuple)) and many_steps:
234 steps = self.maxsteps[1]
235 elif isinstance(self.maxsteps, (list, tuple)) and not many_steps:
236 steps = self.maxsteps[0]
237 else:
238 steps = self.maxsteps
240 if isinstance(self.fmax, (list, tuple)) and many_steps:
241 fmax = self.fmax[1]
242 elif isinstance(self.fmax, (list, tuple)) and not many_steps:
243 fmax = self.fmax[0]
244 else:
245 fmax = self.fmax
246 qn.run(fmax=fmax, steps=steps)
248 # Remove the calculators and replace them with single
249 # point calculators and update all the nodes for
250 # preperration for next iteration
251 neb.distribute = types.MethodType(store_E_and_F_in_spc, neb)
252 neb.distribute()
254 def run(self):
255 '''Run the AutoNEB optimization algorithm.'''
256 n_cur = self.__initialize__()
257 while len(self.all_images) < self.n_simul + 2:
258 if isinstance(self.k, (float, int)):
259 self.k = [self.k] * (len(self.all_images) - 1)
260 if self.world.rank == 0:
261 print('Now adding images for initial run')
262 # Insert a new image where the distance between two images is
263 # the largest
264 spring_lengths = []
265 for j in range(n_cur - 1):
266 spring_vec = self.all_images[j + 1].get_positions() - \
267 self.all_images[j].get_positions()
268 spring_lengths.append(np.linalg.norm(spring_vec))
269 jmax = np.argmax(spring_lengths)
271 if self.world.rank == 0:
272 print('Max length between images is at ', jmax)
274 # The interpolation used to make initial guesses
275 # If only start and end images supplied make all img at ones
276 if len(self.all_images) == 2:
277 n_between = self.n_simul
278 else:
279 n_between = 1
281 toInterpolate = [self.all_images[jmax]]
282 for i in range(n_between):
283 toInterpolate += [toInterpolate[0].copy()]
284 toInterpolate += [self.all_images[jmax + 1]]
286 neb = NEB(toInterpolate)
287 neb.interpolate(method=self.interpolate_method)
289 tmp = self.all_images[:jmax + 1]
290 tmp += toInterpolate[1:-1]
291 tmp.extend(self.all_images[jmax + 1:])
293 self.all_images = tmp
295 # Expect springs to be in equilibrium
296 k_tmp = self.k[:jmax]
297 k_tmp += [self.k[jmax] * (n_between + 1)] * (n_between + 1)
298 k_tmp.extend(self.k[jmax + 1:])
299 self.k = k_tmp
301 # Run the NEB calculation with the new image included
302 n_cur += n_between
304 # Determine if any images do not have a valid energy yet
305 energies = self.get_energies()
307 n_non_valid_energies = len([e for e in energies if e != e])
309 if self.world.rank == 0:
310 print('Start of evaluation of the initial images')
312 while n_non_valid_energies != 0:
313 if isinstance(self.k, (float, int)):
314 self.k = [self.k] * (len(self.all_images) - 1)
316 # First do one run since some energie are non-determined
317 to_run, climb_safe = self.which_images_to_run_on()
318 self.execute_one_neb(n_cur, to_run, climb=False)
320 energies = self.get_energies()
321 n_non_valid_energies = len([e for e in energies if e != e])
323 if self.world.rank == 0:
324 print('Finished initialisation phase.')
326 # Then add one image at a time until we have n_max images
327 while n_cur < self.n_max:
328 if isinstance(self.k, (float, int)):
329 self.k = [self.k] * (len(self.all_images) - 1)
330 # Insert a new image where the distance between two images
331 # is the largest OR where a higher energy reselution is needed
332 if self.world.rank == 0:
333 print('****Now adding another image until n_max is reached',
334 '({0}/{1})****'.format(n_cur, self.n_max))
335 spring_lengths = []
336 for j in range(n_cur - 1):
337 spring_vec = self.all_images[j + 1].get_positions() - \
338 self.all_images[j].get_positions()
339 spring_lengths.append(np.linalg.norm(spring_vec))
341 total_vec = self.all_images[0].get_positions() - \
342 self.all_images[-1].get_positions()
343 tl = np.linalg.norm(total_vec)
345 fR = max(spring_lengths) / tl
347 e = self.get_energies()
348 ed = []
349 emin = min(e)
350 enorm = max(e) - emin
351 for j in range(n_cur - 1):
352 delta_E = (e[j + 1] - e[j]) * (e[j + 1] + e[j] - 2 *
353 emin) / 2 / enorm
354 ed.append(abs(delta_E))
356 gR = max(ed) / enorm
358 if fR / gR > self.space_energy_ratio:
359 jmax = np.argmax(spring_lengths)
360 t = 'spring length!'
361 else:
362 jmax = np.argmax(ed)
363 t = 'energy difference between neighbours!'
365 if self.world.rank == 0:
366 print('Adding image between {0} and'.format(jmax),
367 '{0}. New image point is selected'.format(jmax + 1),
368 'on the basis of the biggest ' + t)
370 toInterpolate = [self.all_images[jmax]]
371 toInterpolate += [toInterpolate[0].copy()]
372 toInterpolate += [self.all_images[jmax + 1]]
374 neb = NEB(toInterpolate)
375 neb.interpolate(method=self.interpolate_method)
377 tmp = self.all_images[:jmax + 1]
378 tmp += toInterpolate[1:-1]
379 tmp.extend(self.all_images[jmax + 1:])
381 self.all_images = tmp
383 # Expect springs to be in equilibrium
384 k_tmp = self.k[:jmax]
385 k_tmp += [self.k[jmax] * 2] * 2
386 k_tmp.extend(self.k[jmax + 1:])
387 self.k = k_tmp
389 # Run the NEB calculation with the new image included
390 n_cur += 1
391 to_run, climb_safe = self.which_images_to_run_on()
393 self.execute_one_neb(n_cur, to_run, climb=False)
395 if self.world.rank == 0:
396 print('n_max images has been reached')
398 # Do a single climb around the top-point if requested
399 if self.climb:
400 if isinstance(self.k, (float, int)):
401 self.k = [self.k] * (len(self.all_images) - 1)
402 if self.world.rank == 0:
403 print('****Now doing the CI-NEB calculation****')
404 to_run, climb_safe = self.which_images_to_run_on()
406 assert climb_safe, 'climb_safe should be true at this point!'
407 self.execute_one_neb(n_cur, to_run, climb=True, many_steps=True)
409 if not self.smooth_curve:
410 return self.all_images
412 # If a smooth_curve is requsted ajust the springs to follow two
413 # gaussian distributions
414 e = self.get_energies()
415 peak = self.get_highest_energy_index()
416 k_max = 10
418 d1 = np.linalg.norm(self.all_images[peak].get_positions() -
419 self.all_images[0].get_positions())
420 d2 = np.linalg.norm(self.all_images[peak].get_positions() -
421 self.all_images[-1].get_positions())
422 l1 = -d1 ** 2 / log(0.2)
423 l2 = -d2 ** 2 / log(0.2)
425 x1 = []
426 x2 = []
427 for i in range(peak):
428 v = (self.all_images[i].get_positions() +
429 self.all_images[i + 1].get_positions()) / 2 - \
430 self.all_images[0].get_positions()
431 x1.append(np.linalg.norm(v))
433 for i in range(peak, len(self.all_images) - 1):
434 v = (self.all_images[i].get_positions() +
435 self.all_images[i + 1].get_positions()) / 2 - \
436 self.all_images[0].get_positions()
437 x2.append(np.linalg.norm(v))
438 k_tmp = []
439 for x in x1:
440 k_tmp.append(k_max * exp(-((x - d1) ** 2) / l1))
441 for x in x2:
442 k_tmp.append(k_max * exp(-((x - d1) ** 2) / l2))
444 self.k = k_tmp
445 # Roll back to start from the top-point
446 if self.world.rank == 0:
447 print('Now moving from top to start')
448 highest_energy_index = self.get_highest_energy_index()
449 nneb = highest_energy_index - self.n_simul - 1
450 while nneb >= 0:
451 self.execute_one_neb(n_cur, range(nneb, nneb + self.n_simul + 2),
452 climb=False)
453 nneb -= 1
455 # Roll forward from the top-point until the end
456 nneb = self.get_highest_energy_index()
458 if self.world.rank == 0:
459 print('Now moving from top to end')
460 while nneb <= self.n_max - self.n_simul - 2:
461 self.execute_one_neb(n_cur, range(nneb, nneb + self.n_simul + 2),
462 climb=False)
463 nneb += 1
464 return self.all_images
466 def __initialize__(self):
467 '''Load files from the filesystem.'''
468 if not os.path.isfile('%s000.traj' % self.prefix):
469 raise IOError('No file with name %s000.traj' % self.prefix,
470 'was found. Should contain initial image')
472 # Find the images that exist
473 index_exists = [i for i in range(self.n_max) if
474 os.path.isfile('%s%03d.traj' % (self.prefix, i))]
476 n_cur = index_exists[-1] + 1
478 if self.world.rank == 0:
479 print('The NEB initially has %d images ' % len(index_exists),
480 '(including the end-points)')
481 if len(index_exists) == 1:
482 raise Exception('Only a start point exists')
484 for i in range(len(index_exists)):
485 if i != index_exists[i]:
486 raise Exception('Files must be ordered sequentially',
487 'without gaps.')
488 if self.world.rank == 0:
489 for i in index_exists:
490 filename_ref = self.iter_trajpath(i, 0)
491 if os.path.isfile(filename_ref):
492 try:
493 os.rename(filename_ref, str(filename_ref) + '.bak')
494 except IOError:
495 pass
496 filename = '%s%03d.traj' % (self.prefix, i)
497 try:
498 shutil.copy2(filename, filename_ref)
499 except IOError:
500 pass
501 # Wait for file system on all nodes is syncronized
502 self.world.barrier()
503 # And now lets read in the configurations
504 for i in range(n_cur):
505 if i in index_exists:
506 filename = '%s%03d.traj' % (self.prefix, i)
507 newim = read(filename)
508 self.all_images.append(newim)
509 else:
510 self.all_images.append(self.all_images[0].copy())
512 self.iteration = 0
513 return n_cur
515 def get_energies(self):
516 """Utility method to extract all energies and insert np.NaN at
517 invalid images."""
518 energies = []
519 for a in self.all_images:
520 try:
521 energies.append(a.get_potential_energy())
522 except RuntimeError:
523 energies.append(np.NaN)
524 return energies
526 def get_energies_one_image(self, image):
527 """Utility method to extract energy of an image and return np.NaN
528 if invalid."""
529 try:
530 energy = image.get_potential_energy()
531 except RuntimeError:
532 energy = np.NaN
533 return energy
535 def get_highest_energy_index(self):
536 """Find the index of the image with the highest energy."""
537 energies = self.get_energies()
538 valid_entries = [(i, e) for i, e in enumerate(energies) if e == e]
539 highest_energy_index = max(valid_entries, key=lambda x: x[1])[0]
540 return highest_energy_index
542 def which_images_to_run_on(self):
543 """Determine which set of images to do a NEB at.
544 The priority is to first include all images without valid energies,
545 secondly include the highest energy image."""
546 n_cur = len(self.all_images)
547 energies = self.get_energies()
548 # Find out which image is the first one missing the energy and
549 # which is the last one missing the energy
550 first_missing = n_cur
551 last_missing = 0
552 n_missing = 0
553 for i in range(1, n_cur - 1):
554 if energies[i] != energies[i]:
555 n_missing += 1
556 first_missing = min(first_missing, i)
557 last_missing = max(last_missing, i)
559 highest_energy_index = self.get_highest_energy_index()
561 nneb = highest_energy_index - 1 - self.n_simul // 2
562 nneb = max(nneb, 0)
563 nneb = min(nneb, n_cur - self.n_simul - 2)
564 nneb = min(nneb, first_missing - 1)
565 nneb = max(nneb + self.n_simul, last_missing) - self.n_simul
566 to_use = list(range(nneb, nneb + self.n_simul + 2))
568 while self.get_energies_one_image(self.all_images[to_use[0]]) != \
569 self.get_energies_one_image(self.all_images[to_use[0]]):
570 to_use[0] -= 1
571 while self.get_energies_one_image(self.all_images[to_use[-1]]) != \
572 self.get_energies_one_image(self.all_images[to_use[-1]]):
573 to_use[-1] += 1
575 return to_use, (highest_energy_index in to_use[1: -1])
578class seriel_writer:
579 def __init__(self, traj, i, num):
580 self.traj = traj
581 self.i = i
582 self.num = num
584 def write(self):
585 if self.num % (self.i + 1) == 0:
586 self.traj.write()
589def store_E_and_F_in_spc(self):
590 """Collect the energies and forces on all nodes and store as
591 single point calculators"""
592 # Make sure energies and forces are known on all nodes
593 self.get_forces()
594 images = self.images
595 if self.parallel:
596 energy = np.empty(1)
597 forces = np.empty((self.natoms, 3))
599 for i in range(1, self.nimages - 1):
600 # Determine which node is the leading for image i
601 root = (i - 1) * self.world.size // (self.nimages - 2)
602 # If on this node, extract the calculated numbers
603 if self.world.rank == root:
604 forces = images[i].get_forces()
605 energy[0] = images[i].get_potential_energy()
606 # Distribute these numbers to other nodes
607 self.world.broadcast(energy, root)
608 self.world.broadcast(forces, root)
609 # On all nodes, remove the calculator, keep only energy
610 # and force in single point calculator
611 self.images[i].calc = SinglePointCalculator(
612 self.images[i],
613 energy=energy[0],
614 forces=forces)