init

electrochemistry/echem/neb/autoneb.py

@@ -0,0 +1,651 @@
+"""
+AutoNEB realization from ASE package
+E. L. Kolsbjerg, M. N. Groves, and B. Hammer, J. Chem. Phys, 145, 094107, 2016. (doi: 10.1063/1.4961868)
+modified by: Sergey Pavlov
+"""
+
+import numpy as np
+import shutil
+import os
+import types
+from math import log
+from math import exp
+from contextlib import ExitStack
+from pathlib import Path
+from warnings import warn
+
+from ase.io import Trajectory
+from ase.io import read
+from ase.mep.neb import NEB
+from ase.optimize import BFGS
+from ase.optimize import FIRE
+from ase.mep.neb import NEBOptimizer
+from ase.calculators.singlepoint import SinglePointCalculator
+import ase.parallel as mpi
+
+import functools
+print = functools.partial(print, flush=True)
+
+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 or path
+        All files that the AutoNEB method reads and writes are prefixed with
+        prefix
+    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: object
+        Optimizer object, defaults to FIRE
+        Use of the valid strings 'BFGS' and 'FIRE' is deprecated.
+    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.
+    """
+
+    def __init__(self, attach_calculators, prefix, n_simul, n_max,
+                 iter_folder='iterations',
+                 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 = Path(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 isinstance(optimizer, str):
+            try:
+                self.optimizer = {
+                    'BFGS': BFGS, 'FIRE': FIRE, 'NEB': NEBOptimizer}[optimizer]
+            except KeyError:
+                raise Exception('Optimizer needs to be BFGS, FIRE or NEB')
+        else:
+            self.optimizer = optimizer
+
+        self.iter_folder = Path(self.prefix) / iter_folder
+        self.iter_folder.mkdir(exist_ok=True)
+
+    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 iter_trajpath(self, i, iiter):
+        """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."""
+        (self.iter_folder / f'iter_{iiter}').mkdir(exist_ok=True)
+        return self.iter_folder / f'iter_{iiter}' / f'i{i:03d}iter{iiter:03d}.traj'
+
+    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 = self.prefix / f'{i:03d}.traj'
+                    with Trajectory(filename, mode='w',
+                                    atoms=self.all_images[i]) as traj:
+                        traj.write()
+                    filename_ref = self.iter_trajpath(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], to_run, self.iteration)
+        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
+        logpath = (self.iter_folder / f'iter_{self.iteration}'
+                   / f'log_iter{self.iteration:03d}.log')
+        qn = closelater(self.optimizer(neb, logfile=logpath))
+
+        # 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(
+                self.prefix / f'{j + nneb:03d}.traj', 'w',
+                self.all_images[j + nneb],
+                master=(self.world.rank % n == 0)
+            ))
+            filename_ref = self.iter_trajpath(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_trajpath(j, self.iteration)
+                trajhist = closelater(Trajectory(
+                    filename_ref, 'w', self.all_images[j]
+                ))
+                qn.attach(seriel_writer(trajhist, i, num).write)
+
+                traj = closelater(Trajectory(
+                    self.prefix / f'{j:03d}.traj', '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()
+        energies = self.get_energies()
+        print(f'Energies after iteration {self.iteration}: {energies}')
+
+    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 energies 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)
+                for i in range(n_cur):
+                    if i <= jmax:
+                        folder_from = self.iter_folder / f'iter_{self.iteration}' / f'{i}'
+                        folder_to = self.iter_folder / f'iter_{self.iteration + 1}' / f'{i}'
+                    else:
+                        folder_from = self.iter_folder / f'iter_{self.iteration}' / f'{i}'
+                        folder_to = self.iter_folder / f'iter_{self.iteration + 1}' / f'{i + 1}'
+                    (self.iter_folder / f'iter_{self.iteration + 1}').mkdir(exist_ok=True)
+                    shutil.copytree(folder_from, folder_to, dirs_exist_ok=True)
+
+            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****')
+
+                for i in range(n_cur):
+                    folder_from = self.iter_folder / f'iter_{self.iteration}' / f'{i}'
+                    folder_to = self.iter_folder / f'iter_{self.iteration + 1}' / f'{i}'
+                    (self.iter_folder / f'iter_{self.iteration + 1}').mkdir(exist_ok=True)
+                    shutil.copytree(folder_from, folder_to, dirs_exist_ok=True)
+
+            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)
+            to_run = list(range(nneb, nneb + self.n_simul + 2))
+
+            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')
+
+            for i in range(n_cur):
+                folder_from = self.iter_folder / f'iter_{self.iteration}' / f'{i}'
+                folder_to = self.iter_folder / f'iter_{self.iteration + 1}' / f'{i}'
+                (self.iter_folder / f'iter_{self.iteration + 1}').mkdir(exist_ok=True)
+                shutil.copytree(folder_from, folder_to, dirs_exist_ok=True)
+
+
+        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
+        energies = self.get_energies()
+        print(f'Energies after iteration {self.iteration}: {energies}')
+        return self.all_images
+
+
+
+    def __initialize__(self):
+        '''Load files from the filesystem.'''
+        if not os.path.isfile(self.prefix / '000.traj'):
+            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(self.prefix / f'{i:03d}.traj')]
+        print(f'Traj files with the following indexes were initially found: {index_exists}')
+        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_trajpath(i, 0)
+                if os.path.isfile(filename_ref):
+                    try:
+                        os.rename(filename_ref, str(filename_ref) + '.bak')
+                    except IOError:
+                        pass
+                filename = self.prefix / f'{i:03d}.traj'
+                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 = self.prefix / f'{i:03d}.traj'
+                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)
+
+        # If all images missing the energy
+        if last_missing - first_missing + 1 == n_cur - 2:
+            return list(range(0, n_cur)), False
+
+        # Other options
+        else:
+            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 = list(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:
+                forces = images[i].get_forces()
+                energy[0] = images[i].get_potential_energy()
+            # 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)

electrochemistry/echem/neb/calculators.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+import tempfile
+import numpy as np
+from ase.calculators.calculator import Calculator
+from echem.core.constants import Hartree2eV, Angstrom2Bohr, Bohr2Angstrom
+from echem.core.useful_funcs import shell
+from pathlib import Path
+import logging
+
+
+# Atomistic Simulation Environment (ASE) calculator interface for JDFTx
+# See http://jdftx.org for JDFTx and https://wiki.fysik.dtu.dk/ase/ for ASE
+# Authors: Deniz Gunceler, Ravishankar Sundararaman
+# Modified: Vitaliy Kislenko
+class JDFTx(Calculator):
+    def __init__(self,
+                 path_jdftx_executable: str | Path,
+                 path_rundir: str | Path | None = None,
+                 commands: list[tuple[str, str]] = None,
+                 jdftx_prefix: str = 'jdft',
+                 output_name: str = 'output.out'):
+
+        self.logger = logging.getLogger(self.__class__.__name__ + ':')
+
+        if isinstance(path_jdftx_executable, str):
+            self.path_jdftx_executable = Path(path_jdftx_executable)
+        else:
+            self.path_jdftx_executable = path_jdftx_executable
+
+        if isinstance(path_rundir, str):
+            self.path_rundir = Path(path_rundir)
+        elif isinstance(path_rundir, Path):
+            self.path_rundir = path_rundir
+        elif path_rundir is None:
+            self.path_rundir = Path(tempfile.mkdtemp())
+        else:
+            raise ValueError(f'path_rundir should be str or Path or None, however you set {path_rundir=}'
+                             f' with type {type(path_rundir)}')
+
+        self.jdftx_prefix = jdftx_prefix
+        self.output_name = output_name
+
+        self.dumps = []
+        self.input = [('dump-name', f'{self.jdftx_prefix}.$VAR'),
+                      ('initial-state', f'{self.jdftx_prefix}.$VAR')]
+
+        if commands is not None:
+            for com, val in commands:
+                if com == 'dump-name':
+                    self.logger.debug(f'{self.path_rundir} You set \'dump-name\' command in commands = \'{val}\', '
+                                      f'however it will be replaced with \'{self.jdftx_prefix}.$VAR\'')
+                elif com == 'initial-state':
+                    self.logger.debug(f'{self.path_rundir} You set \'initial-state\' command in commands = \'{val}\', '
+                                      f'however it will be replaced with \'{self.jdftx_prefix}.$VAR\'')
+                elif com == 'coords-type':
+                    self.logger.debug(f'{self.path_rundir} You set \'coords-type\' command in commands = \'{val}\', '
+                                      f'however it will be replaced with \'cartesian\'')
+                elif com == 'include':
+                    self.logger.debug(f'{self.path_rundir} \'include\' command is not supported, ignore it')
+                elif com == 'coulomb-interaction':
+                    self.logger.debug(f'{self.path_rundir} \'coulomb-interaction\' command will be replaced in accordance with ase atoms')
+                elif com == 'dump':
+                    self.addDump(val.split()[0], val.split()[1])
+                else:
+                    self.addCommand(com, val)
+
+        if ('End', 'State') not in self.dumps:
+            self.addDump("End", "State")
+        if ('End', 'Forces') not in self.dumps:
+            self.addDump("End", "Forces")
+        if ('End', 'Ecomponents') not in self.dumps:
+            self.addDump("End", "Ecomponents")
+
+        # Current results
+        self.E = None
+        self.forces = None
+
+        # History
+        self.lastAtoms = None
+        self.lastInput = None
+
+        self.global_step = None
+
+        self.logger.debug(f'Successfully initialized JDFTx calculator in \'{self.path_rundir}\'')
+
+    def validCommand(self, command) -> bool:
+        """Checks whether the input string is a valid jdftx command by comparing to the input template (jdft -t)"""
+        if type(command) != str:
+            raise IOError('Please enter a string as the name of the command!\n')
+        return True
+
+    def addCommand(self, cmd, v) -> None:
+        if not self.validCommand(cmd):
+            raise IOError(f'{cmd} is not a valid JDFTx command!\n'
+                          'Look at the input file template (jdftx -t) for a list of commands.')
+        self.input.append((cmd, v))
+
+    def addDump(self, when, what) -> None:
+        self.dumps.append((when, what))
+
+    def __readEnergy(self,
+                     filepath: str | Path) -> float:
+        Efinal = None
+        for line in open(filepath):
+            tokens = line.split()
+            if len(tokens) == 3:
+                Efinal = float(tokens[2])
+        if Efinal is None:
+            raise IOError('Error: Energy not found.')
+        return Efinal * Hartree2eV  # Return energy from final line (Etot, F or G)
+
+    def __readForces(self,
+                     filepath: str | Path) -> np.array:
+        idxMap = {}
+        symbolList = self.lastAtoms.get_chemical_symbols()
+        for i, symbol in enumerate(symbolList):
+            if symbol not in idxMap:
+                idxMap[symbol] = []
+            idxMap[symbol].append(i)
+        forces = [0] * len(symbolList)
+        for line in open(filepath):
+            if line.startswith('force '):
+                tokens = line.split()
+                idx = idxMap[tokens[1]].pop(0)  # tokens[1] is chemical symbol
+                forces[idx] = [float(word) for word in tokens[2:5]]  # tokens[2:5]: force components
+
+        if len(forces) == 0:
+            raise IOError('Error: Forces not found.')
+        return (Hartree2eV / Bohr2Angstrom) * np.array(forces)
+
+    def calculation_required(self, atoms, quantities) -> bool:
+        if (self.E is None) or (self.forces is None):
+            return True
+        if (self.lastAtoms != atoms) or (self.input != self.lastInput):
+            return True
+        return False
+
+    def get_forces(self, atoms) -> np.array:
+        if self.calculation_required(atoms, None):
+            self.update(atoms)
+        return self.forces
+
+    def get_potential_energy(self, atoms, force_consistent=False):
+        if self.calculation_required(atoms, None):
+            self.update(atoms)
+        return self.E
+
+    def update(self, atoms):
+        self.runJDFTx(self.constructInput(atoms))
+
+    def runJDFTx(self, inputfile):
+        """ Runs a JDFTx calculation """
+        file = open(self.path_rundir / 'input.in', 'w')
+        file.write(inputfile)
+        file.close()
+
+        if self.global_step is not None:
+            self.logger.info(f'Step: {self.global_step:2}. Run in {self.path_rundir}')
+        else:
+            self.logger.info(f'Run in {self.path_rundir}')
+
+        shell(f'cd {self.path_rundir} && srun {self.path_jdftx_executable} -i input.in -o {self.output_name}')
+
+        self.E = self.__readEnergy(self.path_rundir / f'{self.jdftx_prefix}.Ecomponents')
+
+        if self.global_step is not None:
+            self.logger.debug(f'Step: {self.global_step}. E = {self.E:.4f}')
+        else:
+            self.logger.debug(f'E = {self.E:.4f}')
+
+        self.forces = self.__readForces(self.path_rundir / f'{self.jdftx_prefix}.force')
+
+    def constructInput(self, atoms) -> str:
+        """Constructs a JDFTx input string using the input atoms and the input file arguments (kwargs) in self.input"""
+        inputfile = ''
+
+        lattice = atoms.get_cell() * Angstrom2Bohr
+        inputfile += 'lattice \\\n'
+        for i in range(3):
+            for j in range(3):
+                inputfile += '%f  ' % (lattice[j, i])
+            if i != 2:
+                inputfile += '\\'
+            inputfile += '\n'
+
+        inputfile += '\n'
+
+        inputfile += "".join(["dump %s %s\n" % (when, what) for when, what in self.dumps])
+
+        inputfile += '\n'
+        for cmd, v in self.input:
+            inputfile += '%s %s\n' % (cmd, str(v))
+
+        coords = [x * Angstrom2Bohr for x in list(atoms.get_positions())]
+        species = atoms.get_chemical_symbols()
+        inputfile += '\ncoords-type cartesian\n'
+        for i in range(len(coords)):
+            inputfile += 'ion %s %f %f %f \t 1\n' % (species[i], coords[i][0], coords[i][1], coords[i][2])
+
+        inputfile += '\ncoulomb-interaction '
+        pbc = list(atoms.get_pbc())
+        if sum(pbc) == 3:
+            inputfile += 'periodic\n'
+        elif sum(pbc) == 0:
+            inputfile += 'isolated\n'
+        elif sum(pbc) == 1:
+            inputfile += 'wire %i%i%i\n' % (pbc[0], pbc[1], pbc[2])
+        elif sum(pbc) == 2:
+            inputfile += 'slab %i%i%i\n' % (not pbc[0], not pbc[1], not pbc[2])
+        # --- add truncation center:
+        if sum(pbc) < 3:
+            center = np.mean(np.array(coords), axis=0)
+            inputfile += 'coulomb-truncation-embed %g %g %g\n' % tuple(center.tolist())
+
+        # Cache this calculation to history
+        self.lastAtoms = atoms.copy()
+        self.lastInput = list(self.input)
+
+        return inputfile

electrochemistry/echem/neb/neb.py

@@ -0,0 +1,587 @@
+from __future__ import annotations
+from ase.mep.neb import NEB
+from ase.optimize.sciopt import OptimizerConvergenceError
+from ase.io.trajectory import Trajectory, TrajectoryWriter
+from ase.io import read
+from echem.neb.calculators import JDFTx
+from echem.neb.autoneb import AutoNEB
+from echem.io_data.jdftx import Ionpos, Lattice, Input
+from echem.core.useful_funcs import shell
+from pathlib import Path
+import numpy as np
+import logging
+import os
+from typing import Literal, Callable
+logging.basicConfig(filename='logfile_NEB.log', filemode='a', level=logging.INFO,
+                    format="%(asctime)s %(levelname)8s %(name)14s %(message)s",
+                    datefmt='%d/%m/%Y %H:%M:%S')
+
+
+class NEBOptimizer:
+    def __init__(self,
+                 neb: NEB,
+                 trajectory_filepath: str | Path | None = None,
+                 append_trajectory: bool = True):
+
+        self.neb = neb
+        self.logger = logging.getLogger(self.__class__.__name__ + ':')
+        self.E_image_first = None
+        self.E_image_last = None
+
+        if trajectory_filepath is not None:
+            if append_trajectory:
+                self.trj_writer = TrajectoryWriter(trajectory_filepath, mode='a')
+            else:
+                self.trj_writer = TrajectoryWriter(trajectory_filepath, mode='w')
+        else:
+            self.trj_writer = None
+
+    def converged(self, fmax):
+        return self.neb.get_residual() <= fmax
+
+    def update_positions(self, X):
+        positions = X.reshape((self.neb.nimages - 2) * self.neb.natoms, 3)
+        self.neb.set_positions(positions)
+
+    def get_forces(self):
+        return self.neb.get_forces().reshape(-1)
+
+    def get_energies(self, first: bool = False, last: bool = False):
+        if not first and not last:
+            return [image.calc.E for image in self.neb.images[1:-1]]
+        elif first and not last:
+            return [image.calc.E for image in self.neb.images[:-1]]
+        elif not first and last:
+            return [image.calc.E for image in self.neb.images[1:]]
+        elif first and last:
+            return [image.calc.E for image in self.neb.images]
+
+    def dump_trajectory(self):
+        if self.trj_writer is not None:
+            for image in self.neb.images:
+                self.trj_writer.write(image)
+
+    def dump_positions_vasp(self, prefix='last_img'):
+        length = len(str(self.neb.nimages + 1))
+        for i, image in enumerate(self.neb.images):
+            image.write(f'{prefix}_{str(i).zfill(length)}.vasp', format='vasp')
+
+    def set_step_in_calculators(self, step, first: bool = False, last: bool = False):
+        if not first and not last:
+            for image in self.neb.images[1:-1]:
+                image.calc.global_step = step
+        elif first and not last:
+            for image in self.neb.images[:-1]:
+                image.calc.global_step = step
+        elif not first and last:
+            for image in self.neb.images[1:]:
+                image.calc.global_step = step
+        elif first and last:
+            for image in self.neb.images:
+                image.calc.global_step = step
+
+    def run_static(self,
+                   fmax: float = 0.1,
+                   max_steps: int = 100,
+                   alpha: float = 0.02,
+                   dE_max: float = None,
+                   construct_calc_fn: Callable = None):
+        self.logger.info('Static method of optimization was chosen')
+        max_new_images_at_step = 1
+        min_steps_after_insertion = 3
+        steps_after_insertion = 0
+
+        if dE_max is not None:
+            self.logger.info(f'AutoNEB with max {dE_max} eV difference between images was set')
+            self.logger.info(f'Initial number of images is {self.neb.nimages}, '
+                             f'including initial and final images')
+
+        if max_steps < 1:
+            raise ValueError('max_steps must be greater or equal than one')
+
+        if dE_max is not None:
+            self.set_step_in_calculators(0, first=True, last=True)
+            self.E_image_first = self.neb.images[0].get_potential_energy()
+            self.E_image_last = self.neb.images[-1].get_potential_energy()
+
+        length_step = len(str(max_steps))
+        #X = self.neb.get_positions().reshape(-1)
+        for step in range(max_steps):
+            self.dump_trajectory()
+            self.dump_positions_vasp(prefix=f'Step-{step}-1-')
+            if dE_max is not None:
+                self.set_step_in_calculators(step, first=True, last=True)
+            else:
+                self.set_step_in_calculators(step)
+
+            F = self.get_forces()
+            self.logger.info(f'Step: {step:{length_step}}. Energies = '
+                             f'{[np.round(en, 4) for en in self.get_energies()]}')
+
+            R = self.neb.get_residual()
+            if R <= fmax:
+                self.logger.info(f'Step: {step:{length_step}}. Optimization terminates successfully. Residual R = {R:.4f}')
+                return True
+            else:
+                self.logger.info(f'Step: {step:{length_step}}. Residual R = {R:.4f}')
+
+            X = self.neb.get_positions().reshape(-1)
+            X += alpha * F
+            self.update_positions(X)
+            self.dump_positions_vasp(prefix=f'Step:-{step}-2-')
+
+            if dE_max is not None:
+                energies = self.get_energies(first=True, last=True)
+                self.logger.debug(f'Energies raw: {energies}')
+                self.logger.info(f'Step: {step:{length_step}}. Energies = '
+                                 f'{[np.round(en, 4) for en in energies]}')
+                diff = np.abs(np.diff(energies))
+                self.logger.debug(f'diff: {diff}')
+                idxs = np.where(diff > dE_max)[0]
+                self.logger.debug(f'Idxs where diff > dE_max: {idxs}')
+
+                if len(idxs) > max_new_images_at_step:
+                    idxs = np.flip(np.argsort(diff))[:max_new_images_at_step]
+                    self.logger.debug(f'Images will be added for idxs {idxs} since more than {max_new_images_at_step} '
+                                      f'diffs were large than {dE_max} eV')
+
+                if (len(idxs) > 0) and (steps_after_insertion > min_steps_after_insertion):
+                    steps_after_insertion = -1
+                    for idx in reversed(idxs):
+                        self.logger.debug(f'Start working with idx: {idx}')
+                        length_prev = len(str(len(self.neb.images) - 1))
+                        length_new = len(str(len(self.neb.images)))
+                        self.logger.debug(f'{length_prev=} {length_new=}')
+                        tmp_images = [self.neb.images[idx].copy(),
+                                      self.neb.images[idx].copy(),
+                                      self.neb.images[idx + 1].copy()]
+                        tmp_neb = NEB(tmp_images)
+                        tmp_neb.interpolate()
+                        images_new = self.neb.images.copy()
+                        images_new.insert(idx + 1, tmp_neb.images[1])
+
+                        energies = self.get_energies(first=True, last=True)
+                        energies.insert(idx + 1, None)
+                        self.neb = NEB(images_new,
+                                       k=self.neb.k[0],
+                                       climb=self.neb.climb)
+                        self.dump_positions_vasp(prefix=f'Step: {step} 3 ')
+
+                        zfill_length = len(str(len(self.neb.images)))
+                        for k, image in enumerate(self.neb.images):
+                            self.logger.debug(f'Trying to attach the calc to {k} image '
+                                              f'with the length: {len(str(len(self.neb.images)))}')
+                            image.calc = construct_calc_fn(str(k).zfill(zfill_length))
+                            image.calc.E = energies[k]
+
+                        if length_prev != length_new:
+                            self.logger.debug('Trying to rename due to the change in length')
+                            for i in range(0, self.neb.nimages):
+                                shell(f'mv {str(i).zfill(length_prev)} {str(i).zfill(length_new)}')
+                                self.logger.debug(f'Trying to execute the following command: '
+                                                  f'mv {str(i).zfill(length_prev)} {str(i + 1).zfill(length_new)}')
+
+                        self.logger.debug('Trying to rename due to the insertion')
+                        for i in range(len(self.neb.images) - 2, idx, -1):
+                            self.logger.debug(f'{i=}')
+                            self.logger.debug(f'Trying to execute the following command: '
+                                              f'mv {str(i).zfill(length_new)} {str(i + 1).zfill(length_new)}')
+                            shell(f'mv {str(i).zfill(length_new)} {str(i + 1).zfill(length_new)}')
+
+                        self.logger.debug(f'Trying to create the new folder: {str(idx + 1).zfill(length_new)}')
+                        folder = Path(str(idx + 1).zfill(length_new))
+                        folder.mkdir()
+                        self.dump_positions_vasp(prefix=f'Step: {step} 4 ')
+
+                steps_after_insertion += 1
+                self.logger.debug(f'Step after insertion: {steps_after_insertion}')
+
+        self.logger.warning(f'convergence was not achieved after max iterations = {max_steps}, '
+                            f'residual R = {R:.4f} > {fmax}')
+        return False
+
+    def run_ode(self,
+                fmax: float = 0.1,
+                max_steps: int = 100,
+                C1: float = 1e-2,
+                C2: float = 2.0,
+                extrapolation_scheme: Literal[1, 2, 3] = 3,
+                h: float | None = None,
+                h_min: float = 1e-10,
+                R_max: float = 1e3,
+                rtol: float = 0.1):
+        """
+            fmax : float
+               convergence tolerance for residual force
+            max_steps : int
+                maximum number of steps
+            C1 : float
+               sufficient contraction parameter
+            C2 : float
+               residual growth control (Inf means there is no control)
+            extrapolation_scheme : int
+               extrapolation style (3 seems the most robust)
+            h : float
+               initial step size, if None an estimate is used based on ODE12
+            h_min : float
+               minimal allowed step size
+            R_max: float
+               terminate if residual exceeds this value
+            rtol : float
+               relative tolerance
+        """
+
+        if max_steps < 2:
+            raise ValueError('max_steps must be greater or equal than two')
+        length = len(str(max_steps))
+
+        self.set_step_in_calculators(0)
+
+        F = self.get_forces()
+        self.logger.info(f'Step: {0:{length}}. Energies = {[np.round(en, 4) for en in self.get_energies()]}')
+
+        R = self.neb.get_residual()  # pick the biggest force
+
+        if R >= R_max:
+            self.logger.info(f'Step: {0:{length}}. Residual {R:.4f} >= R_max {R_max}')
+            raise OptimizerConvergenceError(f'Step: 0. Residual {R:.4f} >= R_max {R_max}')
+        else:
+            self.logger.info(f'Step: {0:{length}}. Residual R = {R:.4f}')
+
+        if h is None:
+            h = 0.5 * rtol ** 0.5 / R  # Chose a step size based on that force
+            h = max(h, h_min)  # Make sure the step size is not too big
+        self.logger.info(f'Step: {0:{length}}. Step size h = {h}')
+
+        X = self.neb.get_positions().reshape(-1)
+
+        for step in range(1, max_steps):
+            X_new = X + h * F  # Pick a new position
+            self.update_positions(X_new)
+
+            self.set_step_in_calculators(step)
+            F_new = self.get_forces()  # Calculate the new forces at this position
+            self.logger.info(f'Step: {step:{length}}. Energies = {[np.round(en, 4) for en in self.get_energies()]}')
+
+            R_new = self.neb.get_residual()
+            self.logger.info(f'Step: {step:{length}}. At new coordinates R = {R:.4f} -> R_new = {R_new:.4f}')
+
+            e = 0.5 * h * (F_new - F)  # Estimate the area under the forces curve
+            err = np.linalg.norm(e, np.inf)  # Error estimate
+
+            # Accept step if residual decreases sufficiently and/or error acceptable
+            condition_1 = R_new <= R * (1 - C1 * h)
+            condition_2 = R_new <= R * C2
+            condition_3 = err <= rtol
+            accept = condition_1 or (condition_2 and condition_3)
+            self.logger.info(f'Step: {step:{length}}. {"R_new <= R * (1 - C1 * h)":26} \t is {condition_1}')
+            self.logger.info(f'Step: {step:{length}}. {"R_new <= R * C2":26} is {condition_2}')
+            self.logger.info(f'Step: {step:{length}}. {"err <= rtol":26} is {condition_3}')
+
+            # Pick an extrapolation scheme for the system & find new increment
+            y = F - F_new
+            if extrapolation_scheme == 1:  # F(xn + h Fp)
+                h_ls = h * (F @ y) / (y @ y)
+            elif extrapolation_scheme == 2:  # F(Xn + h Fp)
+                h_ls = h * (F @ F_new) / (F @ y + 1e-10)
+            elif extrapolation_scheme == 3:  # min | F(Xn + h Fp) |
+                h_ls = h * (F @ y) / (y @ y + 1e-10)
+            else:
+                raise ValueError(f'Invalid extrapolation_scheme: {extrapolation_scheme}. Must be 1, 2 or 3')
+
+            if np.isnan(h_ls) or h_ls < h_min:  # Rejects if increment is too small
+                h_ls = np.inf
+
+            h_err = h * 0.5 * np.sqrt(rtol / err)
+
+            if accept:
+                self.logger.info(f'Step: {step:{length}}. The displacement is accepted')
+
+                X = X_new
+                R = R_new
+                F = F_new
+
+                self.dump_trajectory()
+                self.dump_positions_vasp()
+
+                # We check the residuals again
+                if self.converged(fmax):
+                    self.logger.info(f"Step: {step:{length}}. Optimization terminates successfully")
+                    return True
+
+                if R > R_max:
+                    self.logger.info(f"Step: {step:{length}}. Optimization fails, R = {R:.4f} > R_max = {R_max}")
+                    return False
+
+                # Compute a new step size.
+                # Based on the extrapolation and some other heuristics
+                h = max(0.25 * h, min(4 * h, h_err, h_ls))  # Log steep-size analytic results
+                self.logger.info(f'Step: {step:{length}}. New step size h = {h}')
+
+            else:
+                self.logger.info(f'Step: {step:{length}}. The displacement is rejected')
+                h = max(0.1 * h, min(0.25 * h, h_err, h_ls))
+                self.logger.info(f'Step: {step:{length}}. New step size h = {h}')
+
+            if abs(h) < h_min:  # abort if step size is too small
+                self.logger.info(f'Step: {step:{length}}. Stop optimization since step size h = {h} < h_min = {h_min}')
+                return True
+
+        self.logger.warning(f'Step: {step:{length}}. Convergence was not achieved after max iterations = {max_steps}')
+        return True
+
+
+class NEB_JDFTx:
+    def __init__(self,
+                 path_jdftx_executable: str | Path,
+                 nimages: int = 5,
+                 input_filepath: str | Path = 'input.in',
+                 output_name: str = 'output.out',
+                 input_format: Literal['jdftx', 'vasp'] = 'jdftx',
+                 cNEB: bool = True,
+                 spring_constant: float = 5.0,
+                 interpolation_method: Literal['linear', 'idpp'] = 'idpp',
+                 restart: Literal[False, 'from_traj', 'from_vasp'] = False,
+                 dE_max: float = None):
+
+        if isinstance(path_jdftx_executable, str):
+            self.path_jdftx_executable = Path(path_jdftx_executable)
+        else:
+            self.path_jdftx_executable = path_jdftx_executable
+
+        if isinstance(input_filepath, str):
+            input_filepath = Path(input_filepath)
+        self.jdftx_input = Input.from_file(input_filepath)
+
+        self.nimages = nimages
+        self.path_rundir = Path.cwd()
+        self.output_name = output_name
+        self.input_format = input_format.lower()
+        self.cNEB = cNEB
+        self.restart = restart
+        self.spring_constant = spring_constant
+        self.interpolation_method = interpolation_method.lower()
+        self.dE_max = dE_max
+        self.optimizer = None
+        self.logger = logging.getLogger(self.__class__.__name__ + ':')
+
+    def prepare(self):
+        length = len(str(self.nimages + 1))
+
+        if self.restart is False:
+            if self.input_format == 'jdftx':
+                init_ionpos = Ionpos.from_file('init.ionpos')
+                init_lattice = Lattice.from_file('init.lattice')
+                final_ionpos = Ionpos.from_file('final.ionpos')
+                final_lattice = Lattice.from_file('final.lattice')
+                init_poscar = init_ionpos.convert('vasp', init_lattice)
+                init_poscar.to_file('init.vasp')
+                final_poscar = final_ionpos.convert('vasp', final_lattice)
+                final_poscar.to_file('final.vasp')
+
+            initial = read('init.vasp', format='vasp')
+            final = read('final.vasp', format='vasp')
+
+            images = [initial]
+            images += [initial.copy() for _ in range(self.nimages)]
+            images += [final]
+
+            neb = NEB(images,
+                      k=self.spring_constant,
+                      climb=self.cNEB)
+            neb.interpolate(method=self.interpolation_method)
+
+            for i, image in enumerate(images):
+                image.write(f'start_img_{str(i).zfill(length)}.vasp', format='vasp')
+
+        else:
+            images = []
+            if self.restart == 'from_traj':
+                trj = Trajectory('NEB_trajectory.traj')
+                n_iter = int(len(trj) / (self.nimages + 2))
+                for i in range(self.nimages + 2):
+                    trj[(n_iter - 1) * (self.nimages + 2) + i].write(f'start_img_{str(i).zfill(length)}.vasp',
+                                                                     format='vasp')
+                trj.close()
+
+            if self.restart == 'from_traj' or self.restart == 'from_vasp':
+                for i in range(self.nimages + 2):
+                    img = read(f'start_img_{str(i).zfill(length)}.vasp', format='vasp')
+                    images.append(img)
+
+            else:
+                raise ValueError(f'restart must be False or \'from_traj\', '
+                                 f'or \'from_vasp\' but you set {self.restart=}')
+
+            neb = NEB(images,
+                      k=self.spring_constant,
+                      climb=self.cNEB)
+
+        for i in range(self.nimages):
+            folder = Path(str(i+1).zfill(length))
+            folder.mkdir(exist_ok=True)
+        if self.dE_max is not None:
+            self.logger.debug(f'Trying to create the folder {str(0).zfill(length)}')
+            folder = Path(str(0).zfill(length))
+            folder.mkdir(exist_ok=True)
+            self.logger.debug(f'Trying to create the folder {str(self.nimages + 1).zfill(length)}')
+            folder = Path(str(self.nimages + 1).zfill(length))
+            folder.mkdir(exist_ok=True)
+
+        for i, image in enumerate(images[1:-1]):
+            image.calc = JDFTx(self.path_jdftx_executable,
+                               path_rundir=self.path_rundir / str(i+1).zfill(length),
+                               commands=self.jdftx_input.commands)
+
+        if self.dE_max is not None:
+            images[0].calc = JDFTx(self.path_jdftx_executable,
+                                   path_rundir=self.path_rundir / str(0).zfill(length),
+                                   commands=self.jdftx_input.commands)
+            images[-1].calc = JDFTx(self.path_jdftx_executable,
+                                    path_rundir=self.path_rundir / str(self.nimages + 1).zfill(length),
+                                    commands=self.jdftx_input.commands)
+
+        self.optimizer = NEBOptimizer(neb=neb,
+                                      trajectory_filepath='NEB_trajectory.traj')
+
+    def run(self,
+            fmax: float = 0.1,
+            method: Literal['ode', 'static'] = 'ode',
+            max_steps: int = 100,
+            **kwargs):
+
+        self.prepare()
+
+        if self.dE_max is not None:
+            def calc_fn(folder_name) -> JDFTx:
+                return JDFTx(self.path_jdftx_executable,
+                             path_rundir=self.path_rundir / folder_name,
+                             commands=self.jdftx_input.commands)
+        else:
+            calc_fn = None
+        if method == 'ode':
+            self.optimizer.run_ode(fmax, max_steps)
+        elif method == 'static':
+            self.optimizer.run_static(fmax, max_steps, dE_max=self.dE_max, construct_calc_fn=calc_fn)
+        else:
+            raise ValueError(f'Method must be ode or static but you set {method=}')
+
+
+class AutoNEB_JDFTx:
+    """
+    Class for running AutoNEB with JDFTx calculator
+
+    Parameters:
+
+    prefix: string or Path
+        path to folder with initial files. Basically could be os.getcwd()
+        In this folder required:
+            1) init.vasp file with initial configuration
+            2) final.vasp file with final configuration
+            3) in file with JDFTx calculation parameters
+    path_jdftx_executable: string or Path
+        path to jdftx executable
+    n_start: int
+        Starting number of images between starting and final for NEB
+    n_max: int
+        Maximum number of images, including starting and final
+    climb: boolean
+        Whether it is necessary to use cNEB or not
+    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 or object
+        Set optimizer for NEB: FIRE, BFGS or NEB
+    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.
+    interpolation_method: string
+        method for interpolation
+    smooth_curve: boolean
+    """
+
+    def __init__(self,
+                 prefix,
+                 path_jdftx_executable,
+                 n_start=3,
+                 n_simul=3,
+                 n_max=10,
+                 climb=True,
+                 fmax=0.05,
+                 maxsteps=100,
+                 k=0.1,
+                 restart=False,
+                 method='eb',
+                 optimizer='FIRE',
+                 space_energy_ratio=0.5,
+                 interpolation_method='idpp',
+                 smooth_curve=False):
+        self.restart = restart
+        self.path_jdftx_executable = path_jdftx_executable
+        self.prefix = Path(prefix)
+        self.n_start = n_start
+        self.n_max = n_max
+        self.commands = Input.from_file(Path(prefix) / 'in').commands
+        self.interpolation_method = interpolation_method
+        self.autoneb = AutoNEB(self.attach_calculators,
+                               prefix=prefix,
+                               n_simul=n_simul,
+                               n_max=n_max,
+                               climb=climb,
+                               fmax=fmax,
+                               maxsteps=maxsteps,
+                               k=k,
+                               method=method,
+                               space_energy_ratio=space_energy_ratio,
+                               world=None, parallel=False, smooth_curve=smooth_curve,
+                               interpolate_method=interpolation_method, optimizer=optimizer)
+
+    def prepare(self):
+        if not self.restart:
+            initial = read(self.prefix / 'init.vasp', format='vasp')
+            final = read(self.prefix / 'final.vasp', format='vasp')
+            images = [initial]
+            if self.n_start != 0:
+                images += [initial.copy() for _ in range(self.n_start)]
+            images += [final]
+            if self.n_start != 0:
+                neb = NEB(images)
+                neb.interpolate(method=self.interpolation_method)
+            for i, image in enumerate(images):
+                image.write(self.prefix / f'{i:03d}.traj', format='traj')
+                image.write(self.prefix / f'{i:03d}.vasp', format='vasp')
+        else:
+            index_exists = [i for i in range(self.n_max) if
+                            os.path.isfile(self.prefix / f'{i:03d}.traj')]
+            for i in index_exists:
+                image = Trajectory(self.prefix / f'{i:03d}.traj')
+                image[-1].write(self.prefix / f'{i:03d}.vasp', format='vasp')
+                img = read(self.prefix / f'{i:03d}.vasp', format='vasp')
+                img.write(self.prefix / f'{i:03d}.traj', format='traj')
+
+    def attach_calculators(self, images, indexes, iteration):
+        for image, index in zip(images, indexes):
+            path_rundir = self.autoneb.iter_folder / f'iter_{iteration}' / str(index)
+            path_rundir.mkdir(exist_ok=True)
+            image.calc = JDFTx(self.path_jdftx_executable,
+                               path_rundir=path_rundir,
+                               commands=self.commands)
+
+    def run(self):
+        self.prepare()
+        self.autoneb.run()