Addition of JDFTx code

jdftxcov.xml

@@ -0,0 +1,101 @@
+<?xml version="1.0" ?>
+<coverage version="7.6.1" timestamp="1727213408713" lines-valid="68" lines-covered="58" line-rate="0.8529" branches-valid="12" branches-covered="10" branch-rate="0.8333" complexity="0">
+	<!-- Generated by coverage.py: https://coverage.readthedocs.io/en/7.6.1 -->
+	<!-- Based on https://raw.githubusercontent.com/cobertura/web/master/htdocs/xml/coverage-04.dtd -->
+	<sources>
+		<source>/Users/sophi/envs/jobflow_py310/src/atomate2/src/atomate2/jdftx/jobs</source>
+	</sources>
+	<packages>
+		<package name="." line-rate="0.8529" branch-rate="0.8333" complexity="0">
+			<classes>
+				<class name="__init__.py" filename="__init__.py" complexity="0" line-rate="1" branch-rate="1">
+					<methods/>
+					<lines/>
+				</class>
+				<class name="base.py" filename="base.py" complexity="0" line-rate="0.9744" branch-rate="0.75">
+					<methods/>
+					<lines>
+						<line number="3" hits="1"/>
+						<line number="5" hits="1"/>
+						<line number="6" hits="1"/>
+						<line number="7" hits="1"/>
+						<line number="8" hits="1"/>
+						<line number="10" hits="1"/>
+						<line number="11" hits="1"/>
+						<line number="12" hits="1"/>
+						<line number="14" hits="1"/>
+						<line number="16" hits="1" branch="true" condition-coverage="50% (1/2)" missing-branches="17"/>
+						<line number="17" hits="0"/>
+						<line number="18" hits="1"/>
+						<line number="19" hits="1"/>
+						<line number="24" hits="1"/>
+						<line number="25" hits="1"/>
+						<line number="26" hits="1"/>
+						<line number="28" hits="1"/>
+						<line number="37" hits="1"/>
+						<line number="44" hits="1"/>
+						<line number="55" hits="1"/>
+						<line number="70" hits="1"/>
+						<line number="73" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="74" hits="1"/>
+						<line number="93" hits="1"/>
+						<line number="94" hits="1"/>
+						<line number="97" hits="1"/>
+						<line number="98" hits="1"/>
+						<line number="99" hits="1"/>
+						<line number="101" hits="1"/>
+						<line number="102" hits="1"/>
+						<line number="116" hits="1"/>
+						<line number="119" hits="1"/>
+						<line number="121" hits="1"/>
+						<line number="123" hits="1"/>
+						<line number="124" hits="1"/>
+						<line number="126" hits="1"/>
+						<line number="130" hits="1"/>
+						<line number="137" hits="1"/>
+						<line number="139" hits="1"/>
+					</lines>
+				</class>
+				<class name="core.py" filename="core.py" complexity="0" line-rate="0.9524" branch-rate="0.875">
+					<methods/>
+					<lines>
+						<line number="3" hits="1"/>
+						<line number="5" hits="1"/>
+						<line number="6" hits="1"/>
+						<line number="7" hits="1"/>
+						<line number="9" hits="1"/>
+						<line number="10" hits="1"/>
+						<line number="13" hits="1" branch="true" condition-coverage="50% (1/2)" missing-branches="14"/>
+						<line number="14" hits="0"/>
+						<line number="17" hits="1"/>
+						<line number="20" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="21" hits="1"/>
+						<line number="24" hits="1"/>
+						<line number="25" hits="1"/>
+						<line number="27" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="28" hits="1"/>
+						<line number="31" hits="1"/>
+						<line number="32" hits="1"/>
+						<line number="34" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="35" hits="1"/>
+						<line number="38" hits="1"/>
+						<line number="39" hits="1"/>
+					</lines>
+				</class>
+				<class name="test_run.py" filename="test_run.py" complexity="0" line-rate="0" branch-rate="1">
+					<methods/>
+					<lines>
+						<line number="1" hits="0"/>
+						<line number="2" hits="0"/>
+						<line number="3" hits="0"/>
+						<line number="4" hits="0"/>
+						<line number="6" hits="0"/>
+						<line number="7" hits="0"/>
+						<line number="8" hits="0"/>
+						<line number="9" hits="0"/>
+					</lines>
+				</class>
+			</classes>
+		</package>
+	</packages>
+</coverage>

src/atomate2/jdftx/__init__.py

@@ -0,0 +1 @@
+"""Module for JDFTx workflows."""

src/atomate2/jdftx/files.py

@@ -0,0 +1,38 @@
+"""File operations and default JDFTx file names."""
+
+import logging
+
+# if TYPE_CHECKING:
+from pathlib import Path
+
+from pymatgen.core import Structure
+
+from atomate2.jdftx.sets.base import JdftxInputGenerator
+
+logger = logging.getLogger(__name__)
+
+
+def write_jdftx_input_set(
+    structure: Structure,
+    input_set_generator: JdftxInputGenerator,
+    directory: str | Path = ".",
+    **kwargs,
+) -> None:
+    """
+    Write JDFTx input set.
+
+    Parameters
+    ----------
+    structure : .Structure
+        A structure.
+    input_set_generator : .JdftxInputGenerator
+        A JDFTx input set generator.
+    directory : str or Path
+        The directory to write the input files to.
+    **kwargs
+        Keyword arguments to pass to :obj:`.JdftxInputSet.write_input`.
+    """
+    cis = input_set_generator.get_input_set(structure)
+
+    logger.info("Writing JDFTx input set.")
+    cis.write_input(directory, **kwargs)

src/atomate2/jdftx/jobs/__init__.py

@@ -0,0 +1 @@
+"""Module for JDFTx jobs."""

src/atomate2/jdftx/jobs/adsorption.py

@@ -0,0 +1,42 @@
+"""Core jobs for running JDFTx calculations."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from atomate2.jdftx.jobs.base import BaseJdftxMaker
+from atomate2.jdftx.sets.core import IonicMinSetGenerator
+
+if TYPE_CHECKING:
+    from atomate2.jdftx.sets.core import JdftxInputGenerator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SurfaceMinMaker(BaseJdftxMaker):
+    """Maker to create surface relaxation job."""
+
+    name: str = "surface_ionic_min"
+    input_set_generator: JdftxInputGenerator = field(
+        default_factory=lambda: IonicMinSetGenerator(
+            coulomb_truncation=True,
+            auto_kpoint_density=1000,
+            calc_type="surface",
+        )
+    )
+
+
+@dataclass
+class MolMinMaker(BaseJdftxMaker):
+    """Maker to create molecule relaxation job."""
+
+    name: str = "surface_ionic_min"
+    input_set_generator: JdftxInputGenerator = field(
+        default_factory=IonicMinSetGenerator(
+            coulomb_truncation=True,
+            calc_type="molecule",
+        )
+    )

src/atomate2/jdftx/jobs/base.py

@@ -0,0 +1,140 @@
+"""Definition of base JDFTx job maker."""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import Maker, Response, job
+
+from atomate2.jdftx.sets.base import JdftxInputGenerator
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from pathlib import Path
+
+    from pymatgen.core import Structure
+from pymatgen.core.trajectory import Trajectory
+from pymatgen.electronic_structure.bandstructure import (
+    BandStructure,
+    BandStructureSymmLine,
+)
+
+from atomate2.jdftx.files import write_jdftx_input_set
+from atomate2.jdftx.run import run_jdftx, should_stop_children
+from atomate2.jdftx.schemas.task import TaskDoc
+
+logger = logging.getLogger(__name__)
+
+_DATA_OBJECTS = [  # TODO update relevant list for JDFTx
+    BandStructure,
+    BandStructureSymmLine,
+    Trajectory,
+    "force_constants",
+    "normalmode_eigenvecs",
+    "bandstructure",  # FIX: BandStructure is not currently MSONable
+]
+
+_INPUT_FILES = [
+    "init.in",
+    "init.lattice",
+    "init.ionpos",
+]
+
+# Output files.
+_OUTPUT_FILES = [  # TODO finish this list
+    "output.out",
+    "Ecomponents",
+    "wfns",
+    "bandProjections",
+    "boundCharge",
+    "lattice",
+    "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
+class BaseJdftxMaker(Maker):
+    """
+    Base JDFTx job maker.
+
+    Parameters
+    ----------
+    name : str
+        The job name.
+    input_set_generator : .JdftxInputGenerator
+        A generator used to make the input set.
+    write_input_set_kwargs : dict
+        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:
+        """Run a JDFTx calculation.
+
+        Parameters
+        ----------
+        structure : Structure
+            A pymatgen structure object.
+
+        Returns
+        -------
+            Response: A response object containing the output, detours and stop
+                commands of the JDFTx run.
+        """
+        # write jdftx input files
+        write_jdftx_input_set(
+            structure, self.input_set_generator, **self.write_input_set_kwargs
+        )
+        logger.info("Wrote JDFTx input files.")
+        # run jdftx
+        run_jdftx(**self.run_jdftx_kwargs)
+
+        current_dir = os.getcwd()
+        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={},
+            output=task_doc,
+        )
+
+
+def get_jdftx_task_document(path: Path | str, **kwargs) -> TaskDoc:
+    """Get JDFTx Task Document using atomate2 settings."""
+    return TaskDoc.from_directory(path, **kwargs)

src/atomate2/jdftx/jobs/core.py

@@ -0,0 +1,50 @@
+"""Core jobs for running JDFTx calculations."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from atomate2.jdftx.jobs.base import BaseJdftxMaker
+from atomate2.jdftx.sets.core import (
+    IonicMinSetGenerator,
+    LatticeMinSetGenerator,
+    SinglePointSetGenerator,
+)
+
+if TYPE_CHECKING:
+    from atomate2.jdftx.sets.base import JdftxInputGenerator
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SinglePointMaker(BaseJdftxMaker):
+    """Maker to create JDFTx ionic optimization job."""
+
+    name: str = "single_point"
+    input_set_generator: JdftxInputGenerator = field(
+        default_factory=SinglePointSetGenerator
+    )
+
+
+@dataclass
+class IonicMinMaker(BaseJdftxMaker):
+    """Maker to create JDFTx ionic optimization job."""
+
+    name: str = "ionic_min"
+    input_set_generator: JdftxInputGenerator = field(
+        default_factory=IonicMinSetGenerator
+    )
+
+
+@dataclass
+class LatticeMinMaker(BaseJdftxMaker):
+    """Maker to create JDFTx lattice optimization job."""
+
+    name: str = "lattice_min"
+    input_set_generator: JdftxInputGenerator = field(
+        default_factory=LatticeMinSetGenerator
+    )