Merge pull request #8 from cote3804/jdft_sets

Initial Jobs scripts via custodian

src/atomate2/jdftx/jobs/base.py

@@ -21,9 +21,9 @@
 
 from atomate2 import SETTINGS #TODO we can add JDFTx workflow default settings this way
 from atomate2.common.files import gzip_output_folder
+from atomate2.jdftx.sets.base import JdftxInputGenerator
 from atomate2.vasp.files import copy_vasp_outputs, write_vasp_input_set
 from atomate2.vasp.run import run_vasp, should_stop_children
-from atomate2.jdftx.sets.base import JdftxInputGenerator
 
 if TYPE_CHECKING:
     from pymatgen.core import Structure
@@ -249,4 +249,4 @@ def get_jdftx_task_document(path: Path | str, **kwargs) -> TaskDoc:
 
     # kwargs.setdefault("store_volumetric_data", SETTINGS.VASP_STORE_VOLUMETRIC_DATA)
 
-    return TaskDoc.from_directory(path, **kwargs)
+    return TaskDoc.from_directory(path, **kwargs)

src/atomate2/jdftx/jobs/core.py

@@ -64,4 +64,4 @@ class RelaxMaker(BaseVaspMaker):
     """
 
     name: str = "static"
-    input_set_generator: VaspInputGenerator = field(default_factory=StaticSetGenerator)
+    input_set_generator: VaspInputGenerator = field(default_factory=StaticSetGenerator)

src/atomate2/jdftx/jobs/jobs.py

@@ -0,0 +1,100 @@
+"""This module implements basic kinds of jobs for JDFTx runs."""
+
+import logging
+import os
+import subprocess
+
+from custodian.custodian import Job
+
+logger = logging.getLogger(__name__)
+
+
+class JDFTxJob(Job):
+    """
+    A basic JDFTx job. Runs whatever is in the working directory.
+    """
+
+    # If testing, use something like:
+    # job = JDFTxJob()
+    # job.run()  # assumes input files already written to directory
+
+    # Used Cp2kJob developed by Nick Winner as a template.
+
+    def __init__(
+        self,
+        jdftx_cmd,
+        input_file="jdftx.in",
+        output_file="jdftx.out",
+        stderr_file="std_err.txt",
+    ) -> None:
+        """
+        This constructor is necessarily complex due to the need for
+        flexibility. For standard kinds of runs, it's often better to use one
+        of the static constructors. The defaults are usually fine too.
+
+        Args:
+            jdftx_cmd (str): Command to run JDFTx as a string.
+            input_file (str): Name of the file to use as input to JDFTx
+                executable. Defaults to "input.in"
+            output_file (str): Name of file to direct standard out to.
+                Defaults to "jdftx.out".
+            stderr_file (str): Name of file to direct standard error to.
+                Defaults to "std_err.txt".
+        """
+        self.jdftx_cmd = jdftx_cmd
+        self.input_file = input_file
+        self.output_file = output_file
+        self.stderr_file = stderr_file
+
+    def setup(self, directory="./") -> None:
+        """
+        No setup required.
+        """
+        pass
+
+    def run(self, directory="./"):
+        """
+        Perform the actual JDFTx run.
+
+        Returns:
+            (subprocess.Popen) Used for monitoring.
+        """
+        cmd = self.jdftx_cmd + " -i " + self.input_file
+        logger.info(f"Running {cmd}")
+        with (
+            open(os.path.join(directory, self.output_file), "w") as f_std,
+            open(os.path.join(directory, self.stderr_file), "w", buffering=1) as f_err,
+        ):
+            result = subprocess.run([cmd], cwd=directory, stdout=f_std, stderr=f_err, shell=True)
+
+        # Review the return code
+        if result.returncode == 0:
+            logger.info(f"Command executed successfully with return code {result.returncode}.")
+        else:
+            logger.error(f"Command failed with return code {result.returncode}.")
+        # Optionally, you can log or print additional information here
+            with open(os.path.join(directory, self.stderr_file), 'r') as f_err:
+                error_output = f_err.read()
+            logger.error(f"Standard Error Output:\n{error_output}")
+
+        return result
+
+            # use line buffering for stderr
+        #    return subprocess.run([cmd], cwd=directory, stdout=f_std, stderr=f_err, shell=True)
+
+
+    def postprocess(self, directory="./") -> None:
+        """No post-processing required."""
+        pass
+
+    def terminate(self, directory="./") -> None:
+        """Terminate JDFTx."""
+        # This will kill any running process with "jdftx" in the name,
+        # this might have unintended consequences if running multiple jdftx processes
+        # on the same node.
+        for cmd in self.jdftx_cmd:
+            if "jdftx" in cmd:
+                try:
+                    os.system(f"killall {cmd}")
+                except Exception:
+                    pass

src/atomate2/jdftx/jobs/sample-move-later/Dockerfile

@@ -0,0 +1,21 @@
+# install Docker
+# run "docker build . -t jdftx" from directory containing Dockefile
+
+FROM ubuntu:24.04
+
+RUN apt-get -y update && apt-get -y --no-install-recommends install g++ cmake libgsl0-dev libopenmpi-dev openmpi-bin libfftw3-dev libatlas-base-dev liblapack-dev wget unzip ca-certificates make && \
+	cd /root && \
+	wget https://github.com/shankar1729/jdftx/archive/refs/heads/master.zip && unzip master.zip && rm master.zip && \
+	cd jdftx-master && mkdir build && cd build && \
+	cmake ../jdftx && make all && make install && \
+#        make test && \
+        cd /root && rm -rf /root/jdftx-master && \
+        echo 'export PATH="$PATH:/usr/local/share/jdftx/scripts"' >> /root/.bashrc && mkdir /root/research
+
+WORKDIR /root/research
+
+#Use it like this:
+#docker run -it --rm -v $PWD:/root/research jdftx
+
+#Or even better, put the following line at the end of your .bashrc and/or .zshrc so that you can just run 'jdftx' :
+#function jdftx () { docker run -it --rm -v $PWD:/root/research jdftx ; }

src/atomate2/jdftx/jobs/sample-move-later/input-tutorial.in

@@ -0,0 +1,27 @@
+# The input file is a list of commands, one per line
+# The commands may appear in any order; group them to your liking
+# Everything on a line after a # is treated as a comment and ignored
+# Whitespace separates words; extra whitespace is ignored
+# --------------- Water molecule example ----------------
+
+# Set up the unit cell - each column is a bravais lattice vector in bohrs
+# Hence this is a cubic box of side 10 bohr (Note that \ continues lines)
+lattice \
+    10 0  0 \
+    0  10 0 \
+    0  0  10
+
+elec-cutoff 20 100    #Plane-wave kinetic energy cutoff for wavefunctions and charge density in Hartrees
+
+# Specify the pseudopotentials (this defines species O and H):
+ion-species GBRV/h_pbe.uspp
+ion-species GBRV/o_pbe.uspp
+
+# Specify coordinate system and atom positions:
+coords-type cartesian       #the other option is lattice (suitable for solids)
+ion O   0.00 0.00  0.00  0  # The last 0 holds this atom fixed
+ion H   0.00 1.13 +1.45  1  # while the 1 allows this one to move
+ion H   0.00 1.13 -1.45  1  # during ionic minimization
+
+dump-name water.$VAR              #Filename pattern for outputs
+dump End Ecomponents ElecDensity  #Output energy components and electron density at the end

src/atomate2/jdftx/jobs/sample-move-later/test_run.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+
+# Optional for debugging, use via "pip install stackprinter" first:
+import stackprinter
+stackprinter.set_excepthook(style='darkbg2')
+
+# if running from directory containing jobs.py for testing, can also use "from jobs import JDFTxJob"
+from atomate2.jdftx.jobs.jobs import JDFTxJob
+
+# assumes running this script from directory containing already-generated input files
+# if input files are in a different directory, change "$PWD" below
+job = JDFTxJob(jdftx_cmd="docker run -t --rm -v $PWD:/root/research jdftx jdftx", input_file="input-tutorial.in")
+job.run()