1"""
  2==============================================================================
  3Module Name: pyiron_NEB_task.py
  4==============================================================================
  5Description:
  6    This module provides a framework for setting up and executing NEB (Nudged
  7    Elastic Band) calculations using the pyiron framework with VASP. It includes
  8    classes and methods to manage initial and final states, reorder atoms,
  9    create NEB jobs, handle image interpolation, and collect and fit energy data.
 10    
 11    WARNING: It assumes remote and manual execution of VASP (on the cluster), it 
 12    only handles transfer of the input/output files, not the execution (as should 
 13    be ideally the case with pyiron).
 14
 15Dependencies:
 16    - pyiron_atomistics
 17    - pyiron_base
 18    - NumPy
 19    - SciPy
 20
 21Usage:
 22    from NEB import NEB
 23    # Create a NEB object, set initial and final states, and run the NEB calculation
 24    neb = NEB(pr, 'NEB_Calculation')
 25    neb.set_initial_state(initial_job)
 26    neb.set_final_state(final_job)
 27    neb_job = neb.create_neb_job(job_name='neb', nImages=5)
 28    neb_job.write_and_transfer_to_remote()
 29    # run the job manually on HPC
 30    neb_job.transfer_from_remote_and_collect()
 31    # get the barrier
 32    print(job['output/neb/barrier_forward'])
 33
 34==============================================================================
 35"""
 36
 37__author__ = "David Holec"
 38__version__ = "0.0.1"
 39__email__ = "david.holec@unileoben.ac.ay"
 40__status__ = "Prototype"  # Options: "Development", "Production", "Prototype"
 41__date__ = "2024-08-09"
 42__license__ = "BSD 3-Clause License"
 43
 44
 45from pyiron_atomistics.vasp.vasp import Vasp as Vasp_job
 46from pyiron_atomistics.vasp.structure import read_atoms
 47from pyiron_atomistics.vasp.parser.oszicar import Oszicar
 48from pyiron_atomistics.atomistics.job.atomistic import Trajectory
 49from pyiron_atomistics.atomistics.structure.atoms import Atoms
 50from pyiron_base import state
 51from pyiron_base.jobs.job.extension.jobstatus import job_status_successful_lst
 52from glob import glob
 53import os
 54import os.path
 55from shutil import rmtree
 56import numpy as np
 57from scipy.optimize import minimize
 58from numpy.polynomial import Polynomial
 59
 60def get_species_symbols_new(self):
 61    """
 62    MONKEY PATCH to get correct order of POTCARs.
 63    
 64    Returns:
 65        numpy.ndarray: List of the symbols of the species.
 66    """
 67    sp = [self.elements[0]]
 68    for el in self.elements[1:]:
 69        if el != sp[-1]:
 70            sp.append(el)
 71    return np.array([el.Abbreviation for el in sp])
 72
 73# Patch the `get_species_symbols` method of the `Atoms` class
 74Atoms.get_species_symbols = get_species_symbols_new
 75
 76
 77def _wrap_positions(struct, wrap):
 78    """
 79    Wrap atomic positions into the unit cell.
 80
 81    Parameters:
 82        struct (Atoms): Atomic structure.
 83        wrap (float or list): Wrap value(s) for each dimension.
 84
 85    Returns:
 86        Atoms: Structure with wrapped positions.
 87    """
 88    if isinstance(wrap, float):
 89        wrap = [wrap] * 3
 90
 91    new_pos = []
 92    for pos in struct.get_scaled_positions():
 93        for i in range(3):
 94            if pos[i] > wrap[i]:
 95                pos[i] -= 1.0
 96        new_pos.append(pos)
 97
 98    struct.set_scaled_positions(new_pos)
 99    return struct
100
101
102class NEB():
103    """
104    Class to handle the setup and execution of NEB (Nudged Elastic Band) calculations.
105    """
106    
107    def __init__(self, pr, NEB_calc_name):
108        """
109        Initialize the NEB class.
110
111        Parameters:
112            pr (Project): Pyiron project object.
113            NEB_calc_name (str): Name of the NEB calculation.
114        """
115        self.pr = pr.create_group(NEB_calc_name)
116        self.initial = None
117        
118    def add_initial_state(self, initial_job, delete_existing=False, wrap=0.98):
119        """
120        Add the initial state for the NEB calculation.
121
122        Parameters:
123            initial_job (Vasp_job): Initial state as a VASP job.
124            delete_existing (bool, optional): Whether to delete existing files.
125            wrap (float, optional): Wrap value for the atomic positions.
126        
127        Raises:
128            TypeError: If `initial_job` is not of type `Vasp_job`.
129        """
130        if not isinstance(initial_job, Vasp_job):
131            raise TypeError("Initial job must be a VASP job (pyiron_atomistics.vasp.vasp.Vasp)")
132        
133        self.initial_job = initial_job
134        self.initial = initial_job.get_structure()
135        if wrap is not False:
136            self.initial = _wrap_positions(initial_job.get_structure(), wrap)
137        
138        # Setup symbolic links and directory structure
139        rel_path_to_initial = os.path.relpath(self.initial_job.path, self.pr.path)
140        cwd = os.getcwd()
141        os.chdir(self.pr.path)
142        if delete_existing:
143            try:
144                os.remove('initial.h5')
145                rmtree('initial_hdf5')
146            except:
147                pass
148        os.symlink(rel_path_to_initial+'.h5', 'initial.h5')
149        os.mkdir('initial_hdf5')
150        os.chdir('initial_hdf5')
151        os.symlink(os.path.join('..', rel_path_to_initial+'_hdf5', initial_job.name), 'initial', target_is_directory=True)
152        os.chdir(cwd)
153        
154    def set_initial_state(self, initial_job, wrap=0.98):
155        """
156        Set the initial state from an existing VASP job.
157
158        Parameters:
159            initial_job (Vasp_job): Initial state as a VASP job.
160            wrap (float, optional): Wrap value for the atomic positions.
161        """
162        self.initial_job = initial_job
163        struct = initial_job.get_structure()
164        self.initial = struct
165        if wrap is not False:
166            self.initial = _wrap_positions(struct, wrap)
167    
168    def set_final_state(self, final_job, wrap=0.98):
169        """
170        Set the final state for the NEB calculation.
171
172        Parameters:
173            final_job (Vasp_job): Final state as a VASP job.
174            wrap (float, optional): Wrap value for the atomic positions.
175        """
176        self.final_job = final_job
177        struct = final_job.get_structure()
178        self.final = struct
179        if wrap is not False:
180            self.final = _wrap_positions(struct, wrap)
181        
182        
183    def reorder_final(self, moving_index, max_shift=0.05, verbose=True, wrap=0.98):
184        """
185        Reorder the final structure based on the initial structure.
186
187        Parameters:
188            moving_index (int): Index of the moving atom.
189            max_shift (float, optional): Maximum allowable shift for atom mapping.
190            verbose (bool, optional): If True, prints information about the atom positions.
191            wrap (float, optional): Wrap value for the atomic positions.
192        
193        Raises:
194            Exception: If there are issues with the atom mapping.
195        """
196        
197        initial = self.initial
198        final = self.final
199        if wrap is not False:
200            initial = _wrap_positions(initial, wrap)
201            final = _wrap_positions(final, wrap)
202            
203        mapped = [False for _ in range(len(final))]
204        mapping = [None for _ in range(len(initial))]
205        
206        for iat, at in enumerate(initial):
207            if iat == moving_index and verbose:
208                print('position in the initial structure:', at.position)
209            else:
210                nn = final.get_neighborhood(at.position, cutoff_radius=max_shift)
211                if len(nn.distances) != 1:
212                    raise Exception(f"Problem with mapping: atom {iat} in the initial structure was identified with {len(nn.distances)} images")
213                elif mapped[nn.indices[0]]:
214                    raise Exception(f"Problem with mapping: atom {iat} in the initial should be mapped to {nn.indices[0]} in the final structure but this atom has already been mapped.")
215                else:
216                    mapped[nn.indices[0]] = True
217                    mapping[iat] = nn.indices[0]
218        
219        unmapped = [i for i in range(len(final)) if not mapped[i]]
220        if len(unmapped) != 1:
221            raise Exception(f"Problem with mapping: only 1 atom must remain after mapping! Unmapped atoms: {unmapped}")
222        
223        mapping[moving_index] = unmapped[0]
224        if verbose:
225            print('position in the final structure:', self.final.positions[unmapped[0]])
226        
227        new_positions = [final.positions[mapping[i]] for i in range(len(initial))]
228        new_elements = [final.elements[mapping[i]] for i in range(len(initial))]
229        new_final = self.pr.create.structure.atoms(
230            cell=final.cell, 
231            positions=new_positions, 
232            elements=new_elements
233        )
234        self.final = new_final
235        
236    def create_neb_job(self, job_name='neb', nImages=3):
237        """
238        Create a NEB job.
239
240        Parameters:
241            job_name (str, optional): Name of the NEB job.
242            nImages (int, optional): Number of images for the NEB calculation.
243        
244        Returns:
245            NEBjob: Created NEB job object.
246        """
247        try:
248            job = self.pr[job_name]
249        except:
250            job = self.pr.create_job(job_type=self.pr.job_type.Vasp, job_name=job_name)
251            job.__class__ = NEBjob
252            job.structure = self.initial         
253            job.set_kpoints(mesh=[3, 3, 3], scheme='GC')
254            job.calc_minimize(
255                electronic_steps=60,
256                retain_charge_density=False,
257            )
258            job.input.incar['ISIF'] = 2
259            job.input.incar['ENCUT'] = 520  # from MP
260            job.input.incar['IMAGES'] = nImages - 2
261            job.input.incar['IBRION'] = 1
262            job.input.incar['NFREE'] = 2
263            job.save()
264            
265            job['input'].create_group('neb')            
266            job['input/neb'].put('nImages', nImages)
267            job['input/neb'].put('initial_image', self.initial)
268            job['input/neb'].put('final_image', self.final)
269            job['input/neb'].put('initial_energy', self.initial_job['output/generic/energy_pot'][-1])
270            job['input/neb'].put('final_energy', self.final_job['output/generic/energy_pot'][-1])
271            job._create_images()
272            job.structure = job['input/neb/initial_image'].to_object()
273            job.write_input()
274        else:
275            state.logger.warning(f'Job `{job_name}` exists. Loading from database instead of creating it!')
276        return job
277
278
279class NEBjob(Vasp_job):
280    """
281    Class for handling specific NEB (Nudged Elastic Band) job operations.
282    """
283    
284
285    
286    def _create_images(self, wrap=0.98):
287        """
288        Create interpolated images between the initial and final states.
289
290        Parameters:
291            wrap (float, optional): Wrap value for the atomic positions.
292        """
293        cells = []
294        positions = []
295        initial = _wrap_positions(self['input/neb/initial_image'].to_object(), wrap)
296        final = _wrap_positions(self['input/neb/final_image'].to_object(), wrap)
297        nImages = self['input/neb/nImages']
298        
299        for x in np.linspace(0, 1, nImages):
300            cells.append(initial.cell * (1 - x) + final.cell * x)
301            positions.append(initial.positions * (1 - x) + final.positions * x)
302        
303        self['input/neb'].put('initial_cells', cells)
304        self['input/neb'].put('initial_positions', positions)
305        
306    def get_initial_images(self):
307        """
308        Retrieve the initial interpolated images.
309
310        Returns:
311            Trajectory: Trajectory object containing the initial images.
312        """
313        initial = self['input/neb/initial_image'].to_object()
314        positions = self['input/neb/initial_positions']
315        cells = self['input/neb/initial_cells']
316        return Trajectory(positions, initial, cells=cells)
317    
318    def get_final_images(self):
319        """
320        Retrieve the final interpolated images after relaxation.
321
322        Returns:
323            Trajectory: Trajectory object containing the final images.
324        """
325        initial = self['input/neb/initial_image'].to_object()
326        positions = self['output/neb/final_positions']
327        cells = self['output/neb/final_cells']
328        return Trajectory(positions, initial, cells=cells)
329    
330    def _write_images(self):
331        """
332        Write the interpolated images to the appropriate directories.
333        """
334        path = self.path + '_hdf5/' + self.job_name + '/'
335        nImages = self['input/neb/nImages']
336        initial_structs = self.get_initial_images()
337        
338        for i in range(nImages):
339            os.chdir(path)
340            case = f'{i:02d}'
341            if not os.path.exists(case):
342                os.mkdir(case)
343            os.chdir(case)
344            struct = initial_structs[i]
345            struct.write('POSCAR', format='vasp')
346    
347    def write_and_transfer_to_remote(self):
348        """
349        Write the input files and transfer them to the remote cluster.
350
351        Warnings:
352            If the job is already marked as finished, a warning is issued and no action is taken.
353        """
354        if self.status not in job_status_successful_lst:
355            self._write_images()
356            self.write_input()
357            
358            filename = state.queue_adapter.convert_path_to_remote(
359                path=self.project_hdf5.file_name
360            )
361            working_directory = state.queue_adapter.convert_path_to_remote(
362                path=self.working_directory
363            )
364            for filename in glob(self.project_hdf5.path + '_hdf5/**', recursive=True):
365                if os.path.isfile(filename):                
366                    state.queue_adapter.transfer_file_to_remote(
367                        file=filename, transfer_back=False
368                    )
369        else:
370            state.logger.warning(f'Job `{self.job_name}` is marked as finished, uploading doesn\'t make sense')
371            
372    def collect_data(self, wrap=0.98):
373        """
374        Collect the data from the NEB calculation.
375
376        Parameters:
377            wrap (float or bool, optional): Wrap value for atomic positions. Set to `False` to disable wrapping.
378        """
379        path = self.path + '_hdf5/' + self.job_name + '/'
380        nImages = self['input/neb/nImages']
381        initial = self['input/neb/initial_image'].to_object()
382        final = self['input/neb/final_image'].to_object()
383        
384        if wrap is not False:
385            initial = _wrap_positions(initial, wrap)
386            final = _wrap_positions(final, wrap)
387        
388        cells = [initial.cell]
389        positions = [initial.positions]
390        energy_pot = []
391        
392        for i in range(1, nImages - 1):
393            case = f'{i:02d}'
394            out = Oszicar()
395            out.from_file(filename=path + case + '/OSZICAR')
396            energy_pot.append(out.parse_dict['energy_pot'])
397            struct = read_atoms(filename=path + case + '/CONTCAR')
398            if wrap is not False:
399                struct = _wrap_positions(struct, wrap)
400            cells.append(struct.cell)
401            positions.append(struct.positions)
402        
403        relaxation_steps = len(out.parse_dict['energy_pot'])
404        energy_pot = np.array(
405            [[self['input/neb/initial_energy']] * relaxation_steps] + 
406            energy_pot +
407            [[self['input/neb/final_energy']] * relaxation_steps]
408        ).T
409        
410        cells.append(final.cell)
411        positions.append(final.positions)
412        
413        self['output'].create_group('neb')
414        self['output/neb'].put('energy_pot', energy_pot)
415        self['output/neb'].put('final_cells', cells)
416        self['output/neb'].put('final_positions', positions)
417        ene = energy_pot[-1]
418        self['output/neb'].put('barrier_forward', max(ene) - ene[0])
419        self['output/neb'].put('barrier_backward', max(ene) - ene[-1])
420        self._polynomial_fit_with_derivative_constraints(n=nImages + 1)
421        
422    def transfer_from_remote_and_collect(self):
423        """
424        Transfer files from the remote cluster and collect data locally.
425        """
426        self.transfer_from_remote()
427        self.status.collect = True
428        self.collect_data()
429        self.status.finished = True
430    
431    def _polynomial_fit_with_derivative_constraints(self, n=4):
432        """
433        Fit a polynomial to the NEB energy profile with derivative constraints.
434
435        The polynomial is constrained such that its derivatives at x=0 and x=1 are zero.
436
437        Parameters:
438            n (int, optional): The order of the polynomial.
439
440        Returns:
441            Polynomial: The fitted polynomial.
442        """
443        y = self['output/neb/energy_pot'][-1]
444        x = np.linspace(0, 1, len(y))
445        
446        def poly_val_and_deriv(coeffs, x):
447            """
448            Evaluate the polynomial and its derivative.
449
450            Parameters:
451                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
452                x (float or numpy.ndarray): Point(s) at which to evaluate.
453
454            Returns:
455                tuple: Tuple containing the polynomial value and its derivative.
456            """
457            p = Polynomial(coeffs)
458            dp = p.deriv()
459            return p(x), dp(x)
460
461        def objective(coeffs, x, y):
462            """
463            Objective function to minimize.
464
465            Parameters:
466                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
467                x (float or numpy.ndarray): x-values of the data points.
468                y (float or numpy.ndarray): y-values of the data points.
469
470            Returns:
471                float: Sum of squared residuals.
472            """
473            p_vals, _ = poly_val_and_deriv(coeffs, x)
474            return np.sum((y - p_vals) ** 2)
475
476        def constraint_deriv_0(coeffs):
477            """
478            Constraint for the derivative at x=0 to be zero.
479
480            Parameters:
481                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
482
483            Returns:
484                float: Derivative at x=0.
485            """
486            _, dp0 = poly_val_and_deriv(coeffs, 0)
487            return dp0
488
489        def constraint_deriv_1(coeffs):
490            """
491            Constraint for the derivative at x=1 to be zero.
492
493            Parameters:
494                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
495
496            Returns:
497                float: Derivative at x=1.
498            """
499            _, dp1 = poly_val_and_deriv(coeffs, 1)
500            return dp1
501
502        def constraint_val_0(coeffs):
503            """
504            Constraint for the polynomial value at x=0 to match the first data point.
505
506            Parameters:
507                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
508
509            Returns:
510                float: Difference between polynomial value at x=0 and y[0].
511            """
512            p0, _ = poly_val_and_deriv(coeffs, 0)
513            return p0 - y[0]
514
515        def constraint_val_1(coeffs):
516            """
517            Constraint for the polynomial value at x=1 to match the last data point.
518
519            Parameters:
520                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
521
522            Returns:
523                float: Difference between polynomial value at x=1 and y[-1].
524            """
525            p1, _ = poly_val_and_deriv(coeffs, 1)
526            return p1 - y[-1]
527
528        initial_guess = np.ones(n + 1)
529
530        def constraint_max_val(coeffs):
531            """
532            Constraint for the polynomial maximum to match the maximum y value.
533
534            Parameters:
535                coeffs (list or numpy.ndarray): Coefficients of the polynomial.
536
537            Returns:
538                float: Difference between maximum polynomial value and maximum y value.
539            """
540            p = Polynomial(coeffs)
541            x_vals = np.linspace(0.2, 0.8, 10)
542            p_vals = p(x_vals)
543            return max(p_vals) - max(y)
544
545        constraints = [
546            {'type': 'eq', 'fun': constraint_deriv_0},
547            {'type': 'eq', 'fun': constraint_deriv_1},
548            {'type': 'eq', 'fun': constraint_val_0},
549            {'type': 'eq', 'fun': constraint_val_1},
550            {'type': 'eq', 'fun': constraint_max_val}
551        ]
552
553        result = minimize(objective, initial_guess, args=(x, y), constraints=constraints)
554        optimized_coeffs = result.x
555        fitted_polynomial = Polynomial(optimized_coeffs)
556        self['output/neb'].put('fit', fitted_polynomial)
557        
558    def get_energy_fit(self):
559        """
560        Get the fitted polynomial for the NEB energy profile.
561
562        Returns:
563            Polynomial: The fitted polynomial.
564        """
565        return self['output/neb/fit']