Merge pull request #585 from jmmshn/patch_prv

Allow prv_dir to be used more in defect wf

src/atomate2/common/flows/defect.py

@@ -5,7 +5,6 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from jobflow import Flow, Job, Maker, OutputReference
@@ -21,9 +20,12 @@
 )
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     import numpy.typing as npt
     from pymatgen.analysis.defects.core import Defect
     from pymatgen.core.structure import Structure
+    from pymatgen.entries.computed_entries import ComputedStructureEntry
 
 logger = logging.getLogger(__name__)
 
@@ -295,6 +297,7 @@ def make(
                 uc_structure=defect.structure,
                 relax_maker=self.bulk_relax_maker,
                 sc_mat=supercell_matrix,
+                get_planar_locpot=self.get_planar_locpot,
             )
             sc_mat = get_sc_job.output["sc_mat"]
             lattice = get_sc_job.output["sc_struct"].lattice
@@ -305,8 +308,8 @@ def make(
             get_sc_job = get_supercell_from_prv_calc(
                 uc_structure=defect.structure,
                 prv_calc_dir=bulk_supercell_dir,
+                sc_entry_and_locpot_from_prv=self.sc_entry_and_locpot_from_prv,
                 sc_mat_ref=supercell_matrix,
-                structure_from_prv=self.structure_from_prv,
             )
             sc_mat = get_sc_job.output["sc_mat"]
             lattice = get_sc_job.output["lattice"]
@@ -329,12 +332,6 @@ def make(
         jobs.extend([get_sc_job, spawn_output])
 
         if self.collect_defect_entry_data:
-            if isinstance(bulk_supercell_dir, (str, Path)):
-                raise NotImplementedError(
-                    "DefectEntery creation only works when you are explicitly "
-                    "calculating the bulk supercell. This is because the bulk "
-                    "SC energy parsing from previous calculations is not implemented."
-                )
             collection_job = get_defect_entry(
                 charge_state_summary=spawn_output.output,
                 bulk_summary=get_sc_job.output,
@@ -348,8 +345,10 @@ def make(
         )
 
     @abstractmethod
-    def structure_from_prv(self, previous_dir: str) -> Structure:
-        """Copy the output structure from previous directory.
+    def sc_entry_and_locpot_from_prv(
+        self, previous_dir: str
+    ) -> tuple[ComputedStructureEntry, dict]:
+        """Copy the output ComputedStructureEntry and Locpot from previous directory.
 
         Parameters
         ----------
@@ -358,7 +357,25 @@ def structure_from_prv(self, previous_dir: str) -> Structure:
 
         Returns
         -------
-        structure: Structure
+        entry: ComputedStructureEntry
+        """
+
+    @abstractmethod
+    def get_planar_locpot(self, task_doc) -> dict:
+        """Get the Planar Locpot from the TaskDoc.
+
+        This is needed just in case the planar average locpot is stored in different
+        part of the TaskDoc for different codes.
+
+        Parameters
+        ----------
+        task_doc: TaskDoc
+            The task document.
+
+        Returns
+        -------
+        planar_locpot: dict
+            The planar average locpot.
         """
 
     @abstractmethod

src/atomate2/common/jobs/defect.py

@@ -17,6 +17,7 @@
 from pymatgen.entries.computed_entries import ComputedStructureEntry
 
 from atomate2.common.schemas.defects import CCDDocument
+from atomate2.utils.path import strip_hostname
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -185,9 +186,9 @@ def get_ccd_documents(
 @job
 def get_supercell_from_prv_calc(
     uc_structure: Structure,
-    prv_calc_dir: str | Path | None = None,
+    prv_calc_dir: str | Path,
+    sc_entry_and_locpot_from_prv: Callable,
     sc_mat_ref: NDArray | None = None,
-    structure_from_prv: Callable | None = None,
 ) -> dict:
     """Get the supercell from the previous calculation.
 
@@ -201,15 +202,18 @@ def get_supercell_from_prv_calc(
         The directory of the previous calculation.
     sc_mat : NDArray
         The supercell matrix. If not None, use this to validate the extracted supercell.
-    structure_from_prv : Callable
-        Function to get the supercell structure from the previous calculation.
+    sc_entry_and_locpot_from_prv : Callable
+        Function to get the supercell ComputedStructureEntry and Locpot from the
+        previous calculation.
 
     Returns
     -------
     Response:
         Output containing the supercell transformation and the dir_name
     """
-    sc_structure = structure_from_prv(prv_calc_dir)
+    prv_calc_dir = strip_hostname(prv_calc_dir)
+    sc_entry, plnr_locpot = sc_entry_and_locpot_from_prv(prv_calc_dir)
+    sc_structure = sc_entry.structure
     sc_mat_prv, _ = get_matched_structure_mapping(
         uc_struct=uc_structure, sc_struct=sc_structure
     )
@@ -225,14 +229,23 @@ def get_supercell_from_prv_calc(
                 "The supercell matrix extracted from the previous calculation "
                 "does not match the the desired supercell shape."
             )
-    return {"sc_mat": sc_mat_prv, "lattice": Lattice(sc_structure.lattice.matrix)}
+    return {
+        "sc_entry": sc_entry,
+        "sc_struct": sc_structure,
+        "sc_mat": sc_mat_prv,
+        "dir_name": prv_calc_dir,
+        "lattice": Lattice(sc_structure.lattice.matrix),
+        "uuid": None,
+        "locpot_plnr": plnr_locpot,
+    }
 
 
 @job(name="bulk supercell")
 def bulk_supercell_calculation(
     uc_structure: Structure,
     relax_maker: RelaxMaker,
     sc_mat: NDArray | None = None,
+    get_planar_locpot: Callable | None = None,
 ) -> Response:
     """Bulk Supercell calculation.
 
@@ -246,12 +259,19 @@ def bulk_supercell_calculation(
         The relax maker to use.
     sc_mat : NDArray | None
         The supercell matrix used to construct the simulation cell.
+    get_plnr_locpot : Callable | None
+        A function to get the Locpot from the output of the task document.
 
     Returns
     -------
     Response:
         Output a dictionary containing the bulk supercell calculation summary.
     """
+    if get_planar_locpot is None:
+
+        def get_planar_locpot(tdoc):
+            return tdoc.calcs_reversed[0].output.locpot
+
     logger.info("Running bulk supercell calculation. Running...")
     sc_mat = get_sc_fromstruct(uc_structure) if sc_mat is None else sc_mat
     sc_mat = np.array(sc_mat)
@@ -270,7 +290,7 @@ def bulk_supercell_calculation(
         "sc_mat": sc_mat.tolist(),
         "dir_name": relax_output.dir_name,
         "uuid": relax_job.uuid,
-        "locpot_plnr": relax_output.calcs_reversed[0].output.locpot,
+        "locpot_plnr": get_planar_locpot(relax_output),
     }
     flow = Flow([relax_job], output=summary_d)
     return Response(replace=flow)
@@ -409,11 +429,11 @@ def check_charge_state(charge_state: int, task_structure: Structure) -> Response
 @job
 def get_defect_entry(charge_state_summary: dict, bulk_summary: dict) -> list[dict]:
     """Get a defect entry from a defect calculation and a bulk calculation."""
-    bulk_sc_entry = bulk_summary["sc_entry"]
-    bulk_struct_entry = ComputedStructureEntry(
-        structure=bulk_summary["sc_struct"],
-        energy=bulk_sc_entry.energy,
-    )
+    bulk_struct_entry = bulk_summary["sc_entry"]
+    # bulk_struct_entry = ComputedStructureEntry(
+    #     structure=bulk_summary["sc_struct"],
+    #     energy=bulk_sc_entry.energy,
+    # )
     bulk_dir_name = bulk_summary["dir_name"]
     bulk_locpot = bulk_summary["locpot_plnr"]
     defect_ent_res = []

src/atomate2/vasp/flows/defect.py

@@ -4,16 +4,13 @@
 
 import logging
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import TYPE_CHECKING
 
+from emmet.core.tasks import TaskDoc
 from jobflow import Flow, Maker, OutputReference
 from jobflow.core.maker import recursive_call
-from pymatgen.io.vasp.outputs import Vasprun
 
-from atomate2.common.files import get_zfile
 from atomate2.common.flows import defect as defect_flows
-from atomate2.utils.file_client import FileClient
 from atomate2.vasp.flows.core import DoubleRelaxMaker
 from atomate2.vasp.jobs.core import RelaxMaker, StaticMaker
 from atomate2.vasp.jobs.defect import calculate_finite_diff
@@ -26,6 +23,7 @@
 
 if TYPE_CHECKING:
     from pymatgen.core.structure import Structure
+    from pymatgen.entries.computed_entries import ComputedStructureEntry
 
     from atomate2.common.schemas.defects import CCDDocument
     from atomate2.vasp.jobs.base import BaseVaspMaker
@@ -169,7 +167,9 @@ class FormationEnergyMaker(defect_flows.FormationEnergyMaker):
     bulk_relax_maker: BaseVaspMaker | None = None
     name: str = "formation energy"
 
-    def structure_from_prv(self, previous_dir: str) -> Structure:
+    def sc_entry_and_locpot_from_prv(
+        self, previous_dir: str
+    ) -> tuple[ComputedStructureEntry, dict]:
         """Copy the output structure from previous directory.
 
         Read the vasprun.xml file from the previous directory
@@ -182,15 +182,14 @@ def structure_from_prv(self, previous_dir: str) -> Structure:
 
         Returns
         -------
-        structure: Structure
+        ComputedStructureEntry
         """
-        fc = FileClient()
-        # strip off the `hostname:` prefix
-        previous_dir = previous_dir.split(":")[-1]
-        files = fc.listdir(previous_dir)
-        vasprun_file = Path(previous_dir) / get_zfile(files, "vasprun.xml")
-        vasprun = Vasprun(vasprun_file)
-        return vasprun.final_structure
+        task_doc = TaskDoc.from_directory(previous_dir)
+        return task_doc.structure_entry, task_doc.calcs_reversed[0].output.locpot
+
+    def get_planar_locpot(self, task_doc: TaskDoc) -> dict:
+        """Get the planar-averaged electrostatic potential."""
+        return task_doc.calcs_reversed[0].output.locpot
 
     def validate_maker(self) -> None:
         """Check some key settings in the relax maker.

tests/vasp/flows/test_defect.py

@@ -153,7 +153,6 @@ def test_formation_energy_maker(mock_vasp, clean_dir, test_dir, monkeypatch):
         )
     )
 
-    # rmaker = RelaxMaker(input_set_generator=ChargeStateRelaxSetGenerator())
     maker = FormationEnergyMaker(
         relax_radius="auto",
         perturb=0.1,
@@ -180,3 +179,16 @@ def _check_plnr_locpot(name):
 
     for k in ref_paths:
         _check_plnr_locpot(k)
+
+    # make sure the the you can restart the calculation from prv
+    prv_dir = test_dir / "vasp/GaN_Mg_defect/bulk_relax/outputs"
+    flow2 = maker.make(
+        defects[0],
+        bulk_supercell_dir=prv_dir,
+        defect_index=0,
+    )
+    _ = run_locally(
+        flow2,
+        create_folders=True,
+        ensure_success=True,
+    )