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