From bddecc9e1026f14dc7721cab47200875a7c4fe60 Mon Sep 17 00:00:00 2001 From: "@jmmshn" Date: Fri, 20 Oct 2023 11:22:19 -0700 Subject: [PATCH 1/6] prv_dir --- src/atomate2/common/flows/defect.py | 41 +++++++++++++++++++++-------- src/atomate2/common/jobs/defect.py | 32 ++++++++++++++-------- src/atomate2/vasp/flows/defect.py | 23 +++++++--------- tests/vasp/flows/test_defect.py | 14 +++++++++- 4 files changed, 74 insertions(+), 36 deletions(-) diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index 9dca7ca9fd..a8e277e3b6 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,6 +20,8 @@ ) 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 @@ -305,8 +306,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 +330,12 @@ 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." - ) + # 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 +349,8 @@ 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) -> Structure: + """Copy the output ComputedStructureEntry and Locpot from previous directory. Parameters ---------- @@ -358,7 +359,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 769f553b9e..bf5e5fd459 100644 --- a/src/atomate2/common/jobs/defect.py +++ b/src/atomate2/common/jobs/defect.py @@ -185,9 +185,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 +201,17 @@ 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) + 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 +227,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") @@ -409,11 +419,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..3c4eb2efa5 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,7 @@ 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) -> ComputedStructureEntry: """Copy the output structure from previous directory. Read the vasprun.xml file from the previous directory @@ -182,15 +180,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 a78fe897db..9cc0d4a009 100644 --- a/tests/vasp/flows/test_defect.py +++ b/tests/vasp/flows/test_defect.py @@ -160,7 +160,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, @@ -187,3 +186,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, + ) From c5a103fdb430da0c2272c80fe2c71e8bb24f2b84 Mon Sep 17 00:00:00 2001 From: "@jmmshn" Date: Fri, 20 Oct 2023 11:29:39 -0700 Subject: [PATCH 2/6] prv_dir --- src/atomate2/common/flows/defect.py | 1 + src/atomate2/common/jobs/defect.py | 10 +++++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index a8e277e3b6..0e28a6232a 100644 --- a/src/atomate2/common/flows/defect.py +++ b/src/atomate2/common/flows/defect.py @@ -296,6 +296,7 @@ def make( uc_structure=defect.structure, relax_maker=self.bulk_relax_maker, sc_mat=supercell_matrix, + get_locpot=self.get_planar_locpot, ) sc_mat = get_sc_job.output["sc_mat"] lattice = get_sc_job.output["sc_struct"].lattice diff --git a/src/atomate2/common/jobs/defect.py b/src/atomate2/common/jobs/defect.py index bf5e5fd459..e5ee1d92fc 100644 --- a/src/atomate2/common/jobs/defect.py +++ b/src/atomate2/common/jobs/defect.py @@ -243,6 +243,7 @@ def bulk_supercell_calculation( uc_structure: Structure, relax_maker: RelaxMaker, sc_mat: NDArray | None = None, + get_plnr_locpot: Callable | None = None, ) -> Response: """Bulk Supercell calculation. @@ -256,12 +257,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_plnr_locpot is None: + + def get_plnr_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) @@ -280,7 +288,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_plnr_locpot(relax_output), } flow = Flow([relax_job], output=summary_d) return Response(replace=flow) From 1700c81885c42300b15fa88b717b58ebb4dd66c9 Mon Sep 17 00:00:00 2001 From: "@jmmshn" Date: Fri, 20 Oct 2023 11:30:54 -0700 Subject: [PATCH 3/6] prv_dir --- src/atomate2/common/flows/defect.py | 2 +- src/atomate2/common/jobs/defect.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index 0e28a6232a..3d25ec17ec 100644 --- a/src/atomate2/common/flows/defect.py +++ b/src/atomate2/common/flows/defect.py @@ -296,7 +296,7 @@ def make( uc_structure=defect.structure, relax_maker=self.bulk_relax_maker, sc_mat=supercell_matrix, - get_locpot=self.get_planar_locpot, + get_planar_locpot=self.get_planar_locpot, ) sc_mat = get_sc_job.output["sc_mat"] lattice = get_sc_job.output["sc_struct"].lattice diff --git a/src/atomate2/common/jobs/defect.py b/src/atomate2/common/jobs/defect.py index e5ee1d92fc..3ac472f09b 100644 --- a/src/atomate2/common/jobs/defect.py +++ b/src/atomate2/common/jobs/defect.py @@ -243,7 +243,7 @@ def bulk_supercell_calculation( uc_structure: Structure, relax_maker: RelaxMaker, sc_mat: NDArray | None = None, - get_plnr_locpot: Callable | None = None, + get_planar_locpot: Callable | None = None, ) -> Response: """Bulk Supercell calculation. @@ -265,9 +265,9 @@ def bulk_supercell_calculation( Response: Output a dictionary containing the bulk supercell calculation summary. """ - if get_plnr_locpot is None: + if get_planar_locpot is None: - def get_plnr_locpot(tdoc): + def get_planar_locpot(tdoc): return tdoc.calcs_reversed[0].output.locpot logger.info("Running bulk supercell calculation. Running...") @@ -288,7 +288,7 @@ def get_plnr_locpot(tdoc): "sc_mat": sc_mat.tolist(), "dir_name": relax_output.dir_name, "uuid": relax_job.uuid, - "locpot_plnr": get_plnr_locpot(relax_output), + "locpot_plnr": get_planar_locpot(relax_output), } flow = Flow([relax_job], output=summary_d) return Response(replace=flow) From a2e6c2972e010769fd02053d1ddac4a4f032f6a5 Mon Sep 17 00:00:00 2001 From: "@jmmshn" Date: Fri, 20 Oct 2023 12:02:48 -0700 Subject: [PATCH 4/6] strip hostname --- src/atomate2/common/jobs/defect.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/atomate2/common/jobs/defect.py b/src/atomate2/common/jobs/defect.py index 3ac472f09b..faa6cd2d1d 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 @@ -210,6 +211,7 @@ def get_supercell_from_prv_calc( Response: Output containing the supercell transformation and the dir_name """ + 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( From 7e85aeb327782203e26adb2df40f98c7abfeb4e2 Mon Sep 17 00:00:00 2001 From: "@jmmshn" Date: Fri, 27 Oct 2023 00:00:46 -0700 Subject: [PATCH 5/6] update --- src/atomate2/common/flows/defect.py | 5 ++++- src/atomate2/vasp/flows/defect.py | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index 3d25ec17ec..9f72c4f640 100644 --- a/src/atomate2/common/flows/defect.py +++ b/src/atomate2/common/flows/defect.py @@ -25,6 +25,7 @@ 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__) @@ -350,7 +351,9 @@ def make( ) @abstractmethod - def sc_entry_and_locpot_from_prv(self, previous_dir: str) -> Structure: + def sc_entry_and_locpot_from_prv( + self, previous_dir: str + ) -> tuple[ComputedStructureEntry, dict]: """Copy the output ComputedStructureEntry and Locpot from previous directory. Parameters diff --git a/src/atomate2/vasp/flows/defect.py b/src/atomate2/vasp/flows/defect.py index 3c4eb2efa5..6e5577e5c5 100644 --- a/src/atomate2/vasp/flows/defect.py +++ b/src/atomate2/vasp/flows/defect.py @@ -167,7 +167,9 @@ class FormationEnergyMaker(defect_flows.FormationEnergyMaker): bulk_relax_maker: BaseVaspMaker | None = None name: str = "formation energy" - def sc_entry_and_locpot_from_prv(self, previous_dir: str) -> ComputedStructureEntry: + 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 From 5cbd8c4bc40aada532cc4527dc2dad6bc0b9ef07 Mon Sep 17 00:00:00 2001 From: "@jmmshn" Date: Fri, 27 Oct 2023 00:04:36 -0700 Subject: [PATCH 6/6] remove commented block --- src/atomate2/common/flows/defect.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/atomate2/common/flows/defect.py b/src/atomate2/common/flows/defect.py index 9f72c4f640..3980c42d34 100644 --- a/src/atomate2/common/flows/defect.py +++ b/src/atomate2/common/flows/defect.py @@ -332,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,