Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix failing tests from Pydantic v2 migration #558

Merged
merged 10 commits into from
Oct 10, 2023
4 changes: 2 additions & 2 deletions src/atomate2/common/schemas/defects.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,13 +150,13 @@ class CCDDocument(BaseModel):
"in charge state (q2).",
)

static_uuids1: List[str] = Field(
static_uuids1: Optional[List[str]] = Field(
None,
description="UUIDs of distorted calculations for the defect (supercell) in "
"charge state (q1).",
)

static_uuids2: List[str] = Field(
static_uuids2: Optional[List[str]] = Field(
None,
description="UUIDs of distorted calculations for the defect (supercell) in "
"charge state (q2).",
Expand Down
18 changes: 10 additions & 8 deletions src/atomate2/common/schemas/phonons.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,12 +66,14 @@ class ThermalDisplacementData(BaseModel):
class PhononUUIDs(BaseModel):
"""Collection to save all uuids connected to the phonon run."""

optimization_run_uuid: str = Field(None, description="optimization run uuid")
displacements_uuids: List[str] = Field(
optimization_run_uuid: Optional[str] = Field(
None, description="optimization run uuid"
)
displacements_uuids: Optional[List[str]] = Field(
None, description="The uuids of the displacement jobs."
)
static_run_uuid: str = Field(None, description="static run uuid")
born_run_uuid: str = Field(None, description="born run uuid")
static_run_uuid: Optional[str] = Field(None, description="static run uuid")
born_run_uuid: Optional[str] = Field(None, description="born run uuid")


class ForceConstants(MSONable):
Expand Down Expand Up @@ -153,17 +155,17 @@ class PhononBSDOSDoc(StructureMetadata):
)

# needed, e.g. to compute Grueneisen parameter etc
force_constants: ForceConstants = Field(
force_constants: Optional[ForceConstants] = Field(
None, description="Force constants between every pair of atoms in the structure"
)

born: List[Matrix3D] = Field(
born: Optional[List[Matrix3D]] = Field(
None,
description="born charges as computed from phonopy. Only for symmetrically "
"different atoms",
)

epsilon_static: Matrix3D = Field(
epsilon_static: Optional[Matrix3D] = Field(
None, description="The high-frequency dielectric constant"
)

Expand All @@ -186,7 +188,7 @@ class PhononBSDOSDoc(StructureMetadata):
"Field including all relevant job directories"
)

uuids: PhononUUIDs = Field("Field including all relevant uuids")
uuids: Optional[PhononUUIDs] = Field("Field including all relevant uuids")

@classmethod
def from_forces_born(
Expand Down
10 changes: 5 additions & 5 deletions src/atomate2/cp2k/schemas/calculation.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from typing import Any, Dict, List, Optional, Tuple, Union

from jobflow.utils import ValueEnum
from pydantic import BaseModel, Field, validator
from pydantic import BaseModel, Field, field_validator
from pymatgen.command_line.bader_caller import BaderAnalysis
from pymatgen.core.structure import Molecule, Structure
from pymatgen.core.trajectory import Trajectory
Expand Down Expand Up @@ -82,16 +82,16 @@ class CalculationInput(BaseModel):
description="CP2K global parameters used in the last calc of this task.",
)

@validator("atomic_kind_info")
def remove_unnecessary(self, atomic_kind_info):
@field_validator("atomic_kind_info", mode="before")
def remove_unnecessary(cls, atomic_kind_info):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

self should stay as is.

"""Remove unnecessary entry from atomic_kind_info."""
for k in atomic_kind_info:
if "total_pseudopotential_energy" in atomic_kind_info[k]:
del atomic_kind_info[k]["total_pseudopotential_energy"]
return atomic_kind_info

@validator("dft")
def cleanup_dft(self, dft):
@field_validator("dft", mode="before")
def cleanup_dft(cls, dft):
"""Convert UKS strings to UKS=True."""
if any(v.upper() == "UKS" for v in dft.values()):
dft["UKS"] = True
Expand Down
48 changes: 28 additions & 20 deletions src/atomate2/cp2k/schemas/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,8 +189,8 @@ class OutputSummary(BaseModel):
bandgap: Optional[float] = Field(
None, description="The DFT bandgap for the last calculation"
)
cbm: float = Field(None, description="CBM for this calculation")
vbm: float = Field(None, description="VBM for this calculation")
cbm: Optional[float] = Field(None, description="CBM for this calculation")
vbm: Optional[float] = Field(None, description="VBM for this calculation")
forces: List[Vector3D] = Field(
None, description="Forces on atoms from the last calculation"
)
Expand Down Expand Up @@ -234,51 +234,59 @@ def from_cp2k_calc_doc(cls, calc_doc: Calculation) -> "OutputSummary":
class TaskDocument(StructureMetadata, MoleculeMetadata):
"""Definition of CP2K task document."""

dir_name: str = Field(None, description="The directory for this CP2K task")
dir_name: Optional[str] = Field(
None, description="The directory for this CP2K task"
)
last_updated: str = Field(
default_factory=datetime_str,
description="Timestamp for this task document was last updated",
)
completed_at: str = Field(
completed_at: Optional[str] = Field(
None, description="Timestamp for when this task was completed"
)
input: InputSummary = Field(None, description="The input to the first calculation")
output: OutputSummary = Field(
input: Optional[InputSummary] = Field(
None, description="The input to the first calculation"
)
output: Optional[OutputSummary] = Field(
None, description="The output of the final calculation"
)
structure: Union[Structure, Molecule] = Field(
None, description="Final output structure from the task"
)
state: Status = Field(None, description="State of this task")
included_objects: List[Cp2kObject] = Field(
state: Optional[Status] = Field(None, description="State of this task")
included_objects: Optional[List[Cp2kObject]] = Field(
None, description="List of CP2K objects included with this task document"
)
cp2k_objects: Dict[Cp2kObject, Any] = Field(
cp2k_objects: Optional[Dict[Cp2kObject, Any]] = Field(
None, description="CP2K objects associated with this task"
)
entry: ComputedEntry = Field(
entry: Optional[ComputedEntry] = Field(
None, description="The ComputedEntry from the task doc"
)
analysis: AnalysisSummary = Field(
analysis: Optional[AnalysisSummary] = Field(
None, description="Summary of structural relaxation and forces"
)
run_stats: Dict[str, RunStatistics] = Field(
run_stats: Optional[Dict[str, RunStatistics]] = Field(
None,
description="Summary of runtime statistics for each calculation in this task",
)
orig_inputs: Dict[str, Cp2kInput] = Field(
orig_inputs: Optional[Dict[str, Cp2kInput]] = Field(
None, description="Summary of the original CP2K inputs written by custodian"
)
task_label: str = Field(None, description="A description of the task")
tags: List[str] = Field(None, description="Metadata tags for this task document")
author: str = Field(None, description="Author extracted from transformations")
icsd_id: str = Field(
task_label: Optional[str] = Field(None, description="A description of the task")
tags: Optional[List[str]] = Field(
None, description="Metadata tags for this task document"
)
author: Optional[str] = Field(
None, description="Author extracted from transformations"
)
icsd_id: Optional[str] = Field(
None, description="International crystal structure database id of the structure"
)
calcs_reversed: List[Calculation] = Field(
calcs_reversed: Optional[List[Calculation]] = Field(
None, description="The inputs and outputs for all CP2K runs in this task."
)
transformations: Dict[str, Any] = Field(
transformations: Optional[Dict[str, Any]] = Field(
None,
description="Information on the structural transformations, parsed from a "
"transformations.json file",
Expand All @@ -288,7 +296,7 @@ class TaskDocument(StructureMetadata, MoleculeMetadata):
description="Information on the custodian settings used to run this "
"calculation, parsed from a custodian.json file",
)
additional_json: Dict[str, Any] = Field(
additional_json: Optional[Dict[str, Any]] = Field(
None, description="Additional json loaded from the calculation directory"
)
schema: str = Field(
Expand Down
12 changes: 6 additions & 6 deletions src/atomate2/lobster/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,14 +103,14 @@ class LobsterinModel(BaseModel):
)
cohpendenergy: float = Field(None, description="End energy for COHP computation")

gaussiansmearingwidth: float = Field(
gaussiansmearingwidth: Optional[float] = Field(
None, description="Set the smearing width in eV,default is 0.2 (eV)"
)
usedecimalplaces: int = Field(
usedecimalplaces: Optional[int] = Field(
None,
description="Set the decimal places to print in output files, default is 5",
)
cohpsteps: float = Field(
cohpsteps: Optional[float] = Field(
None, description="Number steps in COHPCAR; similar to NEDOS of VASP"
)
basisset: str = Field(None, description="basis set of computation")
Expand All @@ -121,7 +121,7 @@ class LobsterinModel(BaseModel):
saveprojectiontofile: bool = Field(
None, description="Save the results of projections"
)
lsodos: bool = Field(
lsodos: Optional[bool] = Field(
None, description="Writes DOS output from the orthonormalized LCAO basis"
)
basisfunctions: list = Field(
Expand Down Expand Up @@ -388,7 +388,7 @@ class LobsterTaskDocument(StructureMetadata):
dos: LobsterCompleteDos = Field(
None, description="pymatgen pymatgen.io.lobster.Doscar.completedos data"
)
lso_dos: LobsterCompleteDos = Field(
lso_dos: Optional[LobsterCompleteDos] = Field(
None, description="pymatgen pymatgen.io.lobster.Doscar.completedos data"
)
madelung_energies: dict = Field(
Expand All @@ -409,7 +409,7 @@ class LobsterTaskDocument(StructureMetadata):
"each site as a key and the gross population as a value.",
)

band_overlaps: dict = Field(
band_overlaps: Optional[dict] = Field(
None,
description="Band overlaps data for each k-point from"
" bandOverlaps.lobster file if it exists",
Expand Down
4 changes: 2 additions & 2 deletions tests/forcefields/test_jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ def test_m3gnet_static_maker(si_structure):
# validate job outputs
output1 = responses[job.uuid][1].output
assert isinstance(output1, ForceFieldTaskDocument)
assert output1.output.energy == approx(-10.8454, rel=1e-4)
assert output1.output.energy == approx(-10.8, abs=0.2)
assert output1.output.n_steps == 1


Expand All @@ -85,7 +85,7 @@ def test_m3gnet_relax_maker(si_structure):
# validate job outputs
output1 = responses[job.uuid][1].output
assert isinstance(output1, ForceFieldTaskDocument)
assert output1.output.energy == approx(-10.8441, rel=1e-4)
assert output1.output.energy == approx(-10.8, abs=0.2)
assert output1.output.n_steps == 14


Expand Down