From e91b507bf9e135e23bf07aaa539f4f297309f624 Mon Sep 17 00:00:00 2001
From: Fabian <46721498+PythonFZ@users.noreply.github.com>
Date: Sun, 30 Oct 2022 13:53:40 +0100
Subject: [PATCH] initial idea of a Lammps Node

---
 poetry.lock                 |  2 +-
 pyproject.toml              |  5 ++-
 tests/test_atomistic.py     | 42 +++++++++++++++++++
 znlib/atomistic/__init__.py |  3 +-
 znlib/atomistic/lammps.py   | 84 +++++++++++++++++++++++++++++++++++++
 znlib/data/__init__.py      |  0
 znlib/data/lammps.jinja2    |  9 ++++
 7 files changed, 142 insertions(+), 3 deletions(-)
 create mode 100644 znlib/atomistic/lammps.py
 create mode 100644 znlib/data/__init__.py
 create mode 100644 znlib/data/lammps.jinja2

diff --git a/poetry.lock b/poetry.lock
index b4ef047..885659d 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2536,7 +2536,7 @@ zntrack = ["zntrack", "matplotlib"]
 [metadata]
 lock-version = "1.1"
 python-versions = ">=3.8,<4.0.0"
-content-hash = "1b048bca5a02d530589e3e61e21657a978f34bb842794b6eba18f3694f04ede7"
+content-hash = "69e6290975e06dbf21c0375943ce2e24e0e54f04bbcf996f10acf8b6e97a25f9"
 
 [metadata.files]
 aiohttp = [
diff --git a/pyproject.toml b/pyproject.toml
index 03c9609..4365685 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,6 +5,7 @@ description = ""
 authors = ["zincwarecode <zincwarecode@gmail.com>"]
 license = "Apache-2.0"
 readme = "README.md"
+include = ["znlib/data"]
 
 [tool.poetry.dependencies]
 python = ">=3.8,<4.0.0"
@@ -12,11 +13,12 @@ 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"
diff --git a/tests/test_atomistic.py b/tests/test_atomistic.py
index b3c5288..0330bee 100644
--- a/tests/test_atomistic.py
+++ b/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
diff --git a/znlib/atomistic/__init__.py b/znlib/atomistic/__init__.py
index 62169ce..9662425 100644
--- a/znlib/atomistic/__init__.py
+++ b/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"]
diff --git a/znlib/atomistic/lammps.py b/znlib/atomistic/lammps.py
new file mode 100644
index 0000000..8545d70
--- /dev/null
+++ b/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)
diff --git a/znlib/data/__init__.py b/znlib/data/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/znlib/data/lammps.jinja2 b/znlib/data/lammps.jinja2
new file mode 100644
index 0000000..e24d88f
--- /dev/null
+++ b/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 }}