Update to Defects WF

src/atomate2/common/flows/defect.py

@@ -5,6 +5,7 @@
 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
@@ -13,14 +14,13 @@
     bulk_supercell_calculation,
     get_ccd_documents,
     get_charged_structures,
+    get_defect_entry,
     get_supercell_from_prv_calc,
     spawn_defect_q_jobs,
     spawn_energy_curve_calcs,
 )
 
 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
@@ -167,6 +167,7 @@
         messy, it is recommended for each implementation of this maker to check
         some of the most important settings in the `relax_maker`.  Please see
         `FormationEnergyMaker.validate_maker` for more details.
+
     bulk_relax_maker: Maker
         If None, the same `defect_relax_maker` will be used for the bulk supercell.
         A maker to used to perform the bulk supercell calculation. For marginally
@@ -184,13 +185,64 @@
             params = ["NGX", "NGY", "NGZ", "NGXF", "NGYF", "NGZF"]
             ng_settings = dict(zip(params, ng + ngf))
             relax_maker = update_user_incar_settings(relax_maker, ng_settings)
+
     name: str
         The name of the flow created by this maker.
+
+    relax_radius:
+        The radius to include around the defect site for the relaxation.
+        If "auto", the radius will be set to the maximum that will fit inside
+        a periodic cell. If None, all atoms will be relaxed.
+
+    perturb:
+        The amount to perturb the sites in the supercell. Only perturb the
+        sites with selective dynamics set to True. So this setting only works
+        with `relax_radius`.
+
+    collect_defect_entry_data: bool
+        Whether to collect the defect entry data at the end of the flow.
+        If True, the output of all the charge states for each symmetry distinct
+        defect will be collected into a list of dictionaries that can be used
+        to create a DefectEntry. The data here can be trivially combined with
+        phase diagram data from the materials project API to create the formation
+        energy diagrams.
+
+        .. note::
+        Once we remove the requirement for explicit bulk supercell calculations,
+        this setting will be removed.  It is only needed because the bulk supercell
+        locpot is currently needed for the finite-size correction calculation.
+
+        Output format for the DefectEntry data:
+        .. code-block:: python
+        [
+            {
+                'bulk_dir_name': 'computer1:/folder1',
+                'bulk_locpot': {...},
+                'bulk_uuid': '48fb6da7-dc2b-4dcb-b1c8-1203c0f72ce3',
+                'defect_dir_name': 'computer1:/folder2',
+                'defect_entry': {...},
+                'defect_locpot': {...},
+                'defect_uuid': 'e9af2725-d63c-49b8-a01f-391540211750'
+            },
+            {
+                'bulk_dir_name': 'computer1:/folder3',
+                'bulk_locpot': {...},
+                'bulk_uuid': '48fb6da7-dc2b-4dcb-b1c8-1203c0f72ce3',
+                'defect_dir_name': 'computer1:/folder4',
+                'defect_entry': {...},
+                'defect_locpot': {...},
+                'defect_uuid': 'a1c31095-0494-4eed-9862-95311f80a993'
+            }
+        ]
+
     """
 
     defect_relax_maker: Maker
     bulk_relax_maker: Maker | None = None
     name: str = "formation energy"
+    relax_radius: float | str | None = None
+    perturb: float | None = None
+    collect_defect_entry_data: bool = False
 
     def __post_init__(self):
         """Apply post init updates."""
@@ -263,11 +315,27 @@
                 "bulk_supercell_matrix": sc_mat,
                 "bulk_supercell_uuid": get_sc_job.uuid,
             },
+            relax_radius=self.relax_radius,
+            perturb=self.perturb,
         )
         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,
+            )
+            jobs.append(collection_job)
+
         return Flow(
             jobs=jobs,
+            output=spawn_output.output,
             name=self.name,
         )
 

src/atomate2/common/jobs/defect.py

@@ -12,7 +12,9 @@
     get_matched_structure_mapping,
     get_sc_fromstruct,
 )
+from pymatgen.analysis.defects.thermo import DefectEntry
 from pymatgen.core import Lattice, Structure
+from pymatgen.entries.computed_entries import ComputedStructureEntry
 
 from atomate2.common.schemas.defects import CCDDocument
 
@@ -276,6 +278,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,
     }
     flow = Flow([relax_job], output=summary_d)
     return Response(replace=flow)
@@ -290,6 +293,8 @@ def spawn_defect_q_jobs(
     defect_index: int | str = "",
     add_info: dict | None = None,
     validate_charge: bool = True,
+    relax_radius: float | str | None = None,
+    perturb: float | None = None,
 ) -> Response:
     """Perform charge defect supercell calculations.
 
@@ -314,6 +319,18 @@ def spawn_defect_q_jobs(
         The lattice of the relaxed supercell. If provided, the lattice parameters
         of the supercell will be set to value specified.  Otherwise, the lattice it will
         only by set by `defect.structure` and `sc_mat`.
+    validate_charge:
+        Whether to validate the charge states of the defect after the atomic relaxation.
+        Assuming the final output of the relaxation is a TaskDoc, we should make sure
+        that the charge state is set properly and matches the expected charge state from
+        the input defect object.
+    relax_radius:
+        The radius to include around the defect site for the relaxation.
+        If "auto", the radius will be set to the maximum that will fit inside a periodic
+        cell. If None, all atoms will be relaxed.
+    perturb:
+        The amount to perturb the sites in the supercell. Only perturb the sites with
+        selective dynamics set to True. So this setting only works with `relax_radius`.
 
     Returns
     -------
@@ -323,7 +340,9 @@ def spawn_defect_q_jobs(
     """
     defect_q_jobs = []
     all_chg_outputs = {}
-    sc_def_struct = defect.get_supercell_structure(sc_mat=sc_mat)
+    sc_def_struct = defect.get_supercell_structure(
+        sc_mat=sc_mat, relax_radius=relax_radius, perturb=perturb
+    )
     sc_def_struct.lattice = relaxed_sc_lattice
     if sc_mat is not None:
         sc_mat = np.array(sc_mat).tolist()
@@ -357,10 +376,12 @@ def spawn_defect_q_jobs(
         defect_q_jobs.append(charged_relax)
         charged_output: TaskDoc = charged_relax.output
         all_chg_outputs[qq] = {
+            "defect": defect,
             "structure": charged_output.structure,
             "entry": charged_output.entry,
             "dir_name": charged_output.dir_name,
             "uuid": charged_relax.uuid,
+            "locpot_plnr": charged_output.calcs_reversed[0].output.locpot,
         }
         # check that the charge state was set correctly
         if validate_charge:
@@ -387,7 +408,46 @@ def check_charge_state(charge_state: int, task_structure: Structure) -> Response
     """
     if int(charge_state) != int(task_structure.charge):
         raise ValueError(
-            f"The charge state of the structure is {task_structure.charge}, "
-            f"but the charge state of the calculation is {charge_state}."
+            f"The charge of the output structure is {task_structure.charge}, "
+            f"but expect charge state from the Defect object is {charge_state}."
         )
     return True
+
+
+@job
+def get_defect_entry(charge_state_summary: dict, bulk_summary: dict):
+    """Get a defect entry from a defect calculation and a bulk calculation."""
+    bulk_c_entry = bulk_summary["sc_entry"]
+    bulk_struct_entry = ComputedStructureEntry(
+        structure=bulk_summary["sc_struct"],
+        energy=bulk_c_entry.energy,
+    )
+    bulk_dir_name = bulk_summary["dir_name"]
+    bulk_locpot = bulk_summary["locpot_plnr"]
+    defect_ent_res = []
+    for qq, qq_summary in charge_state_summary.items():
+        defect_c_entry = qq_summary["entry"]
+        defect_struct_entry = ComputedStructureEntry(
+            structure=qq_summary["structure"],
+            energy=defect_c_entry.energy,
+        )
+        defect_dir_name = qq_summary["dir_name"]
+        defect_locpot = qq_summary["locpot_plnr"]
+        defect_entry = DefectEntry(
+            defect=qq_summary["defect"],
+            charge_state=qq,
+            sc_entry=defect_struct_entry,
+            bulk_entry=bulk_struct_entry,
+        )
+        defect_ent_res.append(
+            {
+                "defect_entry": defect_entry,
+                "defect_dir_name": defect_dir_name,
+                "defect_locpot": defect_locpot,
+                "bulk_dir_name": bulk_dir_name,
+                "bulk_locpot": bulk_locpot,
+                "bulk_uuid": bulk_summary.get("uuid", None),
+                "defect_uuid": qq_summary.get("uuid", None),
+            }
+        )
+    return defect_ent_res

src/atomate2/vasp/flows/defect.py

@@ -91,6 +91,7 @@ class FormationEnergyMaker(defect_flows.FormationEnergyMaker):
         messy, it is recommended for each implementation of this maker to check
         some of the most important settings in the `relax_maker`.  Please see
         `FormationEnergyMaker.validate_maker` for more details.
+
     bulk_relax_maker: Maker
         If None, the same `defect_relax_maker` will be used for the bulk supercell.
         A maker to used to perform the bulk supercell calculation. For marginally
@@ -108,8 +109,55 @@ class FormationEnergyMaker(defect_flows.FormationEnergyMaker):
             params = ["NGX", "NGY", "NGZ", "NGXF", "NGYF", "NGZF"]
             ng_settings = dict(zip(params, ng + ngf))
             relax_maker = update_user_incar_settings(relax_maker, ng_settings)
+
     name: str
         The name of the flow created by this maker.
+
+    relax_radius:
+        The radius to include around the defect site for the relaxation.
+        If "auto", the radius will be set to the maximum that will fit inside
+        a periodic cell. If None, all atoms will be relaxed.
+
+    perturb:
+        The amount to perturb the sites in the supercell. Only perturb the
+        sites with selective dynamics set to True. So this setting only works
+        with `relax_radius`.
+
+    collect_defect_entry_data: bool
+        Whether to collect the defect entry data at the end of the flow.
+        If True, the output of all the charge states for each symmetry distinct
+        defect will be collected into a list of dictionaries that can be used
+        to create a DefectEntry. The data here can be trivially combined with
+        phase diagram data from the materials project API to create the formation
+        energy diagrams.
+
+        .. note::
+        Once we remove the requirement for explicit bulk supercell calculations,
+        this setting will be removed.  It is only needed because the bulk supercell
+        locpot is currently needed for the finite-size correction calculation.
+
+        Output format for the DefectEntry data:
+        .. code-block:: python
+        [
+            {
+                'bulk_dir_name': 'computer1:/folder1',
+                'bulk_locpot': {...},
+                'bulk_uuid': '48fb6da7-dc2b-4dcb-b1c8-1203c0f72ce3',
+                'defect_dir_name': 'computer1:/folder2',
+                'defect_entry': {...},
+                'defect_locpot': {...},
+                'defect_uuid': 'e9af2725-d63c-49b8-a01f-391540211750'
+            },
+            {
+                'bulk_dir_name': 'computer1:/folder3',
+                'bulk_locpot': {...},
+                'bulk_uuid': '48fb6da7-dc2b-4dcb-b1c8-1203c0f72ce3',
+                'defect_dir_name': 'computer1:/folder4',
+                'defect_entry': {...},
+                'defect_locpot': {...},
+                'defect_uuid': 'a1c31095-0494-4eed-9862-95311f80a993'
+            }
+        ]
     """
 
     defect_relax_maker: BaseVaspMaker = field(

tests/vasp/flows/test_defect.py

@@ -165,7 +165,9 @@ def test_formation_energy_maker(mock_vasp, clean_dir, test_dir, monkeypatch):
     )
 
     # rmaker = RelaxMaker(input_set_generator=ChargeStateRelaxSetGenerator())
-    maker = FormationEnergyMaker()
+    maker = FormationEnergyMaker(
+        relax_radius="auto", perturb=0.1, collect_defect_entry_data=True
+    )
     flow = maker.make(
         defects[0],
         supercell_matrix=[[2, 2, 0], [2, -2, 0], [0, 0, 1]],