[WIP] LAMMPS Flows

src/atomate2/lammps/files.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+from pymatgen.core import Structure, Molecule
+from pymatgen.io.lammps.generators import BaseLammpsGenerator
+from ase.io import read
+from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.io.lammps.data import LammpsData, CombinedData
+from pymatgen.core.trajectory import Trajectory as PmgTrajectory
+from ase.io import Trajectory as AseTrajectory
+from typing import Literal
+from monty.serialization import dumpfn
+from monty.json import MSONable
+from emmet.core.vasp.calculation import StoreTrajectoryOption
+
+def write_lammps_input_set(
+    data: Structure | LammpsData | CombinedData,
+    input_set_generator: BaseLammpsGenerator,
+    additional_data : LammpsData | CombinedData | None = None,
+    directory: str | Path = ".",
+):
+    input_set = input_set_generator.get_input_set(data, 
+                                                  additional_data,
+                                                  )
+    input_set.write_input(directory)
+
+
+class DumpConvertor(MSONable):    
+    '''
+    Class to convert LAMMPS dump files to pymatgen or ase Trajectory objects. Based on TrajectoryObserver from atomate2.ase.
+
+    args:
+        dumpfile : str
+            Path to the LAMMPS dump file
+        store_md_outputs : StoreTrajectoryOption
+            Option to store MD outputs in the Trajectory object
+        read_index : str | int
+            Index of the frame to read from the dump file (default is ':', i.e. read all frames). 
+            Use an integer to read a specific frame (practical for large files).
+
+    '''
+    def __init__(self, dumpfile, store_md_outputs : StoreTrajectoryOption  = StoreTrajectoryOption.NO, read_index: str | int = ':') -> None:
+        self.store_md_outputs = store_md_outputs
+        self.traj = read(dumpfile, index=read_index) if isinstance(read_index, str) else [read(dumpfile, index=read_index)]
+        self.is_periodic = any(self.traj[0].pbc)
+        self.frame_properties_keys = ['forces', 'velocities']
+
+    def to_ase_trajectory(self, filename : str | None = None):
+        for idx, atoms in enumerate(self.traj):
+            with AseTrajectory(filename, 'a' if idx > 0 else 'w', atoms = atoms) as file: #check logic here
+                file.write()
+        return AseTrajectory(filename, 'r')
+
+    def to_pymatgen_trajectory(self, filename : str | None = None) -> PmgTrajectory:
+
+        species = AseAtomsAdaptor.get_structure(self.traj[0], cls = Structure if self.is_periodic else Molecule).species
+
+        frames  = []
+        frame_properties = []
+
+        for atoms in self.traj:
+            if self.store_md_outputs == StoreTrajectoryOption.FULL:
+                frame_properties.append({key : getattr(atoms, f'get_{key}')() for key in self.frame_properties_keys})
+
+            if self.is_periodic:
+                frames.append(Structure(lattice = atoms.get_cell(), 
+                                        species = species, 
+                                        coords = atoms.get_positions(), 
+                                        coords_are_cartesian = True)
+                                )
+            else:
+                frames.append(Molecule(species = species, 
+                                        coords = atoms.get_positions(), 
+                                        charge = atoms.get_charges(),
+                                        )
+                                  )
+        traj_method = 'from_structures' if self.is_periodic else 'from_molecules'
+        pmg_traj = getattr(PmgTrajectory, traj_method)(
+            frames,
+            frame_properties=frame_properties if frame_properties else None,
+            constant_lattice=False,
+        )
+
+        if filename:
+            dumpfn(pmg_traj, filename)
+
+        return pmg_traj
+
+    def save(self, filename : str | None = None, fmt : Literal["pmg", "ase"] = "pmg"):
+        filename = str(filename) if filename is not None else None
+        if fmt == "pmg" and filename:
+            return self.to_pymatgen_trajectory(filename=filename) 
+        if fmt == "ase" and filename:
+            return self.to_ase_trajectory(filename=filename)

src/atomate2/lammps/flows/core.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+
+from jobflow import Flow, Maker
+from pymatgen.core import Structure
+
+from atomate2.lammps.jobs.base import BaseLammpsMaker
+from atomate2.lammps.jobs.core import LammpsNVTMaker, LammpsNPTMaker
+
+from copy import deepcopy
+
+import warnings
+
+
+@dataclass
+class MeltQuenchThermalizeMaker(Maker):
+    # potentially remove and replace with single job
+    name: str = "melt-quench-thermalize"
+    melt_maker: BaseLammpsMaker = field(default_factory=LammpsNPTMaker)
+    quench_maker: BaseLammpsMaker = field(default_factory=LammpsNPTMaker)
+    thermalize_maker: BaseLammpsMaker = field(default_factory=LammpsNVTMaker)
+
+    def make(self, structure: Structure):
+        melt = self.melt_maker.make(structure)
+        quench = self.quench_maker.make(melt.output.structure)
+        thermalize = self.thermalize_maker.make(quench.output.structure)
+        return Flow([melt, quench, thermalize], name=self.name)
+
+    @classmethod
+    def from_temperature_steps(
+        cls,
+        start_temperature: float = 300,
+        melt_temperature: float = 3000,
+        quench_temperature: float = 300,
+        n_steps_melt: int = 10000,
+        n_steps_quench: int = 10000,
+        n_steps_thermalize: int = 10000,
+        npt_maker : LammpsNPTMaker = None,
+        nvt_maker : LammpsNVTMaker = None
+    ) -> "MeltQuenchThermalizeMaker":
+        if nvt_maker is None:
+            warnings.warn("No NVT maker provided, using NPT maker for thermalize.")
+        if npt_maker is None:
+            raise ValueError("NPT maker must be provided.")
+
+        melt_maker = deepcopy(npt_maker)
+        melt_maker.name = "melt"
+        melt_maker.input_set_generator.update_settings({"start_temp" : start_temperature,
+                                                        "end_temp": melt_temperature,
+                                                        "nsteps" : n_steps_melt})
+
+        quench_maker = deepcopy(npt_maker)
+        quench_maker.name = "quench"
+        quench_maker.input_set_generator.update_settings({"start_temp" : melt_temperature,
+                                                        "end_temp": quench_temperature,
+                                                        "nsteps" : n_steps_quench})
+
+        thermalize_maker = deepcopy(nvt_maker) if nvt_maker else deepcopy(npt_maker)
+        thermalize_maker.name = "thermalize"
+        thermalize_maker.input_set_generator.update_settings({"start_temp" : quench_temperature,
+                                                "end_temp": quench_temperature,
+                                                "nsteps" : n_steps_thermalize})
+        return cls(
+            melt_maker=melt_maker,
+            quench_maker=quench_maker,
+            thermalize_maker=thermalize_maker,
+        )

src/atomate2/lammps/jobs/__init__.py

@@ -0,0 +1,4 @@
+"""Lammps job makers for atomate2."""
+
+from .base import BaseLammpsMaker
+from .core import CustomLammpsMaker, LammpsNPTMaker, LammpsNVTMaker, MinimizationMaker

src/atomate2/lammps/jobs/base.py

@@ -0,0 +1,148 @@
+"""Base job maker for LAMMPS calculations."""
+
+import glob
+import os
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from emmet.core.vasp.task_valid import TaskState
+from jobflow import Maker, Response, job
+from pymatgen.core import Structure
+from pymatgen.io.lammps.generators import (
+    BaseLammpsSetGenerator,
+    CombinedData,
+    LammpsData,
+)
+
+from atomate2.common.files import gzip_files
+from atomate2.lammps.files import write_lammps_input_set
+from atomate2.lammps.run import run_lammps
+from atomate2.lammps.schemas.task import LammpsTaskDocument, StoreTrajectoryOption
+
+_DATA_OBJECTS: list[str] = [
+    "raw_log_file",
+    "inputs",
+    "trajectories",
+    "dump_files",
+]
+
+__all__ = ("BaseLammpsMaker", "lammps_job")
+
+
+class LammpsRunError(Exception):
+    """Custom exception for LAMMPS jobs."""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+        self.message = message
+
+
+def lammps_job(method: Callable) -> job:
+    """Job decorator for LAMMPS jobs."""
+    return job(method, data=_DATA_OBJECTS, output_schema=LammpsTaskDocument)
+
+
+@dataclass
+class BaseLammpsMaker(Maker):
+    """
+    Basic Maker class for LAMMPS jobs.
+
+    name: str
+        Name of the job
+    input_set_generator: BaseLammpsGenerator
+        Input set generator for the job, default is the BaseLammpsSetGenerator.
+        Check the sets module for more options on input kwargs.
+    write_input_set_kwargs: dict
+        Additional kwargs to write_lammps_input_set
+    run_lammps_kwargs: dict
+        Additional kwargs to run_lammps
+    task_document_kwargs: dict
+        Additional kwargs to TaskDocument.from_directory
+    write_additional_data: dict
+        Additional data to write to the job directory
+    """
+
+    name: str = "Base LAMMPS job"
+    input_set_generator: BaseLammpsSetGenerator = field(
+        default_factory=BaseLammpsSetGenerator
+    )
+    force_field: str | dict | None = field(default=None)
+    write_input_set_kwargs: dict = field(default_factory=dict)
+    run_lammps_kwargs: dict = field(default_factory=dict)
+    task_document_kwargs: dict = field(default_factory=dict)
+    write_additional_data: LammpsData | CombinedData = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        """Post-initialization warnings for the job."""
+        if (
+            self.task_document_kwargs.get("store_trajectory", StoreTrajectoryOption.NO)
+            != StoreTrajectoryOption.NO
+        ):
+            warnings.warn(
+                "Trajectory data might be large, only store if absolutely necessary. \
+                Consider manually parsing the dump files instead.",
+                stacklevel=1,
+            )
+
+        if self.force_field:
+            self.input_set_generator.force_field = self.force_field
+
+    @lammps_job
+    def make(
+        self,
+        input_structure: Structure | Path | LammpsData = None,
+        prev_dir: Path = None,
+    ) -> Response:
+        """Run a LAMMPS calculation."""
+        if prev_dir:
+            restart_files = glob.glob(os.path.join(prev_dir, "*restart*"))
+            if len(restart_files) != 1:
+                raise FileNotFoundError(
+                    "No/More than one restart file found in the previous directory. \
+                        If present, it should have the extension '.restart'!"
+                )
+
+            self.input_set_generator.update_settings(
+                {"read_restart": os.path.join(prev_dir, restart_files[0])}
+            )
+
+        if isinstance(input_structure, Path):
+            input_structure = LammpsData.from_file(
+                input_structure,
+                atom_style=self.input_set_generator.settings.get("atom_style", "full"),
+            )
+
+        write_lammps_input_set(
+            data=input_structure,
+            input_set_generator=self.input_set_generator,
+            additional_data=self.write_additional_data,
+            **self.write_input_set_kwargs,
+        )
+
+        run_lammps(**self.run_lammps_kwargs)
+
+        task_doc = LammpsTaskDocument.from_directory(
+            os.getcwd(), task_label=self.name, **self.task_document_kwargs
+        )
+
+        if task_doc.state == TaskState.ERROR:
+            try:
+                error = ""
+                for index, line in enumerate(task_doc.raw_log_file.split("\n")):
+                    if "ERROR" in line:
+                        error = error.join(task_doc.raw_log_file.split("\n")[index:])
+                        break
+            except ValueError:
+                error = "could not parse log file"
+            raise LammpsRunError(f"Task {task_doc.task_label} failed, error: {error}")
+
+        task_doc.composition = input_structure.composition
+        task_doc.reduced_formula = input_structure.composition.reduced_formula
+        task_doc.task_label = self.name
+        task_doc.inputs = self.input_set_generator.settings.dict
+
+        gzip_files(".")
+
+        return Response(output=task_doc)