initial idea of a Lammps Node

pyproject.toml

@@ -5,18 +5,20 @@ description = ""
 authors = ["zincwarecode <[email protected]>"]
 license = "Apache-2.0"
 readme = "README.md"
+include = ["znlib/data"]
 
 [tool.poetry.dependencies]
 python = ">=3.8,<4.0.0"
 colorama = "^0.4.5"
 zntrack = { version = "^0.4.3", optional = true }
 matplotlib = { version = "^3.6.1", optional = true }
 ase = { version = "^3.22.1", optional = true }
+jinja2 = { version = "^3.1.2", optional = true }
 
 
 [tool.poetry.extras]
 zntrack = ["zntrack", "matplotlib"]
-atomistic = ["ase", "zntrack"]
+atomistic = ["ase", "zntrack", "jinja2"]
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.2.0"
@@ -36,6 +38,7 @@ matplotlib = "^3.6.1"
 
 [tool.poetry.group.atomistic.dependencies]
 ase = "^3.22.1"
+jinja2 = "^3.1.2"
 
 [tool.poetry.urls]
 repository = "https://github.com/zincware/znlib"

tests/test_atomistic.py

@@ -23,3 +23,45 @@ def test_AddData(proj_path, tetraeder_test_traj):
     assert isinstance(loaded_data.atoms[[0, 1]], list)
     assert isinstance(loaded_data.atoms[:], list)
     assert isinstance(loaded_data.atoms.tolist(), list)
+
+
+def test_LammpsSimulator(proj_path):
+    simulator = znlib.atomistic.LammpsSimulator(
+        extras={"temperature": 1234, "npt": {"steps": 100}}
+    )
+    simulator.system.units = "lj"
+
+    template = simulator.get_template()
+
+    # modify the template by extra arguments.
+    with open(template, "a") as file:
+        file.write("velocity       all create {{ temperature }} 132465")
+        file.write("run             {{ npt.steps }}")
+
+    simulator.write_graph()
+    simulator = simulator.load()
+    simulator.render_template()
+    assert simulator.system.units == "lj"
+    assert simulator.extras == {"temperature": 1234, "npt": {"steps": 100}}
+
+    checks = 0  # count all assertions and make sure all where called
+
+    with open(simulator.input_file) as file:
+        for line in file:
+            if "units" in line:
+                assert simulator.system.units in line
+                checks += 1
+            if "boundary" in line:
+                assert simulator.system.boundary in line
+                checks += 1
+            if "atom_style" in line:
+                assert simulator.system.atom_style in line
+                checks += 1
+            if "velocity" in line:
+                assert str(simulator.extras["temperature"]) in line
+                checks += 1
+            if "run" in line:
+                assert str(simulator.extras["npt"]["steps"]) in line
+                checks += 1
+
+    assert checks == 5

znlib/atomistic/__init__.py

@@ -1,5 +1,6 @@
 """The znlib atomistic interface"""
 from znlib.atomistic import ase
 from znlib.atomistic.ase import FileToASE
+from znlib.atomistic.lammps import LammpsSimulator
 
-__all__ = ["ase", "FileToASE"]
+__all__ = ["ase", "FileToASE", "LammpsSimulator"]

znlib/atomistic/lammps.py

@@ -0,0 +1,84 @@
+"""Lammps Node"""
+import dataclasses
+import importlib.resources
+import pathlib
+import shutil
+import subprocess
+
+import jinja2
+from zntrack import Node, dvc, meta, utils, zn
+
+from znlib import data
+
+
+@dataclasses.dataclass
+class LammpsSystem:
+    """Group Lammps parameters in dataclasses"""
+
+    units: str = "metal"
+    atom_style: str = "charge"
+    boundary: str = "p p p"
+
+
+class LammpsSimulator(Node):
+    """A Lammps Node
+
+    References
+    ----------
+    https://docs.lammps.org/
+
+    Attributes
+    ----------
+    lammps_binary: str, default = "lmp"
+        The name / path to the lammps binary
+
+    template: str, default = "lammps.jinja2"
+        The name of the jinja2 template file. Use get_template() as a starting point.
+
+    system: LammpsSystem,
+        a default Dataclass used in most lammps simulations
+    extras: dict
+        A dictionary of additional template variables that are by default not part
+        of the 'lammps.jinja2' template shipped with znlib.
+    """
+
+    lammps_binary = meta.Text("lmp")
+
+    template: str = dvc.deps("lammps.jinja2")
+
+    output: pathlib.Path = dvc.outs(utils.nwd / ".")
+
+    system: LammpsSystem = zn.params(LammpsSystem())
+    extras: dict = zn.params({})
+
+    def post_init(self):
+        self.input_file = self.output / "input.lmp"
+        self.log_file = self.output / "lammps.log"
+
+    @staticmethod
+    def get_template(filename: str = "lammps.jinja2") -> str:
+        """Load a jinja2 template for lammps that ships with znlib
+
+        Parameters
+        ----------
+        filename: str, default = "lammps.jinja2"
+            The jinja2 template to be adapted for the use with LammpsSimulator
+        """
+        with importlib.resources.path(data, "lammps.jinja2") as file:
+            shutil.copy(file, filename)
+        return filename
+
+    def render_template(self):
+        """Write the lammps input file"""
+        with open(self.template, "r", encoding="utf-8") as file:
+            template = jinja2.Template(file.read())
+
+        input_script = template.render(
+            system=self.system, log_file=self.log_file, **self.extras
+        )
+        self.input_file.write_text(input_script)
+
+    def run(self):
+        """Run the lammps simulations"""
+        self.render_template()
+        subprocess.check_call([self.lammps_binary, "-in", self.input_file], shell=True)

znlib/data/lammps.jinja2

@@ -0,0 +1,9 @@
+log             {{ log_file }} append
+
+#############################
+### Configure system type ###
+#############################
+
+units           {{ system.units }}
+boundary        {{ system.boundary }}
+atom_style      {{ system.atom_style }}