diff --git a/emmet-core/emmet/core/openff/tasks.py b/emmet-core/emmet/core/openff/tasks.py index 6bcfdb536..d103176f6 100644 --- a/emmet-core/emmet/core/openff/tasks.py +++ b/emmet-core/emmet/core/openff/tasks.py @@ -2,22 +2,23 @@ from __future__ import annotations +import zlib from dataclasses import dataclass from datetime import datetime from typing import Optional -from typing_extensions import Annotated -import zlib +from monty.json import MSONable from pydantic import ( BaseModel, + ConfigDict, Field, - PlainValidator, PlainSerializer, + PlainValidator, WithJsonSchema, ) -from monty.json import MSONable - from pymatgen.core import Structure +from typing_extensions import Annotated + from emmet.core.vasp.task_valid import TaskState # type: ignore[import-untyped] @@ -56,7 +57,7 @@ class MoleculeSpec(MSONable): openff_mol: str # a tk.Molecule object serialized with to_json -class MDTaskDocument(BaseModel, extra="allow"): # type: ignore[call-arg] +class MDTaskDocument(BaseModel): # type: ignore[call-arg] """Definition of the OpenMM task document.""" tags: Optional[list[str]] = Field( @@ -109,6 +110,8 @@ class MDTaskDocument(BaseModel, extra="allow"): # type: ignore[call-arg] description="Timestamp for the most recent calculation for this task document", ) + model_config = ConfigDict(extra="allow") + class ClassicalMDTaskDocument(MDTaskDocument): """Definition of the OpenMM task document.""" diff --git a/emmet-core/emmet/core/openmm/tasks.py b/emmet-core/emmet/core/openmm/tasks.py index 667e0a43d..4cb3fb128 100644 --- a/emmet-core/emmet/core/openmm/tasks.py +++ b/emmet-core/emmet/core/openmm/tasks.py @@ -5,20 +5,20 @@ import io from pathlib import Path from typing import Optional, Union -import pandas as pd # type: ignore[import-untyped] import openmm +import pandas as pd # type: ignore[import-untyped] from openmm import XmlSerializer from openmm.app import Simulation from openmm.app.pdbfile import PDBFile -from emmet.core.vasp.task_valid import TaskState # type: ignore[import-untyped] -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from emmet.core.openff import MDTaskDocument # type: ignore[import-untyped] from emmet.core.openff.tasks import CompressedStr # type: ignore[import-untyped] +from emmet.core.vasp.task_valid import TaskState # type: ignore[import-untyped] -class CalculationInput(BaseModel, extra="allow"): # type: ignore[call-arg] +class CalculationInput(BaseModel): # type: ignore[call-arg] """OpenMM input settings for a job, these are the attributes of the OpenMMMaker.""" n_steps: Optional[int] = Field( @@ -100,6 +100,8 @@ class CalculationInput(BaseModel, extra="allow"): # type: ignore[call-arg] description="Whether to embed the trajectory blob in CalculationOutput.", ) + model_config = ConfigDict(extra="allow") + class CalculationOutput(BaseModel): """OpenMM calculation output files and extracted data.""" diff --git a/emmet-core/emmet/core/vasp/calculation.py b/emmet-core/emmet/core/vasp/calculation.py index 9eeb5a764..e87c771b6 100644 --- a/emmet-core/emmet/core/vasp/calculation.py +++ b/emmet-core/emmet/core/vasp/calculation.py @@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np -from pydantic import BaseModel, Extra, Field +from pydantic import BaseModel, ConfigDict, Field from pymatgen.command_line.bader_caller import bader_analysis_from_path from pymatgen.command_line.chargemol_caller import ChargemolAnalysis from pymatgen.core.lattice import Lattice @@ -301,7 +301,7 @@ class ElectronPhononDisplacedStructures(BaseModel): ) -class ElectronicStep(BaseModel, extra=Extra.allow): # type: ignore +class ElectronicStep(BaseModel): # type: ignore """Document defining the information at each electronic step. Note, not all the information will be available at every step. @@ -324,8 +324,10 @@ class ElectronicStep(BaseModel, extra=Extra.allow): # type: ignore e_wo_entrp: Optional[float] = Field(None, description="The energy without entropy.") e_0_energy: Optional[float] = Field(None, description="The internal energy.") + model_config = ConfigDict(extra="allow") -class IonicStep(BaseModel, extra=Extra.allow): # type: ignore + +class IonicStep(BaseModel): # type: ignore """Document defining the information at each ionic step.""" e_fr_energy: Optional[float] = Field(None, description="The free energy.") @@ -342,6 +344,8 @@ class IonicStep(BaseModel, extra=Extra.allow): # type: ignore None, description="The structure at this step." ) + model_config = ConfigDict(extra="allow") + class CalculationOutput(BaseModel): """Document defining VASP calculation outputs.""" @@ -582,9 +586,11 @@ def from_vasp_outputs( frequency_dependent_dielectric=freq_dependent_diel, elph_displaced_structures=elph_structures, dos_properties=dosprop_dict, - ionic_steps=vasprun.ionic_steps - if store_trajectory == StoreTrajectoryOption.NO - else None, + ionic_steps=( + vasprun.ionic_steps + if store_trajectory == StoreTrajectoryOption.NO + else None + ), locpot=locpot_avg, outcar=outcar_dict, run_stats=RunStatistics.from_outcar(outcar) if outcar else None,