from ase.io import Trajectory
from ase.io import read
from ase.neb import NEB
from ase.optimize import BFGS
from ase.optimize import FIRE
from ase.calculators.singlepoint import SinglePointCalculator
import ase.parallel as mpi
import numpy as np
import shutil
import os
import types
from math import log
from math import exp
from contextlib import ExitStack
[docs]class AutoNEB:
"""AutoNEB object.
The AutoNEB algorithm streamlines the execution of NEB and CI-NEB
calculations following the algorithm described in:
E. L. Kolsbjerg, M. N. Groves, and B. Hammer, J. Chem. Phys,
145, 094107, 2016. (doi: 10.1063/1.4961868)
The user supplies at minimum the two end-points and possibly also some
intermediate images.
The stages are:
1) Define a set of images and name them sequentially.
Must at least have a relaxed starting and ending image
User can supply intermediate guesses which do not need to
have previously determined energies (probably from another
NEB calculation with a lower level of theory)
2) AutoNEB will first evaluate the user provided intermediate images
3) AutoNEB will then add additional images dynamically until n_max
is reached
4) A climbing image will attempt to locate the saddle point
5) All the images between the highest point and the starting point
are further relaxed to smooth the path
6) All the images between the highest point and the ending point are
further relaxed to smooth the path
Step 4 and 5-6 are optional steps!
Parameters:
attach_calculators:
Function which adds valid calculators to the list of images supplied.
prefix: string
All files that the AutoNEB method reads and writes are prefixed with
this string
n_simul: int
The number of relaxations run in parallel.
n_max: int
The number of images along the NEB path when done.
This number includes the two end-points.
Important: due to the dynamic adding of images around the peak n_max
must be updated if the NEB is restarted.
climb: boolean
Should a CI-NEB calculation be done at the top-point
fmax: float or list of floats
The maximum force along the NEB path
maxsteps: int
The maximum number of steps in each NEB relaxation.
If a list is given the first number of steps is used in the build-up
and final scan phase;
the second number of steps is used in the CI step after all images
have been inserted.
k: float
The spring constant along the NEB path
method: str (see neb.py)
Choice betweeen three method:
'aseneb', standard ase NEB implementation
'improvedtangent', published NEB implementation
'eb', full spring force implementation (default)
optimizer: str
Which optimizer to use in the relaxation. Valid values are 'BFGS'
and 'FIRE' (default)
space_energy_ratio: float
The preference for new images to be added in a big energy gab
with a preference around the peak or in the biggest geometric gab.
A space_energy_ratio set to 1 will only considder geometric gabs
while one set to 0 will result in only images for energy
resolution.
The AutoNEB method uses a fixed file-naming convention.
The initial images should have the naming prefix000.traj, prefix001.traj,
... up until the final image in prefix00N.traj
Images are dynamically added in between the first and last image until
n_max images have been reached.
When doing the i'th NEB optimization a set of files
prefixXXXiter00i.traj exists with XXX ranging from 000 to the N images
currently in the NEB.
The most recent NEB path can always be monitored by:
$ ase-gui -n -1 neb???.traj
"""
def __init__(self, attach_calculators, prefix, n_simul, n_max,
iter_folder='AutoNEB_iter',
fmax=0.025, maxsteps=10000, k=0.1, climb=True, method='eb',
optimizer='FIRE',
remove_rotation_and_translation=False, space_energy_ratio=0.5,
world=None,
parallel=True, smooth_curve=False, interpolate_method='idpp'):
self.attach_calculators = attach_calculators
self.prefix = prefix
self.n_simul = n_simul
self.n_max = n_max
self.climb = climb
self.all_images = []
self.parallel = parallel
self.maxsteps = maxsteps
self.fmax = fmax
self.k = k
self.method = method
self.remove_rotation_and_translation = remove_rotation_and_translation
self.space_energy_ratio = space_energy_ratio
if interpolate_method not in ['idpp', 'linear']:
self.interpolate_method = 'idpp'
print('Interpolation method not implementet.',
'Using the IDPP method.')
else:
self.interpolate_method = interpolate_method
if world is None:
world = mpi.world
self.world = world
self.smooth_curve = smooth_curve
if optimizer == 'BFGS':
self.optimizer = BFGS
elif optimizer == 'FIRE':
self.optimizer = FIRE
else:
raise Exception('Optimizer needs to be BFGS or FIRE')
self.iter_folder = iter_folder
if not os.path.exists(self.iter_folder) and self.world.rank == 0:
os.makedirs(self.iter_folder)
def execute_one_neb(self, n_cur, to_run, climb=False, many_steps=False):
with ExitStack() as exitstack:
self._execute_one_neb(exitstack, n_cur, to_run,
climb=climb, many_steps=many_steps)
def _execute_one_neb(self, exitstack, n_cur, to_run,
climb=False, many_steps=False):
'''Internal method which executes one NEB optimization.'''
closelater = exitstack.enter_context
self.iteration += 1
# First we copy around all the images we are not using in this
# neb (for reproducability purposes)
if self.world.rank == 0:
for i in range(n_cur):
if i not in to_run[1: -1]:
filename = '%s%03d.traj' % (self.prefix, i)
with Trajectory(filename, mode='w',
atoms=self.all_images[i]) as traj:
traj.write()
filename_ref = self.iter_folder + \
'/%s%03diter%03d.traj' % (self.prefix, i,
self.iteration)
if os.path.isfile(filename):
shutil.copy2(filename, filename_ref)
if self.world.rank == 0:
print('Now starting iteration %d on ' % self.iteration, to_run)
# Attach calculators to all the images we will include in the NEB
self.attach_calculators([self.all_images[i] for i in to_run[1: -1]])
neb = NEB([self.all_images[i] for i in to_run],
k=[self.k[i] for i in to_run[0:-1]],
method=self.method,
parallel=self.parallel,
remove_rotation_and_translation=self
.remove_rotation_and_translation,
climb=climb)
# Do the actual NEB calculation
qn = closelater(
self.optimizer(neb,
logfile=self.iter_folder +
'/%s_log_iter%03d.log' % (self.prefix,
self.iteration))
)
# Find the ranks which are masters for each their calculation
if self.parallel:
nneb = to_run[0]
nim = len(to_run) - 2
n = self.world.size // nim # number of cpu's per image
j = 1 + self.world.rank // n # my image number
assert nim * n == self.world.size
traj = closelater(Trajectory(
'%s%03d.traj' % (self.prefix, j + nneb), 'w',
self.all_images[j + nneb],
master=(self.world.rank % n == 0)
))
filename_ref = self.iter_folder + \
'/%s%03diter%03d.traj' % (self.prefix,
j + nneb, self.iteration)
trajhist = closelater(Trajectory(
filename_ref, 'w',
self.all_images[j + nneb],
master=(self.world.rank % n == 0)
))
qn.attach(traj)
qn.attach(trajhist)
else:
num = 1
for i, j in enumerate(to_run[1: -1]):
filename_ref = self.iter_folder + \
'/%s%03diter%03d.traj' % (self.prefix, j, self.iteration)
trajhist = closelater(Trajectory(
filename_ref, 'w', self.all_images[j]
))
qn.attach(seriel_writer(trajhist, i, num).write)
traj = closelater(Trajectory(
'%s%03d.traj' % (self.prefix, j), 'w',
self.all_images[j]
))
qn.attach(seriel_writer(traj, i, num).write)
num += 1
if isinstance(self.maxsteps, (list, tuple)) and many_steps:
steps = self.maxsteps[1]
elif isinstance(self.maxsteps, (list, tuple)) and not many_steps:
steps = self.maxsteps[0]
else:
steps = self.maxsteps
if isinstance(self.fmax, (list, tuple)) and many_steps:
fmax = self.fmax[1]
elif isinstance(self.fmax, (list, tuple)) and not many_steps:
fmax = self.fmax[0]
else:
fmax = self.fmax
qn.run(fmax=fmax, steps=steps)
# Remove the calculators and replace them with single
# point calculators and update all the nodes for
# preperration for next iteration
neb.distribute = types.MethodType(store_E_and_F_in_spc, neb)
neb.distribute()
def run(self):
'''Run the AutoNEB optimization algorithm.'''
n_cur = self.__initialize__()
while len(self.all_images) < self.n_simul + 2:
if isinstance(self.k, (float, int)):
self.k = [self.k] * (len(self.all_images) - 1)
if self.world.rank == 0:
print('Now adding images for initial run')
# Insert a new image where the distance between two images is
# the largest
spring_lengths = []
for j in range(n_cur - 1):
spring_vec = self.all_images[j + 1].get_positions() - \
self.all_images[j].get_positions()
spring_lengths.append(np.linalg.norm(spring_vec))
jmax = np.argmax(spring_lengths)
if self.world.rank == 0:
print('Max length between images is at ', jmax)
# The interpolation used to make initial guesses
# If only start and end images supplied make all img at ones
if len(self.all_images) == 2:
n_between = self.n_simul
else:
n_between = 1
toInterpolate = [self.all_images[jmax]]
for i in range(n_between):
toInterpolate += [toInterpolate[0].copy()]
toInterpolate += [self.all_images[jmax + 1]]
neb = NEB(toInterpolate)
neb.interpolate(method=self.interpolate_method)
tmp = self.all_images[:jmax + 1]
tmp += toInterpolate[1:-1]
tmp.extend(self.all_images[jmax + 1:])
self.all_images = tmp
# Expect springs to be in equilibrium
k_tmp = self.k[:jmax]
k_tmp += [self.k[jmax] * (n_between + 1)] * (n_between + 1)
k_tmp.extend(self.k[jmax + 1:])
self.k = k_tmp
# Run the NEB calculation with the new image included
n_cur += n_between
# Determine if any images do not have a valid energy yet
energies = self.get_energies()
n_non_valid_energies = len([e for e in energies if e != e])
if self.world.rank == 0:
print('Start of evaluation of the initial images')
while n_non_valid_energies != 0:
if isinstance(self.k, (float, int)):
self.k = [self.k] * (len(self.all_images) - 1)
# First do one run since some energie are non-determined
to_run, climb_safe = self.which_images_to_run_on()
self.execute_one_neb(n_cur, to_run, climb=False)
energies = self.get_energies()
n_non_valid_energies = len([e for e in energies if e != e])
if self.world.rank == 0:
print('Finished initialisation phase.')
# Then add one image at a time until we have n_max images
while n_cur < self.n_max:
if isinstance(self.k, (float, int)):
self.k = [self.k] * (len(self.all_images) - 1)
# Insert a new image where the distance between two images
# is the largest OR where a higher energy reselution is needed
if self.world.rank == 0:
print('****Now adding another image until n_max is reached',
'({0}/{1})****'.format(n_cur, self.n_max))
spring_lengths = []
for j in range(n_cur - 1):
spring_vec = self.all_images[j + 1].get_positions() - \
self.all_images[j].get_positions()
spring_lengths.append(np.linalg.norm(spring_vec))
total_vec = self.all_images[0].get_positions() - \
self.all_images[-1].get_positions()
tl = np.linalg.norm(total_vec)
fR = max(spring_lengths) / tl
e = self.get_energies()
ed = []
emin = min(e)
enorm = max(e) - emin
for j in range(n_cur - 1):
delta_E = (e[j + 1] - e[j]) * (e[j + 1] + e[j] - 2 *
emin) / 2 / enorm
ed.append(abs(delta_E))
gR = max(ed) / enorm
if fR / gR > self.space_energy_ratio:
jmax = np.argmax(spring_lengths)
t = 'spring length!'
else:
jmax = np.argmax(ed)
t = 'energy difference between neighbours!'
if self.world.rank == 0:
print('Adding image between {0} and'.format(jmax),
'{0}. New image point is selected'.format(jmax + 1),
'on the basis of the biggest ' + t)
toInterpolate = [self.all_images[jmax]]
toInterpolate += [toInterpolate[0].copy()]
toInterpolate += [self.all_images[jmax + 1]]
neb = NEB(toInterpolate)
neb.interpolate(method=self.interpolate_method)
tmp = self.all_images[:jmax + 1]
tmp += toInterpolate[1:-1]
tmp.extend(self.all_images[jmax + 1:])
self.all_images = tmp
# Expect springs to be in equilibrium
k_tmp = self.k[:jmax]
k_tmp += [self.k[jmax] * 2] * 2
k_tmp.extend(self.k[jmax + 1:])
self.k = k_tmp
# Run the NEB calculation with the new image included
n_cur += 1
to_run, climb_safe = self.which_images_to_run_on()
self.execute_one_neb(n_cur, to_run, climb=False)
if self.world.rank == 0:
print('n_max images has been reached')
# Do a single climb around the top-point if requested
if self.climb:
if isinstance(self.k, (float, int)):
self.k = [self.k] * (len(self.all_images) - 1)
if self.world.rank == 0:
print('****Now doing the CI-NEB calculation****')
to_run, climb_safe = self.which_images_to_run_on()
assert climb_safe, 'climb_safe should be true at this point!'
self.execute_one_neb(n_cur, to_run, climb=True, many_steps=True)
if not self.smooth_curve:
return self.all_images
# If a smooth_curve is requsted ajust the springs to follow two
# gaussian distributions
e = self.get_energies()
peak = self.get_highest_energy_index()
k_max = 10
d1 = np.linalg.norm(self.all_images[peak].get_positions() -
self.all_images[0].get_positions())
d2 = np.linalg.norm(self.all_images[peak].get_positions() -
self.all_images[-1].get_positions())
l1 = -d1 ** 2 / log(0.2)
l2 = -d2 ** 2 / log(0.2)
x1 = []
x2 = []
for i in range(peak):
v = (self.all_images[i].get_positions() +
self.all_images[i + 1].get_positions()) / 2 - \
self.all_images[0].get_positions()
x1.append(np.linalg.norm(v))
for i in range(peak, len(self.all_images) - 1):
v = (self.all_images[i].get_positions() +
self.all_images[i + 1].get_positions()) / 2 - \
self.all_images[0].get_positions()
x2.append(np.linalg.norm(v))
k_tmp = []
for x in x1:
k_tmp.append(k_max * exp(-((x - d1) ** 2) / l1))
for x in x2:
k_tmp.append(k_max * exp(-((x - d1) ** 2) / l2))
self.k = k_tmp
# Roll back to start from the top-point
if self.world.rank == 0:
print('Now moving from top to start')
highest_energy_index = self.get_highest_energy_index()
nneb = highest_energy_index - self.n_simul - 1
while nneb >= 0:
self.execute_one_neb(n_cur, range(nneb, nneb + self.n_simul + 2),
climb=False)
nneb -= 1
# Roll forward from the top-point until the end
nneb = self.get_highest_energy_index()
if self.world.rank == 0:
print('Now moving from top to end')
while nneb <= self.n_max - self.n_simul - 2:
self.execute_one_neb(n_cur, range(nneb, nneb + self.n_simul + 2),
climb=False)
nneb += 1
return self.all_images
def __initialize__(self):
'''Load files from the filesystem.'''
if not os.path.isfile('%s000.traj' % self.prefix):
raise IOError('No file with name %s000.traj' % self.prefix,
'was found. Should contain initial image')
# Find the images that exist
index_exists = [i for i in range(self.n_max) if
os.path.isfile('%s%03d.traj' % (self.prefix, i))]
n_cur = index_exists[-1] + 1
if self.world.rank == 0:
print('The NEB initially has %d images ' % len(index_exists),
'(including the end-points)')
if len(index_exists) == 1:
raise Exception('Only a start point exists')
for i in range(len(index_exists)):
if i != index_exists[i]:
raise Exception('Files must be ordered sequentially',
'without gaps.')
if self.world.rank == 0:
for i in index_exists:
filename_ref = self.iter_folder + \
'/%s%03diter000.traj' % (self.prefix, i)
if os.path.isfile(filename_ref):
try:
os.rename(filename_ref, filename_ref + '.bak')
except IOError:
pass
filename = '%s%03d.traj' % (self.prefix, i)
try:
shutil.copy2(filename, filename_ref)
except IOError:
pass
# Wait for file system on all nodes is syncronized
self.world.barrier()
# And now lets read in the configurations
for i in range(n_cur):
if i in index_exists:
filename = '%s%03d.traj' % (self.prefix, i)
newim = read(filename)
self.all_images.append(newim)
else:
self.all_images.append(self.all_images[0].copy())
self.iteration = 0
return n_cur
def get_energies(self):
"""Utility method to extract all energies and insert np.NaN at
invalid images."""
energies = []
for a in self.all_images:
try:
energies.append(a.get_potential_energy())
except RuntimeError:
energies.append(np.NaN)
return energies
def get_energies_one_image(self, image):
"""Utility method to extract energy of an image and return np.NaN
if invalid."""
try:
energy = image.get_potential_energy()
except RuntimeError:
energy = np.NaN
return energy
def get_highest_energy_index(self):
"""Find the index of the image with the highest energy."""
energies = self.get_energies()
valid_entries = [(i, e) for i, e in enumerate(energies) if e == e]
highest_energy_index = max(valid_entries, key=lambda x: x[1])[0]
return highest_energy_index
def which_images_to_run_on(self):
"""Determine which set of images to do a NEB at.
The priority is to first include all images without valid energies,
secondly include the highest energy image."""
n_cur = len(self.all_images)
energies = self.get_energies()
# Find out which image is the first one missing the energy and
# which is the last one missing the energy
first_missing = n_cur
last_missing = 0
n_missing = 0
for i in range(1, n_cur - 1):
if energies[i] != energies[i]:
n_missing += 1
first_missing = min(first_missing, i)
last_missing = max(last_missing, i)
highest_energy_index = self.get_highest_energy_index()
nneb = highest_energy_index - 1 - self.n_simul // 2
nneb = max(nneb, 0)
nneb = min(nneb, n_cur - self.n_simul - 2)
nneb = min(nneb, first_missing - 1)
nneb = max(nneb + self.n_simul, last_missing) - self.n_simul
to_use = range(nneb, nneb + self.n_simul + 2)
while self.get_energies_one_image(self.all_images[to_use[0]]) != \
self.get_energies_one_image(self.all_images[to_use[0]]):
to_use[0] -= 1
while self.get_energies_one_image(self.all_images[to_use[-1]]) != \
self.get_energies_one_image(self.all_images[to_use[-1]]):
to_use[-1] += 1
return to_use, (highest_energy_index in to_use[1: -1])
class seriel_writer:
def __init__(self, traj, i, num):
self.traj = traj
self.i = i
self.num = num
def write(self):
if self.num % (self.i + 1) == 0:
self.traj.write()
def store_E_and_F_in_spc(self):
"""Collect the energies and forces on all nodes and store as
single point calculators"""
# Make sure energies and forces are known on all nodes
self.get_forces()
images = self.images
if self.parallel:
energy = np.empty(1)
forces = np.empty((self.natoms, 3))
for i in range(1, self.nimages - 1):
# Determine which node is the leading for image i
root = (i - 1) * self.world.size // (self.nimages - 2)
# If on this node, extract the calculated numbers
if self.world.rank == root:
energy[0] = images[i].get_potential_energy()
forces = images[i].get_forces()
# Distribute these numbers to other nodes
self.world.broadcast(energy, root)
self.world.broadcast(forces, root)
# On all nodes, remove the calculator, keep only energy
# and force in single point calculator
self.images[i].calc = SinglePointCalculator(
self.images[i],
energy=energy[0],
forces=forces)