[WIP] Basic LAMMPS flows implementation

pyproject.toml

@@ -40,6 +40,7 @@ amset = ["amset>=0.4.15", "pydash"]
 cclib = ["cclib"]
 mp = ["mp-api>=0.27.5"]
 phonons = ["phonopy>=1.10.8", "seekpath"]
+lammps = ["pymatgen-io-lammps"]
 docs = [
     "numpydoc==1.5.0",
     "ipython==8.6.0",

requirements.txt

@@ -0,0 +1,8 @@
+jobflow
+fireworks
+paramiko
+fabric
+requests
+matplotlib
+pymatgen
+atomate2

src/atomate2/lammps/files.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+
+from pymatgen.core import Structure
+from pymatgen.io.lammps_new.sets import BaseLammpsGenerator
+
+
+def write_lammps_input_set(
+    structure: Structure,
+    input_set_generator: BaseLammpsGenerator,
+    directory: str | Path = ".",
+    **kwargs,
+):
+    input_set = input_set_generator.get_input_set(structure)
+    input_set.write_input(directory)

src/atomate2/lammps/flows.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass, field
+
+from jobflow import Flow, Maker
+from pymatgen.core import Structure
+
+from .jobs.base import LammpsMaker
+from .jobs.core import MinimizationMaker
+
+
+@dataclass
+class FMinimizationMaker(Maker):
+    name: str = "minimization"
+    minimization_maker: LammpsMaker = field(default_factory=MinimizationMaker)
+
+    def make(self, structure):
+        return Flow([self.minimization_maker.make(structure)], name=self.name)

src/atomate2/lammps/jobs/base.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable
+
+from jobflow import Maker, Response, job
+from monty.serialization import dumpfn
+from monty.shutil import gzip_dir
+from pymatgen.core import Structure
+from pymatgen.io.lammps_new.sets import BaseLammpsGenerator
+
+from ..files import write_lammps_input_set
+from ..schemas.task import TaskDocument
+
+_DATA_OBJECTS = []
+
+__all__ = ("BaseLammpsMaker", "lammps_job")
+
+from ..run import run_lammps
+
+
+def lammps_job(method: Callable):
+    return job(method, data=_DATA_OBJECTS, output_schema=TaskDocument)
+
+
+@dataclass
+class BaseLammpsMaker(Maker):
+    name: str = "Base LAMMPS job"
+    input_set_generator: BaseLammpsGenerator = field(default_factory=BaseLammpsGenerator)
+    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: dict = field(default_factory=dict)
+
+    @lammps_job
+    def make(self, input_structure: Structure | Path):
+        """Run a LAMMPS calculation."""
+
+        write_lammps_input_set(
+            input_structure, self.input_set_generator, **self.write_input_set_kwargs
+        )
+
+        for filename, data in self.write_additional_data.items():
+            dumpfn(data, filename.replace(":", "."))
+
+        run_lammps(**self.run_lammps_kwargs)
+
+        task_doc = TaskDocument.from_directory(Path.cwd(), **self.task_document_kwargs)
+        task_doc.task_label = self.name
+
+        gzip_dir(".")
+
+        return Response(output=task_doc)
+
+
+@dataclass
+class LammpsMaker(BaseLammpsMaker):
+    name: str = "Simple LAMMPS job"
+    input_set_generator: BaseLammpsGenerator = field(default_factory=BaseLammpsGenerator)

src/atomate2/lammps/jobs/core.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass, field
+
+from pymatgen.io.lammps_new.sets import LammpsMinimization, BaseLammpsGenerator
+
+from ..jobs.base import BaseLammpsMaker
+
+
+@dataclass
+class MinimizationMaker(BaseLammpsMaker):
+    name: str = "minimization"
+    input_set_generator: BaseLammpsGenerator = field(default_factory=LammpsMinimization)

src/atomate2/lammps/run.py

@@ -0,0 +1,82 @@
+import subprocess
+from pathlib import Path
+
+from .settings import LAMMPS_SETTINGS
+
+
+def run_lammps(
+    lammps_input_file: str,
+    lammps_cmd: str = LAMMPS_SETTINGS.LAMMPS_CMD,
+    lammps_suffix: list[str] | str | None = LAMMPS_SETTINGS.LAMMPS_SUFFIX,
+    lammps_pks: list[str] | str | None = LAMMPS_SETTINGS.LAMMPS_PACKAGES,
+    lammps_run_flags: list[str] | str | None = None,
+    mpi_cmd: str | None = LAMMPS_SETTINGS.MPI_CMD,
+    mpi_num_processes: int = 1,
+    mpi_num_processes_flag: str = LAMMPS_SETTINGS.MPI_NUM_PROCESSES_FLAG,
+    max_walltime_hours: float | None = None,
+    stdout_file: str | Path = "stdout.log",
+    stderr_file: str | Path = "stderr.log",
+) -> subprocess.Popen:
+    """Run LAMMPS.
+
+    Parameters:
+        lammps_input_file: The path to the main input file to be passed to the
+            LAMMPS executable with the `-in` command-line option.
+        lammps_cmd: The name or path to the LAMMPS executable.
+        lammps_suffix: The suffix to use that applies style variants at runtime
+            (see `LammpsSettings.LAMMPS_SUFFIX`).
+        lammps_pks: The runtime packages and options to tell LAMMPS to use
+            (see `LammpsSettings.LAMMPS_PACKAGES`).
+        lammps_run_kwargs: Any additional arbitrary flags to invoke LAMMPS with.
+        mpi_cmd: The command to invoke MPI (e.g., `'mpirun'` or `'mpiexec'`).
+            If None, invoke the `lammps_cmd` in serial mode.
+        mpi_num_processes: The number of MPI processes to pass to `mpi_cmd` with
+            the `mpi_num_processes_flag`.
+        mpi_num_processes_flag: The command-line flag to use with MPI to pass the number of processes,
+            e.g., `-n/n` [default] or `np/-np`.
+        max_walltime_hours: The maximum walltime in hours to allow for the task. If provided, attempts will
+            be made to cleanly end the calculation after this amount of time. Note: if using with a queueing system,
+            this value should leave suffcient time for the clean-up of the calculation within the maximum walltime
+            allocated to the job by the queue.
+        stdout_file: The name of or path to a file in which to save the stdout stream.
+        stderr_file: The name of or path to a file in which to save the stderr stream.
+
+    """
+
+    lammps_invocation: list[str] = []
+
+    if mpi_cmd is not None:
+        lammps_invocation.extend(
+            [
+                mpi_cmd,
+                f"-{mpi_num_processes_flag.lstrip('-')}",
+                str(mpi_num_processes),
+            ]
+        )
+
+    if lammps_suffix is not None:
+        if isinstance(lammps_suffix, str):
+            lammps_suffix = [lammps_suffix]
+        for sf in lammps_suffix:
+            lammps_invocation += ["-sf", sf]
+
+    if lammps_pks is not None:
+        if isinstance(lammps_pks, str):
+            lammps_pks = [lammps_pks]
+        for pk in lammps_pks:
+            lammps_invocation += ["-pk", pk]
+
+    if lammps_run_flags is not None:
+        if isinstance(lammps_run_flags, str):
+            lammps_run_flags = [lammps_run_flags]
+        for flag in lammps_run_flags:
+            lammps_invocation += [flag]
+
+    lammps_invocation.extend([lammps_cmd, "-in", lammps_input_file])
+
+    with open(stdout_file, "a") as stdout, open(stderr_file, "a") as stderr:
+        return subprocess.Popen(
+            lammps_invocation,
+            stdout=stdout,
+            stderr=stderr,
+        )

src/atomate2/lammps/schemas/task.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+from typing import Type
+
+from atomate2.common.schemas.structure import StructureMetadata
+from pydantic import Field
+
+
+class TaskDocument(StructureMetadata):
+
+    dir_name: str = Field()
+
+    task_label: str = Field()
+
+    @classmethod
+    def from_directory(
+        cls: Type["TaskDocument"],
+        dir_name: str | Path,
+    ) -> "TaskDocument":
+        return TaskDocument(dir_name=dir_name)
+
+    class Config:
+        extras = "allow"

src/atomate2/lammps/sets/base.py

@@ -0,0 +1,3 @@
+from pymatgen.io.lammps_new.sets import LammpsInputSet
+
+__all__ = "LammpsInputSet"

src/atomate2/lammps/settings.py

@@ -0,0 +1,36 @@
+from pydantic import BaseSettings, Field
+
+
+class LammpsSettings(BaseSettings):
+    """Settings that control the invocation of LAMMPS."""
+
+    LAMMPS_CMD: str = Field("lmp", description="The command to run LAMMPS.")
+    LAMMPS_SUFFIX: list[str] | str | None = Field(
+        None,
+        description=(
+            "The LAMMPS style suffix(es) to use."
+            "See https://docs.lammps.org/Run_options.html#suffix for more information."
+        ),
+        examples=["gpu", "kk", "intel", "omp", "opt", [["gpu", "kk"]]],
+    )
+    LAMMPS_PACKAGES: list[str] | str | None = Field(
+        None,
+        description=(
+            "Options to pass to the package command-line flag that controls subpackage "
+            "styles and parameters, e.g., `'gpu 1'` will call `lmp -pk gpu 1` will tell "
+            "LAMMPS to use 1 GPU for this calculation. "
+            "List values are passed with separate '-pk' invocations, e.g., `lmp -pk gpu 1 -pk omp 4`."
+            "See https://docs.lammps.org/Run_options.html#package for more information."
+        ),
+        examples=["gpu 0", "gpu 1 split 0.75", "gpu 2 split -1.0", "gpu 1 omp 4"],
+    )
+    MPI_CMD: str = Field("mpirun", description="The command to invoke MPI.")
+    MPI_NUM_PROCESSES_FLAG: str = Field(
+        "-n",
+        description="The flag with which to provide the number of processes to use in the MPI execution.",
+    )
+
+
+LAMMPS_SETTINGS = LammpsSettings()
+
+__all__ = ("LAMMPS_SETTINGS", "LammpsSettings")