materialsproject
diff --git a/‎.github/workflows/testing.yml‎
Lines changed: 1 addition & 1 deletion b/‎.github/workflows/testing.yml‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎src/atomate2/lammps/__init__.py‎
Lines changed: 1 addition & 0 deletions b/‎src/atomate2/lammps/__init__.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎src/atomate2/lammps/files.py‎
Lines changed: 129 additions & 0 deletions b/‎src/atomate2/lammps/files.py‎
Lines changed: 129 additions & 0 deletions
diff --git a/‎src/atomate2/lammps/flows/__init__.py‎
Lines changed: 1 addition & 0 deletions b/‎src/atomate2/lammps/flows/__init__.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎src/atomate2/lammps/flows/core.py‎
Lines changed: 79 additions & 0 deletions b/‎src/atomate2/lammps/flows/core.py‎
Lines changed: 79 additions & 0 deletions
diff --git a/‎src/atomate2/lammps/jobs/__init__.py‎
Lines changed: 4 additions & 0 deletions b/‎src/atomate2/lammps/jobs/__init__.py‎
Lines changed: 4 additions & 0 deletions
diff --git a/‎src/atomate2/lammps/jobs/base.py‎
Lines changed: 150 additions & 0 deletions b/‎src/atomate2/lammps/jobs/base.py‎
Lines changed: 150 additions & 0 deletions
@@ -290,7 +290,7 @@ jobs:
           MP_API_KEY: ${{ secrets.MP_API_KEY }}
         run: |
           micromamba activate a2
-          pytest -n auto --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb --ignore=./tutorials/force_fields --ignore=./tutorials/torchsim_tutorial.ipynb
+          pytest -n auto --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb --ignore=./tutorials/force_fields --ignore=./tutorials/torchsim_tutorial.ipynb --ignore=./tutorials/lammps_workflow.ipynb 
 
       - name: Test ASE
         env:
 
@@ -0,0 +1 @@
+"""Define LAMMPS jobs and workflows."""
@@ -0,0 +1,129 @@
+"""File I/O functions for LAMMPS input files."""
+
+from pathlib import Path
+from typing import Any, Literal
+
+from ase.io import Trajectory as AseTrajectory
+from ase.io import read
+from emmet.core.vasp.calculation import StoreTrajectoryOption
+from monty.serialization import dumpfn
+from numpy.typing import ArrayLike
+from pymatgen.core import Lattice, Molecule, Structure
+from pymatgen.core.trajectory import Trajectory as PmgTrajectory
+from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.io.lammps.data import CombinedData, LammpsBox, LammpsData
+from pymatgen.io.lammps.generators import BaseLammpsGenerator
+
+
+def write_lammps_input_set(
+    data: Structure | Molecule | LammpsData | CombinedData,
+    input_set_generator: BaseLammpsGenerator,
+    box_or_lattice: Lattice | LammpsBox | ArrayLike | None = None,
+    additional_data: LammpsData | CombinedData | None = None,
+    directory: str | Path = ".",
+) -> None:
+    """Write LAMMPS input set to a directory."""
+    input_set = input_set_generator.get_input_set(
+        data=data, additional_data=additional_data, box_or_lattice=box_or_lattice
+    )
+    input_set.write_input(directory)
+
+
+class DumpConvertor:
+    """
+    Class to convert LAMMPS dump files to pymatgen or ase Trajectory objects.
+
+    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: str,
+        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) -> AseTrajectory:
+        """Convert to ASE trajectory object."""
+        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:
+        """Convert to pymatgen trajectory object."""
+        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(),
+                        properties={"box": atoms.get_cell().tolist()},
+                    )
+                )
+        traj_method = "from_structures" if self.is_periodic else "from_molecules"
+        pmg_traj = getattr(PmgTrajectory, traj_method)(
+            frames,
+            frame_properties=frame_properties or 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"
+    ) -> Any:
+        """Save the trajectory to a file."""
+        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)
+        return None
@@ -0,0 +1 @@
+"""Lammps flow makers for atomate2."""
@@ -0,0 +1,79 @@
+"""Core LAMMPS flows."""
+
+from copy import deepcopy
+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 LammpsNPTMaker, LammpsNVTMaker
+
+
+@dataclass
+class MeltQuenchThermalizeMaker(Maker):
+    """Melt -> Quench -> Thermalize flow maker."""
+
+    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) -> Flow:
+        """Make the flow for melting, quenching, and thermalizing a structure."""
+        melt = self.melt_maker.make(structure)
+        quench = self.quench_maker.make(
+            melt.output.structure, prev_dir=melt.output.dir_name
+        )
+        thermalize = self.thermalize_maker.make(
+            quench.output.structure, prev_dir=quench.output.dir_name
+        )
+        return Flow([melt, quench, thermalize], name=self.name)
+
+    @classmethod
+    def from_temperature_steps(
+        cls,
+        npt_maker: LammpsNPTMaker,
+        nvt_maker: LammpsNVTMaker = None,
+        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,
+    ) -> "MeltQuenchThermalizeMaker":
+        """Make a melt-quench-thermalize flow maker from temperature and steps."""
+        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,
+        )
@@ -0,0 +1,4 @@
+"""Lammps job makers for atomate2."""
+
+from .base import BaseLammpsMaker
+from .core import CustomLammpsMaker, LammpsNPTMaker, LammpsNVTMaker, MinimizationMaker
@@ -0,0 +1,150 @@
+"""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 Molecule, Structure
+from pymatgen.io.lammps.generators import (
+    BaseLammpsSetGenerator,
+    CombinedData,
+    LammpsData,
+    LammpsForceField,
+)
+
+from atomate2.common.files import gunzip_files, 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] = [
+    "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 | LammpsForceField | 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:
+            if isinstance(self.force_field, dict):
+                self.force_field = LammpsForceField.from_dict(self.force_field)
+            self.input_set_generator.force_field = self.force_field
+
+    @lammps_job
+    def make(
+        self,
+        input_structure: Structure | Molecule | Path | LammpsData = None,
+        prev_dir: Path | str = 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'!"
+                )
+
+            restart_file = restart_files[0]
+            if restart_file.endswith(".restart.gz"):
+                gunzip_files(
+                    directory=prev_dir, include_files=[restart_file], force=True
+                )
+                restart_file = str(Path(restart_file).with_suffix(""))
+            self.input_set_generator.update_settings({"read_restart": restart_file})
+
+        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}")
+
+        # TODO: Only gzip LAMMPS files, not job scheduler related files
+        gzip_files(".")
+
+        return Response(output=task_doc)
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+"""Define LAMMPS jobs and workflows."""`