diff --git a/src/atomate2/common/analysis/defects/jobs.py b/src/atomate2/common/analysis/defects/jobs.py index 43ce409e47..de96562af7 100644 --- a/src/atomate2/common/analysis/defects/jobs.py +++ b/src/atomate2/common/analysis/defects/jobs.py @@ -26,8 +26,14 @@ logger = logging.getLogger(__name__) - -logger = logging.getLogger(__name__) +__all__ = [ + "get_charged_structures", + "spawn_energy_curve_calcs", + "get_ccd_documents", + "get_supercell_from_prv_calc", + "bulk_supercell_calculation", + "spawn_defect_calcs", +] class CCDInput(BaseModel): diff --git a/src/atomate2/common/schemas/defects.py b/src/atomate2/common/schemas/defects.py index 50c99d972b..7bcdc93ebb 100644 --- a/src/atomate2/common/schemas/defects.py +++ b/src/atomate2/common/schemas/defects.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Callable, List, Optional, Tuple, Type +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import numpy as np from pydantic import BaseModel, Field @@ -12,8 +12,6 @@ logger = logging.getLogger(__name__) -__all__ = ["CCDDocument", "FormationEnergyDiagramDocument"] - class FormationEnergyDiagramDocument(BaseModel): """A document for storing a formation energy diagram. @@ -55,11 +53,11 @@ class FormationEnergyDiagramDocument(BaseModel): None, description="The directory name of the pristine supercell calculation." ) - defect_sc_dirs: dict[int, str] = Field( + defect_sc_dirs: Dict[int, str] = Field( None, description="The directory names of the charged defect calculations." ) - dielectric: float | list[list[float]] = Field( + dielectric: Union[float, List[List[float]]] = Field( None, description="The dielectric constant or tensor, can be used to compute corrections.", ) @@ -88,7 +86,7 @@ def from_FormationEnergyDiagram( def as_FormationEnergyDiagram( self, pd_entries: Optional[List[ComputedEntry]] = None - ) -> "FormationEnergyDiagram": + ) -> FormationEnergyDiagram: """Create a `FormationEnergyDiagram` object from the document. Since the `pd_entries` field is optional, this method allows the user @@ -235,7 +233,7 @@ def get_ent(struct, energy, dir_name, uuid): @classmethod def from_entries( - cls: Type["CCDDocument"], + cls, entries1: List[ComputedStructureEntry], entries2: List[ComputedStructureEntry], relaxed_uuid1: Optional[str] = None, @@ -257,7 +255,7 @@ def from_entries( """ - def find_entry(entries, uuid) -> Tuple[int, ComputedStructureEntry]: + def find_entry(entries, uuid) -> tuple[int, ComputedStructureEntry]: """Find the entry with the given UUID.""" for itr, entry in enumerate(entries): if entry.data["uuid"] == uuid: diff --git a/src/atomate2/vasp/jobs/defect.py b/src/atomate2/vasp/jobs/defect.py index 2ea4819c10..9a59e1bd36 100644 --- a/src/atomate2/vasp/jobs/defect.py +++ b/src/atomate2/vasp/jobs/defect.py @@ -8,6 +8,14 @@ from pymatgen.io.vasp import Incar from pymatgen.io.vasp.outputs import WSWQ +from atomate2.common.analysis.defects.jobs import ( + bulk_supercell_calculation, + get_ccd_documents, + get_charged_structures, + get_supercell_from_prv_calc, + spawn_defect_calcs, + spawn_energy_curve_calcs, +) from atomate2.common.files import copy_files, gunzip_files, gzip_files, rename_files from atomate2.utils.file_client import FileClient from atomate2.utils.path import strip_hostname @@ -17,6 +25,17 @@ logger = logging.getLogger(__name__) +__all__ = [ + "bulk_supercell_calculation", + "calculate_finite_diff", + "get_ccd_documents", + "get_charged_structures", + "get_supercell_from_prv_calc", + "spawn_defect_calcs", + "spawn_energy_curve_calcs", +] +# sort the list above + @job(data=WSWQ, output_schema=FiniteDifferenceDocument) def calculate_finite_diff(