Skip to content

Commit

Permalink
precommit
Browse files Browse the repository at this point in the history
  • Loading branch information
@esoteric-ephemera
esoteric-ephemera committed Jan 17, 2025
1 parent 3df7048 commit 6e97455
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 67 deletions.
116 changes: 61 additions & 55 deletions emmet-core/emmet/core/tasks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -325,61 +325,62 @@ def from_vasp_calc_docs(
max_force=max_force,
warnings=warnings,
errors=errors,
)
)


class EmmetComputedEntry(BaseModel):
"""Fixed-schema version of pymatgen ComputedEntry."""

energy : float
composition: dict[str,float]
entry_id : MPID | None = None
correction : float | None = None
energy_adjustments : list[float] | None = None
potcar_spec : list[PotcarSpec] | None = None
run_type : RunType | None = None
is_hubbard : bool = False
hubbards : dict[str,float] | None = None
oxide_type : str | None = None
aspherical : bool = False
last_updated : datetime | None = None
energy: float
composition: dict[str, float]
entry_id: MPID | None = None
correction: float | None = None
energy_adjustments: list[float] | None = None
potcar_spec: list[PotcarSpec] | None = None
run_type: RunType | None = None
is_hubbard: bool = False
hubbards: dict[str, float] | None = None
oxide_type: str | None = None
aspherical: bool = False
last_updated: datetime | None = None

def get_computed_entry(self) -> ComputedEntry:
"""Get pymatgen computed entry."""
return ComputedEntry.from_dict({
"correction": self.correction,
"entry_id": self.entry_id,
"composition": self.composition,
"energy": self.energy,
"parameters": {
"potcar_spec": (
[ps.model_dump() for ps in self.potcar_spec] if self.potcar_spec is not None
else []
),
"run_type": str(self.run_type),
"is_hubbard": self.is_hubbard,
"hubbards": self.hubbards,
},
"data": {
"oxide_type": self.oxide_type,
"aspherical": self.aspherical,
"last_updated": str(self.last_updated),
},
})

return ComputedEntry.from_dict(
{
"correction": self.correction,
"entry_id": self.entry_id,
"composition": self.composition,
"energy": self.energy,
"parameters": {
"potcar_spec": (
[ps.model_dump() for ps in self.potcar_spec]
if self.potcar_spec is not None
else []
),
"run_type": str(self.run_type),
"is_hubbard": self.is_hubbard,
"hubbards": self.hubbards,
},
"data": {
"oxide_type": self.oxide_type,
"aspherical": self.aspherical,
"last_updated": str(self.last_updated),
},
}
)


class DbTaskDoc(StructureMetadata):
"""Calculation-level details about VASP calculations that power Materials Project.
This schema is intended to be fixed for database best practices.
"""

tags: list[str] | None = Field(
None, title="tag", description="Metadata tagged to a given task."
)
dir_name: str | None = Field(
None, description="The directory for this VASP task"
)
dir_name: str | None = Field(None, description="The directory for this VASP task")

calcs_reversed: list[Calculation] | None = Field(
None,
Expand All @@ -391,9 +392,7 @@ class DbTaskDoc(StructureMetadata):
None, description="Final output structure from the task"
)

task_type: TaskType | None = Field(
None, description="The type of calculation."
)
task_type: TaskType | None = Field(None, description="The type of calculation.")

run_type: RunType | None = Field(
None, description="The functional used in the calculation."
Expand Down Expand Up @@ -508,15 +507,17 @@ def model_post_init(self, __context: Any) -> None:
# TODO: remove after imposing TaskDoc schema on older tasks in collection
if self.structure is None:
self.structure = calcs_reversed[0].output.structure

# Set the computed entry if not set
if (
not self.entry
and self.calcs_reversed
and getattr(self.calcs_reversed[0].output, "structure", None)
):
use_pymatgen_rep = getattr(self, "_use_pymatgen_rep", False)
self.entry = self.get_entry(self.calcs_reversed, self.task_id, use_pymatgen_rep=use_pymatgen_rep)
self.entry = self.get_entry(
self.calcs_reversed, self.task_id, use_pymatgen_rep=use_pymatgen_rep
)

# Make sure that the datetime field is properly formatted
# (Unclear when this is not the case, please leave comment if observed)
Expand Down Expand Up @@ -722,7 +723,7 @@ def from_vasprun(
def get_entry(
calcs_reversed: list[Calculation],
task_id: Optional[Union[MPID, str]] = None,
use_pymatgen_rep : bool = False
use_pymatgen_rep: bool = False,
) -> EmmetComputedEntry | ComputedEntry:
"""
Get a computed entry from a list of VASP calculation documents.
Expand All @@ -741,17 +742,20 @@ def get_entry(
ComputedEntry
A computed entry.
"""
inp_kwargs = {k : getattr(calcs_reversed[0].input,k,None) for k in ("potcar_spec","is_hubbard","hubbards")}
inp_kwargs = {
k: getattr(calcs_reversed[0].input, k, None)
for k in ("potcar_spec", "is_hubbard", "hubbards")
}
ce = EmmetComputedEntry(
energy = calcs_reversed[0].output.energy,
energy=calcs_reversed[0].output.energy,
composition=calcs_reversed[0].output.structure.composition.as_dict(),
entry_id=task_id,
correction=0.,
correction=0.0,
run_type=calcs_reversed[0].run_type,
oxide_type=oxide_type(calcs_reversed[0].output.structure),
aspherical=calcs_reversed[0].input.parameters.get("LASPH", False),
last_updated=utcnow(),
**inp_kwargs
**inp_kwargs,
)
if use_pymatgen_rep:
return ce.get_computed_entry()
Expand Down Expand Up @@ -812,6 +816,7 @@ def structure_entry(self) -> ComputedStructureEntry:
entry_id=ce.entry_id,
)


class TaskDoc(DbTaskDoc, extra="allow"):
"""Calculation-level details about VASP calculations that power Materials Project."""

Expand All @@ -835,7 +840,7 @@ class TaskDoc(DbTaskDoc, extra="allow"):
None, description="The ComputedEntry from the task doc"
)

_use_pymatgen_rep : bool = True
_use_pymatgen_rep: bool = True

@classmethod
def from_directory(
Expand Down Expand Up @@ -883,18 +888,19 @@ def from_directory(

db_task = DbTaskDoc.from_directory(
dir_name,
volumetric_files = volumetric_files,
additional_fields = additional_fields,
volume_change_warning_tol = volume_change_warning_tol,
task_names = task_names,
volumetric_files=volumetric_files,
additional_fields=additional_fields,
volume_change_warning_tol=volume_change_warning_tol,
task_names=task_names,
**vasp_calculation_kwargs,
)
config = db_task.model_dump()
config["entry"] = db_task.entry.get_computed_entry()
if additional_json:
config.update(additional_json= additional_json)
config.update(additional_json=additional_json)
return cls(**config)



class TrajectoryDoc(BaseModel):
"""Model for task trajectory data."""

Expand Down
31 changes: 19 additions & 12 deletions emmet-core/emmet/core/vasp/calculation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@

VaspObjectTypes = BandStructure | Dos | VolumetricData | Trajectory | Wavecar | Procar


class VaspObject(ValueEnum):
"""Types of VASP data objects."""

Expand Down Expand Up @@ -78,22 +79,26 @@ class CalculationBaseModel(BaseModel):
def get(self, key: Any, default_value: Optional[Any] = None) -> Any:
return getattr(self, key, default_value)


class PotcarSummaryStats(BaseModel):
"""Specific format of the POTCAR summary statistics used for POTCAR matching."""

keywords : dict[str, list[str]] | None = Field(
None, description=(
keywords: dict[str, list[str]] | None = Field(
None,
description=(
"dict of the form `{'header': [list of POTCAR header keywords],"
"'data': [list of POTCAR body keywords]}`"
)
),
)
stats : dict[str,dict[str,float]] | None = Field(
None, description=(
stats: dict[str, dict[str, float]] | None = Field(
None,
description=(
"dict of the form `{'header': {statistic: value}, 'data': {...}}` "
"based on the data in the POTCAR"
)
),
)


class PotcarSpec(BaseModel):
"""Document defining a VASP POTCAR specification."""

Expand Down Expand Up @@ -219,14 +224,16 @@ def from_vasprun(cls, vasprun: Vasprun) -> "CalculationInput":
class RunStatistics(BaseModel):
"""Summary of the run statistics for a VASP calculation."""

average_memory: float = Field(0., description="The average memory used in kb")
max_memory: float = Field(0., description="The maximum memory used in kb")
elapsed_time: float = Field(0., description="The real time elapsed in seconds")
system_time: float = Field(0., description="The system CPU time in seconds")
average_memory: float = Field(0.0, description="The average memory used in kb")
max_memory: float = Field(0.0, description="The maximum memory used in kb")
elapsed_time: float = Field(0.0, description="The real time elapsed in seconds")
system_time: float = Field(0.0, description="The system CPU time in seconds")
user_time: float = Field(
0., description="The user CPU time spent by VASP in seconds"
0.0, description="The user CPU time spent by VASP in seconds"
)
total_time: float = Field(
0.0, description="The total CPU time for this calculation"
)
total_time: float = Field(0., description="The total CPU time for this calculation")
cores: int = Field(0, description="The number of cores used by VASP")

@classmethod
Expand Down

0 comments on commit 6e97455

Please sign in to comment.