Changes to get to functional run on Perlmutter

src/atomate2/jdftx/io/JDFTXOutfile.py

@@ -5,6 +5,8 @@
 from atomate2.jdftx.io.JDFTXOutfileSlice import JDFTXOutfileSlice
 from dataclasses import dataclass
 from typing import List, Optional
+from monty.io import zopen
+from pathlib import Path
 
 
 class ClassPrintFormatter():
@@ -26,8 +28,9 @@ def check_file_exists(func):
 
     @wraps(func)
     def wrapper(filename):
-        if not os.path.isfile(filename):
-            raise OSError("'" + filename + "' file doesn't exist!")
+        filename = Path(filename)
+        if not filename.is_file():
+            raise OSError(f"'{filename}' file doesn't exist!")
         return func(filename)
 
     return wrapper
@@ -48,7 +51,7 @@ def read_file(file_name: str) -> list[str]:
     text: list[str]
         list of strings from file
     '''
-    with open(file_name, 'r') as f:
+    with zopen(file_name, 'r') as f:
         text = f.readlines()
     return text
 

src/atomate2/jdftx/io/JDFTXOutfileSlice.py

@@ -1,7 +1,6 @@
 import os
 from functools import wraps
 import math
-from ase import Atom, Atoms
 from atomate2.jdftx.io.JMinSettings import JMinSettings, JMinSettingsElectronic, JMinSettingsFluid, JMinSettingsIonic, JMinSettingsLattice
 import numpy as np
 from dataclasses import dataclass, field

src/atomate2/jdftx/io/JOutStructures.py

@@ -1,5 +1,6 @@
 from typing import Any
 from atomate2.jdftx.io.JOutStructure import JOutStructure
+
 from dataclasses import dataclass
 
 

src/atomate2/jdftx/jobs/base.py

@@ -5,7 +5,6 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
-
 from atomate2.jdftx.schemas.task import TaskDoc
 from jobflow import Maker, Response, job
 from pymatgen.core.trajectory import Trajectory
@@ -17,9 +16,7 @@
 
 from atomate2.jdftx.sets.base import JdftxInputGenerator
 from atomate2.jdftx.files import write_jdftx_input_set
-
-
-from atomate2.jdftx.run import run_jdftx
+from atomate2.jdftx.run import run_jdftx, should_stop_children
 
 #if TYPE_CHECKING:
 from pymatgen.core import Structure
@@ -35,14 +32,14 @@
 ]
 
 _INPUT_FILES = [
-    "inputs.in",
-    "inputs.lattice",
-    "inputs.ionpos",
+    "init.in",
+    "init.lattice",
+    "init.ionpos",
 ]
 
 # Output files. Partially from https://www.vasp.at/wiki/index.php/Category:Output_files
 _OUTPUT_FILES = [ # TODO finish this list
-    "out.log",
+    "output.out",
     "Ecomponents",
     "wfns",
     "bandProjections",
@@ -51,7 +48,22 @@
     "ionpos",
 ]
 
+def jdftx_job(method: Callable) -> job:
+    """
+    Decorate the ``make`` method of JDFTx job makers.
 
+    Parameters
+    ----------
+    method : callable
+        A BaseJdftxMaker.make method. This should not be specified directly and is
+        implied by the decorator.
+
+    Returns
+    -------
+    callable
+        A decorated version of the make function that will generate JDFTx jobs.
+    """
+    return job(method, data=_DATA_OBJECTS, output_schema=TaskDoc)
 
 
 @dataclass
@@ -69,16 +81,18 @@ class BaseJdftxMaker(Maker):
         Keyword arguments that will get passed to :obj:`.write_jdftx_input_set`.
     run_jdftx_kwargs : dict
         Keyword arguments that will get passed to :obj:`.run_jdftx`.
-
+    task_document_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`.
 
     """
 
     name: str = "base JDFTx job"
     input_set_generator: JdftxInputGenerator = field(default_factory=JdftxInputGenerator)
     write_input_set_kwargs: dict = field(default_factory=dict)
     run_jdftx_kwargs: dict = field(default_factory=dict)
+    task_document_kwargs: dict = field(default_factory=dict)
 
-
+    @jdftx_job
     def make(
         self, structure: Structure
     ) -> Response:
@@ -105,6 +119,10 @@ def make(
         current_dir = Path.cwd()
         files = [str(f) for f in current_dir.glob('*') if f.is_file()]
 
+        task_doc = get_jdftx_task_document(current_dir, **self.task_document_kwargs)
+
+        stop_children = should_stop_children(task_doc)
+
         return Response(
             stop_children=stop_children,
             stored_data={"custodian": task_doc.custodian},
@@ -115,4 +133,5 @@ def make(
 def get_jdftx_task_document(path: Path | str, **kwargs) -> TaskDoc:
     """Get JDFTx Task Document using atomate2 settings."""
 
-    return TaskDoc.from_directory(path, **kwargs)
+    return TaskDoc.from_directory(path, **kwargs)
+

src/atomate2/jdftx/run.py

@@ -5,6 +5,7 @@
 
 from jobflow.utils import ValueEnum
 from atomate2.jdftx.jobs.jobs import JDFTxJob
+from atomate2.jdftx.schemas.task import TaskDoc, JDFTxStatus
 
 class JobType(ValueEnum):
     """
@@ -24,8 +25,20 @@ def run_jdftx(
     jdftx_job_kwargs = jdftx_job_kwargs or {}
 
     if job_type == JobType.NORMAL:
-        job = JDFTxJob(jdftx_cmd, **jdftx_job_kwargs, input_file="input-tutorial.in")
+        job = JDFTxJob(jdftx_cmd, **jdftx_job_kwargs)
 
     job.run()
 
-#need to call job = run_jdftx() to run calc
+#need to call job = run_jdftx() to run calc
+
+def should_stop_children(
+    task_document: TaskDoc,
+) -> bool:
+    """
+    Parse JDFTx TaskDoc and decide whether to stop child processes.
+    If JDFTx failed, stop child processes.
+    """
+    if task_document.state == JDFTxStatus.SUCCESS:
+        return False
+    else:
+        return True

src/atomate2/jdftx/schemas/calculation.py

@@ -20,14 +20,15 @@
 #     task_type,
 #     calc_type,
 # )
-from atomate2.jdftx.schemas.task import JDFTxStatus
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.qchem.inputs import QCInput
 from pymatgen.io.qchem.outputs import QCOutput
 
 from atomate2.jdftx.io.JDFTXInfile import JDFTXInfile, JDFTXStructure
 from atomate2.jdftx.io.JDFTXOutfile import JDFTXOutfile
+from emmet.core.utils import ValueEnum
+
 
 functional_synonyms = {
     "b97mv": "b97m-v",
@@ -51,6 +52,14 @@
 # as QChem data objects
 
 
+class JDFTxStatus(ValueEnum):
+    """
+    JDFTx Calculation State
+    """
+
+    SUCCESS = "successful"
+    FAILED = "unsuccessful"
+
 class CalculationInput(BaseModel):
     """
     Document defining JDFTx calculation inputs.

src/atomate2/jdftx/schemas/task.py

@@ -10,12 +10,12 @@
 from pymatgen.core import Structure
 from custodian.qchem.jobs import QCJob
 from emmet.core.qchem.calc_types import CalcType, LevelOfTheory, TaskType
-from emmet.core.utils import ValueEnum
 from atomate2.jdftx.schemas.calculation import Calculation, CalculationInput, CalculationOutput
 from emmet.core.structure import StructureMetadata
 from monty.serialization import loadfn
 from pydantic import BaseModel, Field
 
+from atomate2.jdftx.schemas.calculation import JDFTxStatus
 from atomate2.utils.datetime import datetime_str
 
 __author__ = (
@@ -26,14 +26,6 @@
 _T = TypeVar("_T", bound="TaskDoc")
 # _DERIVATIVE_FILES = ("GRAD", "HESS")
 
-class JDFTxStatus(ValueEnum):
-    """
-    JDFTx Calculation State
-    """
-
-    SUCCESS = "successful"
-    FAILED = "unsuccessful"
-
 class OutputDoc(BaseModel):
     initial_structure: Structure = Field(None, description="Input Structure object")
     optimized_structure: Optional[Structure] = Field(
@@ -235,7 +227,7 @@ def from_directory(
         calc_doc = Calculation.from_files(
             dir_name=dir_name,
             jdftxinput_file="inputs.in",
-            jdftxoutput_file="out.log"
+            jdftxoutput_file="output.out"
             )
         # task_files = _find_qchem_files(dir_name)
 

src/atomate2/jdftx/sets/BaseJdftxSet.yaml

@@ -49,6 +49,7 @@ symmetries: none
 ion-species: GBRV_v1.5/$ID_pbe_v1.uspp
 
 ### Output Files ###
+dump-name: jdftx.$VAR
 dump:
   - freq: End
     var: Dtot

src/atomate2/jdftx/sets/base.py

@@ -10,6 +10,7 @@
 
 from monty.serialization import loadfn
 from pymatgen.io.core import InputGenerator, InputSet
+from pymatgen.util.typing import PathLike
 
 from atomate2.jdftx.io.JDFTXInfile import (  # TODO update this to the pymatgen module
     JDFTXInfile,
@@ -43,6 +44,7 @@ def __init__(
     def write_input(
         self,
         directory: str | Path,
+        infile: PathLike = "inputs.in", #TODO I don't think this should be optional
         make_dir: bool = True,
         overwrite: bool = True,
     ) -> None:
@@ -57,15 +59,16 @@ def write_input(
         overwrite
             Whether to overwrite an input file if it already exists.
         """
-        infile = "inputs.in"
         directory = Path(directory)
         if make_dir:
             os.makedirs(directory, exist_ok=True)
 
         if not overwrite and (directory / infile).exists():
             raise FileExistsError(f"{directory / infile} already exists.")
 
-        self.jdftxinput.write_file(filename=(directory / infile))
+        jdftxinput = condense_jdftxinputs(self.jdftxinput, self.jdftxstructure)
+
+        jdftxinput.write_file(filename=(directory / infile))
 
     @staticmethod
     def from_directory(
@@ -128,3 +131,32 @@ def get_input_set(
         jdftxinput = JDFTXInfile.from_dict(jdftxinputs)
 
         return JdftxInputSet(jdftxinput=jdftxinput, jdftxstructure=jdftx_structure)
+
+def condense_jdftxinputs(
+        jdftxinput:JDFTXInfile, 
+        jdftxstructure:JDFTXStructure
+        ) -> JDFTXInfile:
+    """
+    Function to combine a JDFTXInputs class with calculation
+    settings and a JDFTxStructure that defines the structure
+    into one JDFTXInputs instance.
+
+    Parameters
+        ----------
+        jdftxinput: JDFTXInfile
+            A JDFTXInfile object with calculation settings.
+
+        jdftxstructure: JDFTXStructure
+            A JDFTXStructure object that defines the structure.
+
+        Returns
+        -------
+        JDFTXInfile
+            A JDFTXInfile that includes the calculation
+            parameters and input structure.
+    """
+    condensed_inputs = (
+        jdftxinput + 
+        JDFTXInfile.from_str(jdftxstructure.get_str())
+        )
+    return condensed_inputs