wmaee.codes.pyiron.pyiron_GRACE_job

  1# Metadata
  2__author__ = "David Holec, Amin Sakic"
  3__maintainer__ = "David Holec"
  4__email__ = "david.holec@unileoben.ac.at"
  5__date__ = "2025-08-03"
  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
 12
 13# from tensorpotential.calculator import grace_fm
 14# from tensorpotential.calculator import TPCalculator
 15
 16
 17import io
 18import ast
 19import os
 20from ase import Atoms
 21from ase.io import write, read
 22from ase.io.trajectory import Trajectory
 23from ase.constraints import ExpCellFilter
 24from typing import Tuple
 25
 26import pandas as pd
 27import numpy as np
 28from typing import Literal, Optional
 29
 30from wmaee.units import EV_PER_ANG3_TO_GPA
 31
 32_grace_models = [
 33    "GRACE-1L-MP-r6", "GRACE-2L-MP-r5", "GRACE-2L-MP-r6", "GRACE-FS-OAM",
 34    "GRACE-FS-OAM", "GRACE-1L-OAM", "GRACE-2L-OAM",
 35    "GRACE-FS-OMAT", "GRACE-1L-OMAT", "GRACE-2L-OMAT",
 36    "GRACE-1L-OMAT-medium-base", "GRACE-1L-OMAT-medium-ft-E", "GRACE-1L-OMAT-large-base", "GRACE-1L-OMAT-large-ft-E",
 37    "GRACE-2L-OMAT-medium-base", "GRACE-2L-OMAT-medium-ft-E", "GRACE-2L-OMAT-large-base", "GRACE-2L-OMAT-large-ft-E",
 38]
 39
 40class Grace(AtomisticGenericJob):
 41    """
 42    Custom job class for using Grace within the pyiron framework.
 43
 44    This class provides methods to perform static calculations, structure relaxation,
 45    and molecular dynamics simulations using Grace UMLIPs.
 46
 47    Attributes
 48    ----------
 49    structure : pyiron.atomistics.structure.atoms.Atoms
 50        Structure associated with the job.
 51    input : CalcParams
 52        Input parameters for the Grace calculations.
 53    """
 54
 55    def __init__(self, project, job_name):
 56        """
 57        Initialize the Grace job.
 58
 59        Parameters
 60        ----------
 61        project : pyiron.Project
 62            The pyiron project in which the job resides.
 63        job_name : str
 64            Name of the job.
 65        """
 66        super(Grace, self).__init__(project, job_name)
 67        self.structure = None
 68        self.input = CalcParams()
 69        if self.project_hdf5.file_exists:
 70            self.input.from_hdf(self.project_hdf5, group_name='input')
 71        # executable is in pyiron_resources/grace/bin
 72        self._executable_activate(codename='grace')
 73
 74    @property
 75    def fmax(self) -> float:
 76        """
 77        Maximum force tolerance for ionic relaxation.
 78
 79        This property controls the maximum force tolerance (in eV/Å) used during ionic relaxation.
 80
 81        Returns
 82        -------
 83        float
 84            Maximum force tolerance (eV/Å).
 85        """
 86        return self._generic_input['ionic_force_tolerance']
 87
 88    @fmax.setter
 89    def fmax(self, fmax: float):
 90        """
 91        Set the maximum force tolerance for ionic relaxation.
 92
 93        Parameters
 94        ----------
 95        fmax : float
 96            Maximum force tolerance (eV/Å).
 97        """
 98        self._generic_input['ionic_force_tolerance'] = fmax
 99
100    @property
101    def max_iter(self) -> int:
102        """
103        Get the maximum number of iterations for relaxation or MD run.
104
105        Returns
106        -------
107        int
108            Maximum number of iterations.
109        """
110        return self._generic_input['max_iter']
111
112    @max_iter.setter
113    def max_iter(self, max_iter: int):
114        """
115        Set the maximum number of iterations for relaxation or MD run.
116
117        Parameters
118        ----------
119        max_iter : int
120            Maximum number of iterations.
121        """
122        self._generic_input['max_iter'] = max_iter
123
124    @property
125    def model(self) -> str:
126        """
127        Return the current GRACE model.
128    
129        Returns
130        -------
131        str
132            The selected GRACE model name.
133        """
134        return self.input['model']
135
136
137    @model.setter
138    def model(self, model: str):
139        """
140        Set the GRACE model name.
141    
142        Parameters
143        ----------
144        model : str
145            Name of the GRACE model to use. Must be one of the supported
146            GRACE model identifiers:
147    
148            - "GRACE-1L-MP-r6"
149            - "GRACE-2L-MP-r5"
150            - "GRACE-2L-MP-r6"
151            - "GRACE-FS-OAM"
152            - "GRACE-1L-OAM"
153            - "GRACE-2L-OAM"
154            - "GRACE-FS-OMAT"
155            - "GRACE-1L-OMAT"
156            - "GRACE-2L-OMAT"
157            - "GRACE-1L-OMAT-medium-base"
158            - "GRACE-1L-OMAT-medium-ft-E"
159            - "GRACE-1L-OMAT-large-base"
160            - "GRACE-1L-OMAT-large-ft-E"
161            - "GRACE-2L-OMAT-medium-base"
162            - "GRACE-2L-OMAT-medium-ft-E"
163            - "GRACE-2L-OMAT-large-base"
164            - "GRACE-2L-OMAT-large-ft-E"
165    
166        Raises
167        ------
168        ValueError
169            If `model` is not a valid GRACE model identifier.
170        """
171        if model not in _grace_models:
172            raise ValueError(f"Unknown grace model: {model}")
173        self.input["model"] = model
174
175
176    def available_models(self):
177        """
178        Return the list of supported GRACE models.
179    
180        Returns
181        -------
182        list of str
183            All valid GRACE model identifiers.
184        """
185        return _grace_models
186
187    
188    ####################################################
189    ### --- IMPLEMENTATION OF STATIC CALCULATION --- ###
190    ####################################################
191    def calc_static(
192        self,
193        grace_model: Literal[*_grace_models] = "GRACE-2L-OMAT",
194    ):
195        """
196        Perform a static calculation.
197    
198        Parameters
199        ----------
200        grace_model : Literal[*_grace_models], optional
201            GRACE model to use for the relaxation.  
202            Must be one of the supported models, which can be obtained via 
203            :meth:`available_models`. Default is ``"GRACE-2L-OMAT"``.
204    
205        Raises
206        ------
207        ValueError
208            If the structure has not been set before calling this method.
209    
210        Notes
211        -----
212        This method configures the GRACE calculator for a static calculation
213        using the provided GRACE model.  
214        A structure must be assigned to the instance prior to calling
215        ``calc_static``.
216    
217        Returns
218        -------
219        None
220        """
221        if self.structure is None:
222            raise ValueError(
223                'Structure must be set before calling calc_static().')
224        self._generic_input['calc_mode'] = 'static'
225        self.model = grace_model
226
227        
228    def _write_calc_static(self):
229        """
230        Write the calculation script for performing static calculations.
231
232        This method generates a Python script that uses Grace for static calculations
233        and writes it to a file in the working directory.
234        """
235        model = self.input["model"]
236
237        script = [
238            'from tensorpotential.calculator import grace_fm',
239            'from ase.io import read',
240            'from ase.io.cif import write_cif',
241            '',
242            'initial_structure = read("structure.cif")',
243            f'initial_structure.calc = grace_fm("{model}")',
244            '',
245            '# perform the static calculation',
246            'energy_pot = initial_structure.get_total_energy()',
247            'forces = initial_structure.get_forces()',
248            'stresses = initial_structure.get_stress()',
249            '',
250            '# write the results to an output file',
251            'with open("log.out", "w") as f:',
252            '    f.write("step\tenergy_pot\tforces\tstresses\\n")',
253            '    f.write(f"1\\t{energy_pot}\\t{forces.tolist()}\\t{stresses.tolist()}\\n")',
254            '',
255            '# write the final structure to a CIF file',
256            'write_cif("final_structure.cif", initial_structure)',
257        ]
258        with open(os.path.join(self.working_directory, 'calc_script.py'), 'w') as f:
259            f.writelines("\n".join(script))
260
261    def _parse_calc_static(
262        self, 
263        output: str = 'log.out', 
264        skip: int = 0
265    ) -> Optional[pd.DataFrame]:
266        """
267        Parse the output of a static calculation and return optimization results.
268    
269        This method reads a static calculation output file and extracts the 
270        optimization steps and energy information, returning them as a DataFrame.
271    
272        Parameters
273        ----------
274        output : str, optional
275            Name of the output file to parse. Default is 'log.out'.
276        skip : int, optional
277            Number of rows to skip at the start of the file before parsing. Default is 0.
278    
279        Returns
280        -------
281        pd.DataFrame or None
282            A DataFrame containing the optimization results if parsing succeeds,
283            otherwise None.
284        """
285
286        try:
287            df = pd.read_csv(
288                os.path.join(self.working_directory, output),
289                sep=r"\t+", 
290                engine='python'
291            )
292            # convert strings to np.array
293            df['forces'] = df['forces'].apply(lambda x: np.array(ast.literal_eval(x)))
294            df['stresses'] = df['stresses'].apply(lambda x: np.array(ast.literal_eval(x)))
295            return df
296        except:
297            self.logger.warning(f'Cannot parse output from {output} file.')
298            return None        
299
300    
301    ##########################################################
302    ### --- IMPLEMENTATION OF MINIMIZATION CALCULATION --- ###
303    ##########################################################
304    
305    def calc_minimize(
306        self,
307        grace_model: Literal[*_grace_models] = "GRACE-2L-OMAT",
308        algorithm: Literal["FIRE", "BFGS", "LBFGS"] = "FIRE",
309        algorithm_kwargs: dict | None = None,
310        relax_cell: bool = False,
311        relax_cell_kwargs: dict | None = None,
312        ionic_force_tolerance: float = 0.001,
313        max_iter: int = 500,
314        n_print: int = 1,
315        save_path: str | None = "relax.traj",
316    ) -> None:
317        """
318        Perform a structure relaxation (geometry optimization).
319    
320        Parameters
321        ----------
322        grace_model : Literal[*_grace_models], optional
323            GRACE model to use for the relaxation.  
324            Must be one of the supported models, which can be obtained via 
325            :meth:`available_models`. Default is ``"GRACE-2L-OMAT"``.
326    
327        algorithm : {"FIRE", "BFGS", "LBFGS"}, optional
328            Ionic relaxation algorithm from ASE.  
329            Default is ``"FIRE"``.
330    
331        algorithm_kwargs : dict, optional
332            Additional keyword arguments passed to the relaxation algorithm.  
333            Default is ``None``.
334    
335        relax_cell : bool, optional
336            If ``True``, optimize the simulation cell in addition to ionic positions.  
337            Default is ``False``.
338    
339        relax_cell_kwargs : dict, optional
340            Additional keyword arguments for cell relaxation.  
341            Default is ``None``.
342    
343        ionic_force_tolerance : float, optional
344            Force convergence criterion in eV/Å.  
345            Default is ``0.001``.
346    
347        max_iter : int, optional
348            Maximum number of optimization steps.  
349            Default is ``500``.
350    
351        n_print : int, optional
352            Logging frequency for structures and stresses.  
353            Default is ``1``.
354    
355        save_path : str or None, optional
356            Path to save the relaxation trajectory.  
357            If ``None``, no trajectory is written.  
358            Default is ``"relax.traj"``.
359    
360        Returns
361        -------
362        None
363        """
364        
365        super().calc_minimize(
366            ionic_energy_tolerance=0,
367            ionic_force_tolerance=ionic_force_tolerance,
368            max_iter=max_iter,
369        )
370
371        # Saving the parameters in dictionary
372        self._generic_input['calc_mode'] = 'minimize'
373        self.input['model'] = grace_model
374        self.input['algo'] = algorithm
375        self.input['algo_kwargs'] = algorithm_kwargs or {}
376        self.input['relax_cell'] = relax_cell
377        self.input['relax_cell_kwargs'] = relax_cell_kwargs or {}
378        self._generic_input['ionic_force_tolerance'] = ionic_force_tolerance
379        self._generic_input['n_print'] = n_print
380        self._generic_input['max_iter'] = max_iter
381        self.input['save_path'] = save_path
382    
383    def _write_calc_minimize(self):
384        """
385        Write the calculation script for performing structure relaxation.
386
387        This method generates a Python script that uses Grace for structure
388        relaxation and writes it to a file in the working directory.
389        """
390
391        model = self.input['model'] # Grace model
392        algo = self.input.get('algo', 'FIRE') # ASE optimizer
393        if algo not in ['FIRE', 'BFGS', 'LBFGS']:
394            raise ValueError(
395                f"Unsupported optimizer: {algo}. Supported optimizers are: FIRE, BFGS, LBFGS.")
396        # ASE optimizer kwargs
397        algo_kwargs = self.input.get('algo_kwargs', {}).copy()
398        save_path = self.input.get('save_path', 'relax.traj') # Add trajectory parameter
399        algo_kwargs['trajectory'] = f"{save_path}"
400        algo_kwargs['loginterval'] = self._generic_input['n_print']
401        # Reformat dictionary to string
402        algo_kwargs_str = ", ".join(
403            f'{k}="{v}"' if isinstance(v, str) else f"{k}={v}"
404            for k, v in algo_kwargs.items()
405        )
406
407        # Cell relaxation
408        relax_cell = self.input.get('relax_cell', False)
409        # Cell relaxation kwargs ---> see ase.constraints.ExpCellFilter
410        relax_cell_kwargs = self.input.get('relax_cell_kwargs', {})
411        # Reformat dictionary to string
412        relax_cell_kwargs_str = ", ".join(
413            f'{k}="{v}"' if isinstance(v, str) else f"{k}={v}"
414            for k, v in relax_cell_kwargs.items()
415        )
416
417        # Set parameters for minimization
418        params_run = f"fmax={self._generic_input['ionic_force_tolerance']}, steps={self._generic_input['max_iter']}" 
419
420        # Save parameters in        
421        self._generic_input['relax_cell'] = relax_cell
422        self._generic_input['save_path'] = save_path
423        
424        # Somewhat hidden feature of adding custom fixes
425        # TODO: Document? Make it a more "official" paramerer?
426        fixes = self.input.get('fixes', [])
427        if isinstance(fixes, str):
428            fixes = [fixes]
429        elif isinstance(fixes, list):
430            pass
431        else:
432            self.logger.warning('Fixes are neither string or array of strings. Ignoring')
433            fixes = []
434        fix_imports = self.input.get('fix_imports', [])
435        if isinstance(fix_imports, str):                
436            fix_imports = [fix_imports]
437        elif isinstance(fix_imports, list):
438            pass
439        else:
440            self.logger.warning('Fix imports are neither string or array of strings. Ignoring')
441            fix_imports = []
442            
443        script = [
444            'from tensorpotential.calculator import grace_fm',
445            'from ase.io import read',
446            'from ase.io.cif import write_cif',
447            'from ase.io.trajectory import Trajectory',
448            f'from ase.optimize import {algo}',
449        ]
450        if relax_cell:
451            script += ['from ase.constraints import ExpCellFilter as ECF']
452        for f in fix_imports:
453            script += [f]
454        script += ['']
455
456        script += [
457            f'struct = read("structure.cif")',
458            f'struct.calc = grace_fm("{model}")'
459        ]
460        for f in fixes:
461            script.append(f'struct.set_constraint({f})')
462        script += ['']
463        
464        if not relax_cell:
465            script += [
466                f'relaxation = {algo.upper()}(atoms=struct, {algo_kwargs_str}).run({params_run})',]
467        else:
468            script += [
469                f'relaxation = {algo.upper()}(ECF(atoms=struct, {relax_cell_kwargs_str}), {algo_kwargs_str}).run({params_run})']
470        script += ['']
471
472        script += ['write_cif("final_structure.cif", struct)']
473
474        with open(os.path.join(self.working_directory, 'calc_script.py'), 'w') as f:
475            f.writelines("\n".join(script))
476
477    def _parse_calc_minimize(self, output: str = 'error.out') -> Optional[pd.DataFrame]:
478        """
479        Parse the output of a minimization.
480
481        This function reads the minimization output file and returns the
482        corresponding DataFrame containing the optimization step and energy information.
483
484        Parameters
485        ----------
486        output : str, optional
487            The name of the output file to parse, by default 'error.out'.
488
489        Returns
490        -------
491        pd.DataFrame or None
492            A DataFrame containing the optimization results, or None if parsing fails.
493        """
494        # ASE optimizer
495        algo = self.input.get('algo', 'FIRE')
496
497        try:
498            # Find the header line dynamically
499            with open(os.path.join(self.working_directory, output), 'r') as f:
500                lines = f.readlines()
501            for i, line in enumerate(lines):
502                if line.strip().startswith('Step'):
503                    header_idx = i
504                    break
505            else:
506                self.logger.warning(f'No table header found in {output}.')
507                return None
508
509            table_lines = []
510            # Ensure that the line contains the expected algorithm
511            for line in lines[header_idx:]:
512                if algo in line.strip() or "Step" in line.strip():
513                    table_lines.append(line)
514
515            # Read table into DataFrame
516            table_str = ''.join(table_lines)
517            df = pd.read_csv(io.StringIO(table_str),
518                             sep=r'\s+', engine='python')
519
520            return df
521        except Exception as e:
522            self.logger.warning(f'Cannot parse output from {output} file: {e}')
523            return None
524
525
526    ################################################
527    ### --- IMPLEMENTATION OF MD CALCULATION --- ###
528    ################################################
529    def calc_md(
530        self,
531        grace_model: Literal[*_grace_models] = "GRACE-2L-OMAT",
532        ensemble: Literal["NVT", "nvt", "Langevin", "langevin", "NPT", "npt"] = "NVT", 
533        temperature: float = 300.0,
534        starting_temperature: int | None = None,
535        pressure: float | None = None,
536        time_step: float = 2.0,
537        n_ionic_steps: int = 100,
538        n_print: int = 10,
539        temperature_damping_timescale: float = 100.0,
540        pressure_damping_timescale: float | None = None,
541        trajectory: str = "md_run.traj",
542        logfile: str = "md_run.log",
543        random_seed: int | None = None,
544        mask: Tuple[int] | np.ndarray | None = None
545    ):
546        """
547        Run a molecular dynamics simulation using a GRACE foundation model.
548    
549        Parameters
550        ----------
551        grace_model : Literal[*_grace_models], optional
552            GRACE model to use for the relaxation.  
553            Must be one of the supported models, which can be obtained via 
554            :meth:`available_models`. Default is ``"GRACE-2L-OMAT"``.
555    
556        ensemble : {"NVT", "nvt", "Langevin", "langevin", "NPT", "npt"}, optional
557            Type of statistical ensemble:
558            - "NVT" or "Langevin": constant volume and temperature.
559            - "NPT": constant pressure and temperature.
560            Case-insensitive (default is "NVT").
561    
562        temperature : float, optional
563            Target simulation temperature in Kelvin (default is 300.0).
564    
565        starting_temperature : int or None, optional
566            Initial temperature of the system in Kelvin. If None, uses
567            the target `temperature` (default is None).
568    
569        pressure : float or None, optional
570            Target external pressure in GPa for NPT runs. Ignored if ensemble
571            is NVT or Langevin (default is None).
572    
573        time_step : float, optional
574            Integration time step in femtoseconds (default is 2.0 fs).
575    
576        n_ionic_steps : int, optional
577            Number of MD steps to perform (default is 100).
578    
579        n_print : int, optional
580            Frequency of log printing and trajectory writing (default is 10).
581    
582        temperature_damping_timescale : float, optional
583            Thermostat coupling time in femtoseconds (default is 100.0 fs).
584    
585        pressure_damping_timescale : float or None, optional
586            Barostat coupling time in femtoseconds. Only relevant for NPT runs
587            (default is None).
588    
589        trajectory : str, optional
590            Output trajectory file name (default is "md_run.traj").
591    
592        logfile : str, optional
593            Output log file name (default is "md_run.log").
594    
595        random_seed : int or None, optional
596            Seed for random velocity initialization. If None, system clock
597            is used (default is None).
598            
599        mask : None, sequence of 3 ints (0/1), or (3,3) array-like, optional
600            Specify which components of the computational box (strain tensor)
601            are allowed to change under the barostat (default is None equalling
602            to *No mask*). Only relevant for NPT simulations; it has no effect 
603            for NVT / Langevin ensembles.
604        """
605        self._generic_input['calc_mode'] = 'md'
606        ensemble = ensemble.lower()
607        if ensemble in ['langevin', 'nvt']: #requested NVT, ignore NPT settings
608            pressure = None
609            mask = None
610        if pressure == None:
611            pressure_damping_timescale = None
612        super().calc_md(
613            temperature=temperature,
614            pressure=pressure,
615            n_ionic_steps=n_ionic_steps,
616            time_step=time_step,
617            n_print=n_print,
618            temperature_damping_timescale=temperature_damping_timescale,
619            pressure_damping_timescale=pressure_damping_timescale,
620        )
621        self.input['model'] = grace_model
622        self.input['ensemble'] = ensemble.lower()
623        self.input['trajectory'] = trajectory
624        self.input['logfile'] = logfile 
625        if starting_temperature is None:
626            starting_temperature = temperature
627        self.input['starting_temperature'] = starting_temperature
628        self.input['random_seed'] = random_seed
629        self.input['mask'] = mask
630
631    def _write_calc_md(self):
632        """
633        Write the calculation script for performing molecular dynamics (MD).
634
635        This method generates a Python script that uses CHGNet for running MD
636        simulations and writes it to a file in the working directory.
637        """
638        # Somewhat hidden feature of adding custom fixes
639        # TODO: Document? Make it a more "official" paramerer?
640        fixes = self.input.get('fixes', [])
641        if isinstance(fixes, str):
642            fixes = [fixes]
643        elif isinstance(fixes, list):
644            pass
645        else:
646            self.logger.warning(
647                'Fixes are neither string or array of strings. Ignoring')
648            fixes = []
649        fix_imports = self.input.get('fix_imports', [])
650        if isinstance(fix_imports, str):
651            fix_imports = [fix_imports]
652        elif isinstance(fix_imports, list):
653            pass
654        else:
655            self.logger.warning(
656                'fix_imports are neither string or array of strings. Ignoring')
657            fix_imports = []
658        
659        script = [
660            'from tensorpotential.calculator import grace_fm',
661            'from ase.md.velocitydistribution import MaxwellBoltzmannDistribution',
662            'from ase.io import read',
663            'from ase.io.cif import write_cif',
664            'from ase.io.trajectory import Trajectory',
665            'from ase import units'
666        ]
667        for f in fix_imports:
668            script += [f]
669        script += [
670            '',
671            'struct = read("structure.cif")',]
672        for f in fixes:
673            script.append(f'struct.set_constraint({f})')
674        script += [
675            f'calc = grace_fm("{self.input["model"]}")',
676            'struct.calc = calc',
677            '',
678        ]
679        if not self.input['random_seed'] == None:
680            script += [
681                'import numpy as np',
682                f'np.random.seed({self.input["random_seed"]})',
683            ]
684        script += [
685            f'MaxwellBoltzmannDistribution(struct, temperature_K={self.input["starting_temperature"]})',
686            '',
687        ]
688        if self.input['ensemble'] == 'langevin':
689            script += [
690                'from ase.md.langevin import Langevin',
691                f'dyn = Langevin(struct, timestep={self._generic_input["time_step"]}, temperature_K={self._generic_input["temperature"]}, friction={self._generic_input["temperature_damping_timescale"]})',
692                ]
693        else:
694            pressure = self._generic_input["pressure"]
695            if not pressure == None:
696                pressure /= EV_PER_ANG3_TO_GPA
697            pfactor = self._generic_input["pressure_damping_timescale"]
698            if not pfactor == None:
699                pfactor = f'{pfactor}*units.fs'
700            script += [
701                'from ase.md.npt import NPT',                
702                f'dyn = NPT(struct, timestep={self._generic_input["time_step"]}*units.fs, temperature_K={self._generic_input["temperature"]}, ttime={self._generic_input["temperature_damping_timescale"]}*units.fs, externalstress={pressure}, pfactor={pfactor} , mask={self.input["mask"]})',
703                ]
704        script += [
705            f'traj = Trajectory("{self.input["trajectory"]}", "w", struct)',
706            f'dyn.attach(traj.write, interval={self._generic_input["n_print"]})',
707            '',
708            f'logfile = open("{self.input["logfile"]}", "w")',
709            'logfile.write("step\\tTime\\tEtot\\tEpot\\tEkin\\tT\\tVol\\ta\\tb\\tc\\tsigma_xx\\tsigma_yy\\tsigma_zz\\tsigma_yz\\tsigma_xz\\tsigma_xy\\n")',
710            'def print_status():',
711            '    values = [dyn.nsteps, dyn.nsteps*dyn.dt]',
712            '    values += [struct.get_potential_energy()+struct.get_kinetic_energy()]',
713            '    values += [struct.get_potential_energy(), struct.get_kinetic_energy()]',
714            '    values += [struct.get_temperature()]',
715            '    values += [struct.get_volume()]+list(struct.cell.lengths())',
716            '    values += list(struct.get_stress())',
717            '    status = "\\t".join(str(v) for v in values)',
718            '    print(status)',
719            '    logfile.write(status+"\\n")',
720            '    logfile.flush()',
721            f'dyn.attach(print_status, interval={self._generic_input["n_print"]})',
722            '',
723            f'dyn.run({self._generic_input["n_ionic_steps"]})',
724            '',
725            'write_cif("final_structure.cif", struct)'
726        ]
727
728        with open(os.path.join(self.working_directory, 'calc_script.py'), 'w') as f:
729            f.writelines("\n".join(script))
730
731    def _parse_calc_md(self, output: str = 'md_run.log') -> Optional[pd.DataFrame]:
732        """
733        Parse the output of a molecular dynamics (MD) simulation.
734
735        This function reads the MD log file and returns the corresponding DataFrame
736        containing time, total energy, potential energy, kinetic energy, and temperature.
737
738        Parameters
739        ----------
740        output : str, optional
741            The name of the MD log file to parse, by default 'md_run.log'.
742
743        Returns
744        -------
745        pd.DataFrame or None
746            A DataFrame containing the MD simulation results, or None if parsing fails.
747        """
748        try:
749            # df = pd.read_csv(
750            #     os.path.join(self.working_directory, output),
751            #     skiprows=1,
752            #     sep=r'\s+'
753            # )
754            df = pd.read_csv(
755                os.path.join(self.working_directory, output),
756                sep=r"\t+", 
757                engine='python'
758            )
759            return df
760        except:
761            self.logger.warning(f'Cannot parse output from {output} file.')
762            return None
763
764    ###############################
765    ### --- COMMON FUNCTION --- ###
766    ###############################
767    def write_input(self) -> None:
768        """
769        Write input files required for the Grace calculation.
770
771        This function writes the structure file (`structure.cif`), stores the input
772        dictionary to the HDF5 format, and delegates further writing depending on the
773        calculation mode (MD or static/minimize).
774
775        Returns
776        -------
777        None
778        """
779        write(os.path.join(self.working_directory, 'structure.cif'),
780              pyiron_to_ase(self.structure))
781        self.input.to_hdf(self.project_hdf5, group_name='input')
782        if self._generic_input['calc_mode'] == 'static':
783            self._write_calc_static()
784        elif self._generic_input['calc_mode'] == 'minimize':
785            self._write_calc_minimize()
786        elif self._generic_input['calc_mode'] == 'md':
787            self._write_calc_md()
788        else:
789            raise ValueError(
790                "No calculation type set. Use calc_static(), calc_minimize() or calc_md() first.")
791            
792    def collect_output(self) -> None:
793        """
794        Collect and store the results of a calculation.
795    
796        This method parses the final structure, trajectory, and other output properties
797        from the calculation and stores them in HDF5 format. The internal structure
798        is updated, and various properties such as energies, structure, forces, and 
799        stresses are saved to the output group.
800    
801        Returns
802        -------
803        None
804        """
805
806        def voigt_to_tensor(s):
807            # s is a 6-element array from atoms.get_stress()
808            return np.array([
809                [s[0], s[5], s[4]],
810                [s[5], s[1], s[3]],
811                [s[4], s[3], s[2]]
812            ])
813
814        relaxed_ase_structure = read(
815            os.path.join(self.working_directory, 'final_structure.cif'), format='cif')
816        relaxed_pyiron_structure = ase_to_pyiron(relaxed_ase_structure)
817
818        # Store the relaxed structure in the HDF5 group
819        with self.project_hdf5.open('output') as h5out:
820            h5out['structure'] = relaxed_pyiron_structure
821        self.structure = relaxed_pyiron_structure  # Update the job's structure
822
823        # Parse trajectory from minimize
824        if self._generic_input['calc_mode'] in ['minimize', 'md']:
825            # For minimize and md, we expect a trajectory file with the relaxed structure
826            if self._generic_input['calc_mode'] == 'minimize':
827                traj_file = self.input['save_path']
828                traj = Trajectory(os.path.join(self.working_directory, traj_file))
829            else:
830                traj_file = self.input['trajectory']
831                traj = read(os.path.join(self.working_directory, traj_file), index=":")
832            try:
833                with self.project_hdf5.open('output/generic') as h5out:
834                    h5out['cells'] = np.array([s.get_cell() for s in traj])
835                    h5out['positions'] = np.array([s.positions for s in traj])
836                    h5out['stresses'] = np.array(
837                        [voigt_to_tensor(s.get_stress())*EV_PER_ANG3_TO_GPA for s in traj])
838                    h5out['forces'] = np.array([s.get_forces() for s in traj])
839            except:
840                self.logger.warning(
841                    f'Cannot parse output from {traj_file} file.')
842        
843        # Parse other properties
844        if self._generic_input['calc_mode'] == 'static':
845            df = self._parse_calc_static()  
846            df['max_force'] = df['forces'].apply(lambda x: np.max(x))
847            with self.project_hdf5.open('output/generic') as h5out:
848                h5out['energy_pot'] = df['energy_pot'].values
849                h5out['energy_tot'] = df['energy_pot'].values
850                h5out['forces'] = df['forces'].values
851                h5out['max_force'] = df['max_force'].values
852                h5out['stresses'] = np.array(
853                    [voigt_to_tensor(s)*EV_PER_ANG3_TO_GPA for s in df['stresses'].values])
854                h5out['steps'] = df['step'].values
855                h5out['cells'] = np.array([relaxed_pyiron_structure.get_cell()])
856                h5out['positions'] = np.array([relaxed_pyiron_structure.positions])
857        elif self._generic_input['calc_mode'] == 'minimize':
858            df = self._parse_calc_minimize()
859            with self.project_hdf5.open('output/generic') as h5out:
860                h5out['energy_pot'] = df['Energy'].values
861                h5out['energy_tot'] = df['Energy'].values
862                h5out['max_force'] = df['fmax'].values
863                h5out['steps'] = df['Step'].values        
864        elif self._generic_input['calc_mode'] == 'md':
865            df = self._parse_calc_md(self.input['logfile'])
866            with self.project_hdf5.open('output/generic') as h5out:
867                h5out['energy_pot'] = df['Epot'].values
868                h5out['energy_tot'] = df['Etot'].values
869                h5out['energy_kin'] = df['Ekin'].values
870                h5out['temperature'] = df['T'].values
871                h5out['time'] = df['Time'].values
872                h5out['steps'] = df.index.values
873
874        self.compress()
875        self.status.finished = True
876
877    def get_structure(self, frame: int = -1) -> Atoms:
878        """
879        Retrieve the structure at a given frame from the calculation output.
880
881        This function returns the atomic structure (positions, cell, species) at the
882        specified frame index. The frame index is adjusted to ensure it is valid.
883
884        Parameters
885        ----------
886        frame : int, optional
887            The frame index to retrieve. If negative, it is counted from the last frame,
888            by default -1.
889
890        Returns
891        -------
892        Atoms
893            The ASE Atoms object representing the structure at the given frame.
894
895        Raises
896        ------
897        ValueError
898            If no structures are found in the job output.
899        IndexError
900            If the frame index is out of range.
901        """
902        num_structures = self.number_of_structures
903        if num_structures == 0:
904            raise ValueError("No structures found in the job output.")
905
906        # Ensure frame is valid
907        if frame < 0:
908            frame += num_structures
909        if frame >= num_structures:
910            raise IndexError(
911                f"Frame {frame} is out of range for {num_structures} structures.")
912
913        with self.project_hdf5.open('output/generic') as h5out:
914            positions = h5out["positions"][frame]
915            cell = h5out["cells"][frame]
916        # Read positions and cell
917        with self.project_hdf5.open('output/structure') as h5out:
918            species = h5out["species"]
919            indices = h5out["indices"]
920
921        # Generate symbols list
922        symbols = [species[i] for i in indices]
923        return ase_to_pyiron(Atoms(symbols=symbols, positions=positions, cell=cell, pbc=True))
924
925    def compress(self, files_to_compress=None):
926        """
927        Compress the output files of a job object.
928    
929        Args:
930            files_to_compress (list): A list of files to compress (optional)
931        """
932        import os
933        # delete empty files
934        for f in self.files.list():
935            filename = os.path.join(self.working_directory, f)
936            if (os.path.exists(filename) and os.stat(filename).st_size == 0):
937                os.remove(filename)
938        # delete slurm scripts
939        if "run_queue.sh" in self.files.list():
940            filename = os.path.join(self.working_directory, "run_queue.sh")
941            os.remove(filename)
942        # compress all remaining files
943        if files_to_compress is None:
944            files_to_compress = [
945                f
946                for f in self.files.list()
947                if f
948                not in [
949                    "final_structure.cif",
950                ]
951            ]
952        
953        super(Grace, self).compress(files_to_compress=files_to_compress)
954
955
956class CalcParams(GenericParameters):
957    """
958    Class to manage the calculation parameters for Grace calculations.
959
960    This class extends from `GenericParameters` and is used to define and manage
961    parameters specific to Grace calculations. It allows for initialization of
962    parameters and stores them in a table, with a default table name of "grace".
963    The class can be used to load and store various Grace calculation settings and
964    other related configurations.
965
966    Parameters
967    ----------
968    table_name : str, optional
969        The name of the table where the parameters are stored. By default, it is set
970        to "grace", but it can be customized if necessary.
971
972    Methods
973    -------
974    __init__(self, table_name="grace")
975        Initializes the parameters for the calculation.
976    """
977
978    def __init__(self, table_name: str = "grace") -> None:
979        """
980        Initialize the CalcParams object with the given table name.
981
982        Parameters
983        ----------
984        table_name : str, optional
985            The name of the table where the parameters are stored. By default,
986            it is set to "grace".
987        """
988        super(CalcParams, self).__init__(
989            table_name=table_name,
990        )