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