diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index 1ff536ee19..f6d8bce15a 100644 --- a/src/atomate2/common/flows/defect.py +++ b/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 diff --git a/src/atomate2/common/jobs/defect.py b/src/atomate2/common/jobs/defect.py index 610e7ec83a..d1d5d325c0 100644 --- a/src/atomate2/common/jobs/defect.py +++ b/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,7 +229,15 @@ 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") @@ -233,6 +245,7 @@ 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 = [] diff --git a/src/atomate2/vasp/flows/defect.py b/src/atomate2/vasp/flows/defect.py index e60918dc2e..6e5577e5c5 100644 --- a/src/atomate2/vasp/flows/defect.py +++ b/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. diff --git a/tests/vasp/flows/test_defect.py b/tests/vasp/flows/test_defect.py index 9766eac600..4da2fbbad2 100644 --- a/tests/vasp/flows/test_defect.py +++ b/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, + )