wmaee.codes.pyiron.pyiron_CHGNet_job

  1# Metadata
  2__author__ = "David Holec, Amin Sakic"
  3__maintainer__ = "David Holec"
  4__email__ = "david.holec@unileoben.ac.at"
  5__date__ = "2024-12-19"
  6__version__ = "0.11"
  7__license__ = "MIT"
  8
  9from pyiron_atomistics.atomistics.job.atomistic import AtomisticGenericJob
 10from pyiron_atomistics import pyiron_to_ase, ase_to_pyiron
 11from pyiron_base import GenericParameters
 12from ase.io import write, read
 13from os.path import join
 14import pandas as pd
 15from ase import Atoms
 16from ase.io.trajectory import Trajectory
 17import pickle
 18import numpy as np
 19from typing import Literal, Optional
 20
 21
 22class CHGNet(AtomisticGenericJob):
 23    """
 24    Custom job class for using CHGNet within the pyiron framework.
 25
 26    This class provides methods to perform static calculations, structure relaxation,
 27    and molecular dynamics simulations using CHGNet.
 28
 29    Attributes
 30    ----------
 31    structure : pyiron.atomistics.structure.atoms.Atoms
 32        Structure associated with the job.
 33    input : CalcParams
 34        Input parameters for the CHGNet calculations.
 35    """
 36
 37
 38    def __init__(self, project, job_name):
 39        """
 40        Initialize the CHGNet job.
 41
 42        Parameters
 43        ----------
 44        project : pyiron.Project
 45            The pyiron project in which the job resides.
 46        job_name : str
 47            Name of the job.
 48        """
 49        super(CHGNet, self).__init__(project, job_name)
 50        self.structure = None
 51        self.input = CalcParams()
 52        if self.project_hdf5.file_exists:
 53            self.input.from_hdf(self.project_hdf5, group_name='input')
 54        self._executable_activate(codename='chgnet')            
 55
 56    @property
 57    def fmax(self) -> float:
 58        """
 59        Get the maximum force tolerance for ionic relaxation.
 60
 61        Returns
 62        -------
 63        float
 64            Maximum force tolerance (eV/Å).
 65        """
 66        return self._generic_input['ionic_force_tolerance']
 67
 68    @fmax.setter
 69    def fmax(self, fmax: float):
 70        """
 71        Set the maximum force tolerance for ionic relaxation.
 72
 73        Parameters
 74        ----------
 75        fmax : float
 76            Maximum force tolerance (eV/Å).
 77        """
 78        self._generic_input['ionic_force_tolerance'] = fmax
 79
 80    @property
 81    def max_iter(self) -> int:
 82        """
 83        Get the maximum number of iterations for relaxation.
 84
 85        Returns
 86        -------
 87        int
 88            Maximum number of iterations.
 89        """
 90        return self._generic_input['n_print']
 91
 92    @max_iter.setter
 93    def max_iter(self, max_iter: int):
 94        """
 95        Set the maximum number of iterations for relaxation.
 96
 97        Parameters
 98        ----------
 99        max_iter : int
100            Maximum number of iterations.
101        """
102        self._generic_input['n_print'] = max_iter
103
104        
105    def calc_static(
106        self,
107        on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
108    ):        
109        """
110        Perform a static calculation.
111
112        Parameters
113        ----------
114        on_isolated_atoms : {'ignore', 'warn', 'error'}, optional
115            Behavior for isolated atoms (default is 'warn').
116        """
117        self.calc_minimize(max_iter=0, relax_cell=False, on_isolated_atoms=on_isolated_atoms)
118        self._generic_input["calc_mode"] = "static"   
119       
120    def calc_minimize(
121        self,
122        ionic_force_tolerance: float = 0.1,
123        max_iter: int = 500,
124        n_print: int = 1,
125        relax_cell: bool | None = True,
126        save_path: str | None = 'relax.traj',
127        on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
128    ):
129        """
130        Perform a structure relaxation (geometry optimization). More CHGNet parameters can be specified by 
131        explicitly adding them to `job.input` dictionary, see
132        [https://chgnet.lbl.gov/api#class-structoptimizer](https://chgnet.lbl.gov/api#class-structoptimizer).
133
134        Parameters
135        ----------
136        ionic_force_tolerance : float, optional
137            Convergence criterion for forces (default is 0.1 eV/Å).
138        max_iter : int, optional
139            Maximum number of iterations (default is 500).
140        n_print : int, optional
141            Frequency of logging output (default is 1).
142        relax_cell : bool or None, optional
143            Whether to allow relaxation of the simulation cell (default is True).
144        save_path : str or None, optional
145            Path to save the trajectory file (default is 'relax.traj').
146        on_isolated_atoms : {'ignore', 'warn', 'error'}, optional
147            Behavior for isolated atoms (default is 'warn').
148        """
149        super().calc_minimize(
150            ionic_energy_tolerance=0,
151            ionic_force_tolerance=ionic_force_tolerance,
152            max_iter=max_iter,            
153        )       
154        self.input['relax_cell'] = relax_cell
155        self.input['save_path'] = save_path
156        self.input['verbose'] = True
157        self.input['assign_magmoms'] = True                          
158        self.input['on_isolated_atoms'] = on_isolated_atoms 
159        self._generic_input['n_print'] = n_print
160        self._generic_input['ionic_force_tolerance'] = ionic_force_tolerance        
161    
162    def _write_calc_minimize(self):
163        """
164        Write the calculation script for performing structure relaxation.
165
166        This method generates a Python script that uses CHGNet for structure
167        relaxation and writes it to a file in the working directory.
168        """
169        params = ['model=chgnet']
170        allowed_params = ['optimizer_class', 'use_device', 'on_isolated_atoms']
171        for p in self.input.keys():
172            v = self.input.get(p)
173            if p in allowed_params and not v == None:
174                if type(v) == str:
175                    params.append(p + "=\"" + v + "\"")
176                else:
177                    params.append(f'{p}={self.input.get(p)}')
178        model_params = ','.join(params)
179
180        params = []
181        params += [f"fmax={self._generic_input['ionic_force_tolerance']},steps={self._generic_input['max_iter']}"]
182        params += [f"loginterval={self._generic_input['n_print']}"]
183        allowed_params = ['relax_cell', 'ase_filter', 'save_path', 'crystal_feas_save_path']        
184        for p in self.input.keys():
185            v = self.input.get(p)
186            if p in allowed_params and not v == None:
187                if type(v) == str:
188                    params.append(p + "=\"" + v + "\"")
189                else:
190                    params.append(f'{p}={self.input.get(p)}')
191        params = ','.join(params)
192        fix_imports = ''
193        fixes = []
194        if 'fix_imports' in self.input.keys():
195            fix_imports = self.input.get('fix_imports')
196        if 'fixes' in self.input.keys():
197            f = self.input.get('fixes')
198            if isinstance(f, str):                
199                fixes = [self.input.get('fixes')]
200            elif isinstance(f, list):
201                fixes = f
202            else:
203                self.logger.warning('Fixes are neither string or array of strings. Ignoring')         
204        script = [
205            'from chgnet.model import CHGNet, StructOptimizer',
206            'from ase.io import read',
207            'from ase.io.cif import write_cif',
208            'from ase.io.trajectory import Trajectory',
209            'from pymatgen.io.ase import AseAtomsAdaptor',
210            'chgnet = CHGNet.load()',
211            fix_imports,
212            'struct = read("structure.cif")']
213        for f in fixes:
214            script.append(f'struct.set_constraint({f})')
215        script += [            
216            'relaxer = StructOptimizer(' + model_params + ')',
217            'result = relaxer.relax(struct,' + params + ')',
218            'final_structure = result["final_structure"]',
219            'ase_structure = AseAtomsAdaptor.get_atoms(final_structure)',
220            'write_cif("final_structure.cif", ase_structure)'
221        ]
222        with open(join(self.working_directory, 'calc_script.py'), 'w') as f:
223            f.writelines("\n".join(script))
224            
225    def calc_md(
226        self,
227        temperature=300,
228        pressure=None,
229        n_ionic_steps=100,
230        time_step=2.0,  # in femto-seconds
231        n_print=10,
232        temperature_damping_timescale=100.0, # taut: float | None = None,        
233        pressure_damping_timescale=None, # taup: float | None = None,        
234        trajectory="md_run.traj",
235        logfile="md_run.log",        
236        on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
237        starting_temperature: int | None = None,
238        ensemble: str = "nvt",
239    ):        
240        """
241        Set up and perform a molecular dynamics (MD) simulation. More CHGNet parameters can be specified by 
242        explicitly adding them to `job.input` dictionary, see
243        [https://chgnet.lbl.gov/api#class-moleculardynamics](https://chgnet.lbl.gov/api#class-moleculardynamics).
244
245        Parameters
246        ----------
247        temperature : float, optional
248            Target temperature of the simulation in Kelvin (default is 300 K).
249        pressure : float or None, optional
250            Target pressure for the simulation in bar (default is None).
251        n_ionic_steps : int, optional
252            Number of ionic steps to simulate (default is 100).
253        time_step : float, optional
254            Time step for the simulation in femtoseconds (default is 2.0 fs).
255        n_print : int, optional
256            Frequency of printing simulation logs (default is 10).
257        temperature_damping_timescale : float, optional
258            Timescale for temperature damping in femtoseconds (default is 100.0 fs).
259        pressure_damping_timescale : float or None, optional
260            Timescale for pressure damping in femtoseconds (default is None).
261        trajectory : str, optional
262            Name of the output trajectory file (default is 'md_run.traj').
263        logfile : str, optional
264            Name of the output log file (default is 'md_run.log').
265        on_isolated_atoms : {'ignore', 'warn', 'error'}, optional
266            Behavior for isolated atoms (default is 'warn').
267        starting_temperature : int or None, optional
268            Initial temperature of the simulation (default is the target temperature).
269        ensemble : str, optional
270            Ensemble type, e.g., 'nvt' or 'npt' (default is 'nvt').
271        """
272        super().calc_md(
273            temperature=temperature,
274            pressure=pressure,
275            n_ionic_steps=n_ionic_steps,
276            time_step=time_step,
277            n_print=n_print,
278            temperature_damping_timescale=temperature_damping_timescale,
279            pressure_damping_timescale=pressure_damping_timescale,            
280        )
281        self.input['ensemble'] = ensemble
282        self.input['trajectory'] = trajectory
283        self.input['logfile'] = logfile
284        self.input['on_isolated_atoms'] = on_isolated_atoms
285        if starting_temperature==None:
286            starting_temperature=temperature
287        self.input['starting_temperature'] = starting_temperature
288
289    def _write_calc_md(self):        
290        """
291        Write the calculation script for performing molecular dynamics (MD).
292
293        This method generates a Python script that uses CHGNet for running MD
294        simulations and writes it to a file in the working directory.
295        """
296        params = []
297        params += [f"temperature={self._generic_input['temperature']},pressure={self._generic_input['pressure']},timestep={self._generic_input['time_step']}"]
298        params += [f"loginterval={self._generic_input['n_print']},taut={self._generic_input['temperature_damping_timescale']},taup={self._generic_input['pressure_damping_timescale']}"]
299        allowed_params = ['ensemble', 'thermostat', 'starting_temperature', 'bulk_modulus', 'crystal_feas_logfile', 'append_trajectory', 'use_device', 'trajectory', 'logfile', 'on_isolated_atoms']                   
300        for p in self.input.keys():
301            v = self.input.get(p)
302            if p in allowed_params and not v==None:
303                if type(v)==str:
304                    params.append(p+"=\""+v+"\"")
305                else:
306                    params.append(f'{p}={self.input.get(p)}')
307        params = ','.join(params)
308        fix_imports = ''
309        fixes = []
310        if 'fix_imports' in self.input.keys():
311            fix_imports = self.input.get('fix_imports')
312        if 'fixes' in self.input.keys():
313            f = self.input.get('fixes')
314            if isinstance(f, str):                
315                fixes = [self.input.get('fixes')]
316            elif isinstance(f, list):
317                fixes = f
318            else:
319                self.logger.warning('Fixes are neither string or array of strings. Ignoring') 
320                    
321        script = [
322            'from chgnet.model import CHGNet, MolecularDynamics',
323            'chgnet = CHGNet.load()',
324            'from ase.io import read',
325            'from ase.io.cif import write_cif',
326            fix_imports,
327            'struct = read("structure.cif")']
328        for f in fixes:
329            script.append(f'struct.set_constraint({f})')
330        script += [
331            'md = MolecularDynamics(atoms=struct,model=chgnet,'+params+')',
332            f'md.run({self._generic_input["n_ionic_steps"]})',            
333            'traj = read(f"md_run.traj", index=":")',
334            'write_cif("final_structure.cif", traj[-1])'
335        ]
336        with open(join(self.working_directory, 'calc_script.py'), 'w') as f:
337            f.writelines("\n".join(script))
338
339    def write_input(self) -> None:
340        """
341        Write input files required for the CHGNet calculation.
342
343        This function writes the structure file (`structure.cif`), stores the input 
344        dictionary to the HDF5 format, and delegates further writing depending on the 
345        calculation mode (MD or static/minimize).
346
347        Returns
348        -------
349        None
350        """
351        write(join(self.working_directory, 'structure.cif'), pyiron_to_ase(self.structure))
352        self.input.to_hdf(self.project_hdf5, group_name='input')
353        if self._generic_input['calc_mode'] == 'md':
354            self._write_calc_md()
355        else:
356            # calc_mode = static is treated as minimize with only a single ioninc step
357            self._write_calc_minimize()
358
359
360    def _parse_calc_static(self, output: str = 'error.out', skip: int = 3) -> Optional[pd.DataFrame]:
361        """
362        Parse the output of a static calculation.
363
364        This function reads the static calculation output file and returns the 
365        corresponding DataFrame containing the optimization step and energy information.
366
367        Parameters
368        ----------
369        output : str, optional
370            The name of the output file to parse, by default 'error.out'.
371        skip : int, optional
372            The number of rows to skip at the start of the file, by default 3.
373
374        Returns
375        -------
376        pd.DataFrame or None
377            A DataFrame containing the optimization results, or None if parsing fails.
378        """
379        try:
380            header = pd.read_csv(
381                join(self.working_directory, output), 
382                skiprows=skip, 
383                nrows=0, 
384                sep=r'\s+'
385            )
386            df = pd.read_csv(
387                join(self.working_directory, output), 
388                skiprows=skip+1, 
389                sep=r'\s+',
390                header=None
391            )
392            df.columns = ['optimizer_class'] + list(header.columns[:])
393            return df
394        except:
395            self.logger.warning(f'Cannot parse output from {output} file.')
396            return None
397
398
399    def _parse_calc_md(self, output: str = 'md_run.log') -> Optional[pd.DataFrame]:
400        """
401        Parse the output of a molecular dynamics (MD) simulation.
402
403        This function reads the MD log file and returns the corresponding DataFrame 
404        containing time, total energy, potential energy, kinetic energy, and temperature.
405
406        Parameters
407        ----------
408        output : str, optional
409            The name of the MD log file to parse, by default 'md_run.log'.
410
411        Returns
412        -------
413        pd.DataFrame or None
414            A DataFrame containing the MD simulation results, or None if parsing fails.
415        """
416        try:            
417            df = pd.read_csv(
418                join(self.working_directory, output), 
419                skiprows=1, 
420                sep=r'\s+'
421            )
422            df.columns = ['Time', 'Etot', 'Epot', 'Ekin', 'T']
423            return df
424        except:
425            self.logger.warning(f'Cannot parse output from {output} file.')
426            return None
427
428
429    def collect_output(self) -> None:
430        """
431        Collect and store the results of the calculation.
432
433        This function parses the final structure, trajectory, and other output properties,
434        and stores them in the HDF5 format. The structure is updated, and various properties 
435        such as energies, forces, and magnetization are saved to the output group.
436
437        Returns
438        -------
439        None
440        """
441                    
442        relaxed_ase_structure = read(join(self.working_directory, 'final_structure.cif'), format='cif')
443        relaxed_pyiron_structure = ase_to_pyiron(relaxed_ase_structure)
444
445        # Store the relaxed structure in the HDF5 group
446        with self.project_hdf5.open('output') as h5out:
447            h5out['structure'] = relaxed_pyiron_structure
448        self.structure = relaxed_pyiron_structure  # Update the job's structure
449
450        # Parse trajectory        
451        if self._generic_input['calc_mode'] in ['static', 'minimize']:
452            traj_file = self.input['save_path']
453            try:
454                with open(join(self.working_directory, traj_file), 'rb') as file:
455                    traj = pickle.load(file)                   
456                with self.project_hdf5.open('output/generic') as h5out:
457                    h5out['cells'] = np.array(traj['cell'][:-1])
458                    h5out['positions'] = np.array(traj['atom_positions'][:-1])
459                    h5out['stresses'] = np.array([[[s[0],s[5],s[4]],[s[5],s[1],s[3]],[s[4],s[3],s[2]]] for s in traj['stresses'][:-1]])
460                    h5out['forces'] = np.array([f for f in traj['forces'][:-1]])
461                with self.project_hdf5.open('output/dft') as h5out:
462                    h5out['magnetization'] = np.array([f for f in traj['magmoms'][:-1]])
463            except:
464                self.logger.warning(f'Cannot parse output from {traj_file} file.')
465        elif self._generic_input['calc_mode'] == 'md':            
466            traj_file = self.input['trajectory']
467            try:
468                traj = read(join(self.working_directory, traj_file), index=":")
469                with self.project_hdf5.open('output/generic') as h5out:
470                    h5out['cells'] = np.array([s.get_cell() for s in traj])
471                    h5out['positions'] = np.array([s.positions for s in traj])
472                    h5out['stresses'] = np.array([s.get_stress() for s in traj])
473                    h5out['forces'] = np.array([s.get_forces() for s in traj])
474            except:
475                self.logger.warning(f'Cannot parse output from {traj_file} file.')
476
477        # Parse other properties
478        if self._generic_input['calc_mode'] in ['static', 'minimize']:            
479            df = self._parse_calc_static()
480            with self.project_hdf5.open('output/generic') as h5out:
481                h5out['energy_pot'] = df['Energy'].values
482                h5out['energy_tot'] = df['Energy'].values
483                h5out['max_force'] = df['fmax'].values
484                h5out['steps'] = df['Step'].values
485        if self._generic_input['calc_mode'] == 'md':
486            df = self._parse_calc_md(self.input['logfile'])
487            with self.project_hdf5.open('output/generic') as h5out:
488                h5out['energy_pot'] = df['Epot'].values
489                h5out['energy_tot'] = df['Etot'].values                
490                h5out['energy_kin'] = df['Ekin'].values
491                h5out['temperature'] = df['T'].values
492                h5out['time'] = df['Time'].values
493                h5out['steps'] = df.index.values
494
495        self.status.finished = True
496
497
498    def get_structure(self, frame: int = -1) -> Atoms:
499        """
500        Retrieve the structure at a given frame from the calculation output.
501
502        This function returns the atomic structure (positions, cell, species) at the 
503        specified frame index. The frame index is adjusted to ensure it is valid.
504
505        Parameters
506        ----------
507        frame : int, optional
508            The frame index to retrieve. If negative, it is counted from the last frame, 
509            by default -1.
510
511        Returns
512        -------
513        Atoms
514            The ASE Atoms object representing the structure at the given frame.
515
516        Raises
517        ------
518        ValueError
519            If no structures are found in the job output.
520        IndexError
521            If the frame index is out of range.
522        """
523        num_structures = self.number_of_structures
524        if num_structures == 0:
525            raise ValueError("No structures found in the job output.")
526
527        # Ensure frame is valid
528        if frame < 0:
529            frame += num_structures
530        if frame >= num_structures:
531            raise IndexError(
532                f"Frame {frame} is out of range for {num_structures} structures.")
533
534        with self.project_hdf5.open('output/generic') as h5out:
535            positions = h5out["positions"][frame]
536            cell = h5out["cells"][frame]            
537        # Read positions and cell
538        with self.project_hdf5.open('output/structure') as h5out:
539            species = h5out["species"]
540            indices = h5out["indices"]
541
542        # Generate symbols list
543        symbols = [species[i] for i in indices]
544        return ase_to_pyiron(Atoms(symbols=symbols, positions=positions, cell=cell, pbc=True))
545
546
547
548class CalcParams(GenericParameters):
549    """
550    Class to manage the calculation parameters for CHGNet calculations.
551
552    This class extends from `GenericParameters` and is used to define and manage 
553    parameters specific to CHGNet calculations. It allows for initialization of 
554    parameters and stores them in a table, with a default table name of "chgnet". 
555    The class can be used to load and store various CHGNet calculation settings and 
556    other related configurations.
557
558    Parameters
559    ----------
560    table_name : str, optional
561        The name of the table where the parameters are stored. By default, it is set 
562        to "chgnet", but it can be customized if necessary.
563    
564    Methods
565    -------
566    __init__(self, table_name="chgnet")
567        Initializes the parameters for the calculation.
568    """
569    
570    def __init__(self, table_name: str = "chgnet") -> None:
571        """
572        Initialize the CalcParams object with the given table name.
573
574        Parameters
575        ----------
576        table_name : str, optional
577            The name of the table where the parameters are stored. By default, 
578            it is set to "chgnet".
579        """
580        super(CalcParams, self).__init__(
581            table_name=table_name,
582        )