Merge pull request #1167 from esoteric-ephemera/run_stats

Correct typing of TaskDoc.run_stats, add number of electronic steps per ionic step to TaskDoc

Author: esoteric-ephemera

.github/workflows/testing.yml

@@ -81,13 +81,13 @@ jobs:
 
       - name: Install Python dependencies
         shell: bash -l {0}
+        env:
+          CC: gcc-14
+          CXX: g++-14
         run: |
           python${{ matrix.python-version }} -m pip install --upgrade pip pip-tools
           python${{ matrix.python-version }} -m pip install --user `grep numpy ${{ matrix.package }}/requirements/${{ matrix.os }}_py${{ matrix.python-version }}_extras.txt`
           python${{ matrix.python-version }} -m pip install --user -r ${{ matrix.package }}/requirements/${{ matrix.os }}_py${{ matrix.python-version }}_extras.txt
-        env:
-          CC: gcc-10
-          CXX: g++-10
 
       - name: Install editable emmet-core if needed
         shell: bash -l {0}

emmet-core/emmet/core/defect.py

@@ -3,13 +3,13 @@
 from pydantic import Field
 
 from emmet.core.tasks import TaskDoc, _VOLUMETRIC_FILES
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 from pymatgen.analysis.defects.core import Defect
 from monty.json import MontyDecoder
 from pydantic import BaseModel
 
 if TYPE_CHECKING:
-    from typing import Any, Dict, Optional, Tuple, Union
+    from typing import Any, Dict, Tuple, Union
     from pathlib import Path
 
 mdecoder = MontyDecoder().process_decoded
@@ -32,13 +32,13 @@ class DefectInfo(BaseModel):
         description="Unit cell representation of the defect object.",
     )
 
-    charge_state: int = Field(
+    charge_state: Optional[int] = Field(
         None,
         title="Charge State",
         description="Charge state of the defect.",
     )
 
-    supercell_matrix: list = Field(
+    supercell_matrix: Optional[list] = Field(
         None,
         title="Supercell Matrix",
         description="Supercell matrix used to construct the defect supercell.",

emmet-core/emmet/core/elasticity.py

@@ -195,10 +195,10 @@ class ElasticityDoc(PropertyDoc):
     thermal_conductivity: Optional[ThermalConductivity] = Field(
         None, description="Thermal conductivity"
     )
-    young_modulus: float = Field(
+    young_modulus: Optional[float] = Field(
         None, description="Young's modulus (SI units)", alias="youngs_modulus"
     )
-    universal_anisotropy: float = Field(
+    universal_anisotropy: Optional[float] = Field(
         None, description="Universal elastic anisotropy"
     )
     homogeneous_poisson: Optional[float] = Field(

emmet-core/emmet/core/feff/task.py

@@ -69,6 +69,9 @@ def xas_spectrum(self) -> XAS:
             structure = self.structure
             absorbing_index = self.absorbing_atom
             absorbing_element = self.absorbing_element
+            if isinstance(absorbing_element, Species):
+                absorbing_element = absorbing_element.element
+
             edge = self.edge
             spectrum_type = str(self.spectrum_type)
 

emmet-core/emmet/core/molecules/trajectory.py

@@ -23,8 +23,9 @@ class ForcesDoc(PropertyDoc):
 
     forces: List[List[float]] = Field(..., description="Atomic forces (units: Ha/Bohr)")
 
-    precise_forces: Optional[List[List[float]]] = Field(
-        None, description="High-precision atomic forces (units: Ha/Bohr)"
+    precise_forces: List[Optional[List[float]]] = Field(
+        default_factory=list,
+        description="High-precision atomic forces (units: Ha/Bohr)",
     )
 
     pcm_forces: Optional[List[List[float]]] = Field(
@@ -159,44 +160,44 @@ class TrajectoryDoc(PropertyDoc):
     )
 
     pcm_forces: List[Optional[List[List[List[float]]]]] = Field(
-        None,
+        default_factory=list,
         description="Electrostatic atomic forces from polarizable continuum model (PCM) implicit solvation "
         "for each optimization step for each optimization trajectory (units: Ha/Bohr).",
     )
 
     cds_forces: List[Optional[List[List[List[float]]]]] = Field(
-        None,
+        default_factory=list,
         description="Atomic force contributions from cavitation, dispersion, and structural rearrangement in the SMx "
         "family of implicit solvent models, for each optimization step for each optimization trajectory "
         "(units: Ha/Bohr)",
     )
 
     mulliken_partial_charges: List[Optional[List[List[float]]]] = Field(
-        None,
+        default_factory=list,
         description="Partial charges of each atom for each optimization step for each optimization trajectory, using "
         "the Mulliken method",
     )
 
     mulliken_partial_spins: List[Optional[List[List[float]]]] = Field(
-        None,
+        default_factory=list,
         description="Partial spins of each atom for each optimization step for each optimization trajectory, using "
         "the Mulliken method",
     )
 
     resp_partial_charges: List[Optional[List[List[float]]]] = Field(
-        None,
+        default_factory=list,
         description="Partial charges of each atom for each optimization step for each optimization trajectory, using "
         "the restrained electrostatic potential (RESP) method",
     )
 
     dipole_moments: List[Optional[List[List[float]]]] = Field(
-        None,
+        default_factory=list,
         description="Molecular dipole moment for each optimization step for each optimization trajectory, "
         "(units: Debye)",
     )
 
     resp_dipole_moments: List[Optional[List[List[float]]]] = Field(
-        None,
+        default_factory=list,
         description="Molecular dipole moment for each optimization step for each optimization trajectory, "
         "using the restrainted electrostatic potential (RESP) method (units: Debye)",
     )
@@ -247,21 +248,27 @@ def as_trajectories(self) -> List[Trajectory]:
             num_steps = len(mols)
 
             # Frame (structure) properties
-            frame_props = {
-                "energies": self.energies[ii],
-                "dipole_moments": self.dipole_moments[ii],
-                "resp_dipole_moments": self.resp_dipole_moments[ii],
-            }
+            frame_props = {"energies": self.energies[ii]}
+            for prop in (
+                "dipole_moments",
+                "resp_dipole_moments",
+            ):
+                frame_props[prop] = []
+                if (vals := getattr(self, prop, None)) is not None:
+                    frame_props[prop] = vals[ii]
 
             # Site (atomic) properties
-            site_props = {
-                "forces": self.forces[ii],
-                "pcm_forces": self.pcm_forces[ii],
-                "cds_forces": self.cds_forces[ii],
-                "mulliken_partial_charges": self.mulliken_partial_charges[ii],
-                "mulliken_partial_spins": self.mulliken_partial_spins[ii],
-                "resp_partial_charges": self.resp_partial_charges[ii],
-            }
+            site_props = {"forces": self.forces[ii]}
+            for prop in (
+                "pcm_forces",
+                "cds_forces",
+                "mulliken_partial_charges",
+                "mulliken_partial_spins",
+                "resp_partial_charges",
+            ):
+                site_props[prop] = []
+                if (vals := getattr(self, prop, None)) is not None:
+                    site_props[prop] = vals[ii]
 
             # Convert into a Trajectory object
             traj_frame_props = list()

emmet-core/emmet/core/openmm/calculations.py

@@ -13,31 +13,31 @@ class CalculationsDoc(BaseModel):
     In each field, calculations are listed sequentially, in the order they were run.
     """
 
-    task_names: List[str] = Field(None, description="Names of tasks.")
+    task_names: Optional[List[str]] = Field(None, description="Names of tasks.")
 
-    calc_types: List[str] = Field(None, description="Types of calculations.")
+    calc_types: Optional[List[str]] = Field(None, description="Types of calculations.")
 
-    elapsed_times: List[Union[float, None]] = Field(
+    elapsed_times: Optional[List[Union[float, None]]] = Field(
         None, description="Elapsed time for calculations."
     )
 
-    steps: List[Union[float, None]] = Field(
+    steps: Optional[List[Union[float, None]]] = Field(
         None, description="n_steps for calculations."
     )
 
-    step_sizes: List[Union[float, None]] = Field(
+    step_sizes: Optional[List[Union[float, None]]] = Field(
         None, description="Step sizes for each calculations."
     )
 
-    temperatures: List[Union[float, None]] = Field(
+    temperatures: Optional[List[Union[float, None]]] = Field(
         None, description="Temperature for each calculations."
     )
 
-    pressures: List[Union[float, None]] = Field(
+    pressures: Optional[List[Union[float, None]]] = Field(
         None, description="Pressure for each calculations."
     )
 
-    friction_coefficients: List[Union[float, None]] = Field(
+    friction_coefficients: Optional[List[Union[float, None]]] = Field(
         None,
         description="Friction coefficients for each calculations.",
     )

emmet-core/emmet/core/openmm/tasks.py

@@ -249,12 +249,14 @@ class OpenMMInterchange(BaseModel):
     """An object to sit in the place of the Interchance object
     and serialize the OpenMM system, topology, and state."""
 
-    system: str = Field(None, description="An XML file representing the OpenMM system.")
-    state: str = Field(
+    system: Optional[str] = Field(
+        None, description="An XML file representing the OpenMM system."
+    )
+    state: Optional[str] = Field(
         None,
         description="An XML file representing the OpenMM state.",
     )
-    topology: str = Field(
+    topology: Optional[str] = Field(
         None,
         description="An XML file representing an OpenMM topology object."
         "This must correspond to the atom ordering in the system.",

emmet-core/emmet/core/qchem/molecule.py

@@ -161,7 +161,7 @@ class MoleculeDoc(CoreMoleculeDoc):
         None, description="Standardized hash of the InChI for this molecule"
     )
 
-    calc_types: Mapping[str, CalcType] = Field(  # type: ignore
+    calc_types: Optional[Mapping[str, CalcType]] = Field(  # type: ignore
         None,
         description="Calculation types for all the calculations that make up this molecule",
     )

emmet-core/emmet/core/tasks.py

@@ -5,7 +5,7 @@
 from collections import OrderedDict
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Type, TypeVar, Union
 
 import numpy as np
 from monty.json import MontyDecoder
@@ -99,7 +99,7 @@ class OutputDoc(BaseModel):
     density: Optional[float] = Field(None, description="Density of in units of g/cc.")
     energy: Optional[float] = Field(None, description="Total Energy in units of eV.")
     forces: Optional[List[List[float]]] = Field(
-        None, description="The force on each atom in units of eV/A^2."
+        None, description="The force on each atom in units of eV/A."
     )
     stress: Optional[List[List[float]]] = Field(
         None, description="The stress on the cell in units of kB."
@@ -147,10 +147,10 @@ def from_vasp_calc_doc(
         OutputDoc
             The calculation output summary.
         """
-        if calc_doc.output.ionic_steps is not None:
+        if calc_doc.output.ionic_steps:
             forces = calc_doc.output.ionic_steps[-1].forces
             stress = calc_doc.output.ionic_steps[-1].stress
-        elif trajectory is not None:
+        elif trajectory:
             ionic_steps = trajectory.frame_properties
             forces = ionic_steps[-1]["forces"]
             stress = ionic_steps[-1]["stress"]
@@ -434,7 +434,7 @@ class TaskDoc(StructureMetadata, extra="allow"):
         description="Identifier for this calculation; should provide rough information about the calculation origin and purpose.",
     )
 
-    run_stats: Optional[RunStatistics] = Field(
+    run_stats: Optional[Mapping[str, RunStatistics]] = Field(
         None,
         description="Summary of runtime statistics for each calculation in this task",
     )
@@ -982,9 +982,10 @@ def _parse_additional_json(dir_name: Path) -> Dict[str, Any]:
 def _get_max_force(calc_doc: Calculation) -> Optional[float]:
     """Get max force acting on atoms from a calculation document."""
     if calc_doc.output.ionic_steps:
-        forces: Optional[Union[np.ndarray, List]] = calc_doc.output.ionic_steps[
-            -1
-        ].forces
+        forces: Optional[Union[np.ndarray, List]] = None
+        if calc_doc.output.ionic_steps:
+            forces = calc_doc.output.ionic_steps[-1].forces
+
         structure = calc_doc.output.structure
         if forces:
             forces = np.array(forces)

emmet-core/emmet/core/vasp/calculation.py

@@ -9,7 +9,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 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
@@ -341,12 +341,21 @@ class IonicStep(BaseModel):  # type: ignore
     electronic_steps: Optional[List[ElectronicStep]] = Field(
         None, description="The electronic convergence steps."
     )
+    num_electronic_steps: Optional[int] = Field(
+        None, description="The number of electronic steps needed to reach convergence."
+    )
     structure: Optional[Structure] = Field(
         None, description="The structure at this step."
     )
 
     model_config = ConfigDict(extra="allow")
 
+    @model_validator(mode="after")
+    def set_elec_step_count(self):
+        if self.electronic_steps is not None:
+            self.num_electronic_steps = len(self.electronic_steps)
+        return self
+
 
 class CalculationOutput(BaseModel):
     """Document defining VASP calculation outputs."""
@@ -413,6 +422,9 @@ class CalculationOutput(BaseModel):
     ionic_steps: Optional[List[IonicStep]] = Field(
         None, description="Energy, forces, structure, etc. for each ionic step"
     )
+    num_electronic_steps: Optional[List[int]] = Field(
+        None, description="The number of electronic steps in each ionic step."
+    )
     locpot: Optional[Dict[int, List[float]]] = Field(
         None, description="Average of the local potential along the crystal axes"
     )
@@ -576,6 +588,19 @@ def from_vasp_outputs(
                 temp = str(elph_poscar.name).replace("POSCAR.T=", "").replace(".gz", "")
                 elph_structures["temperatures"].append(temp)
                 elph_structures["structures"].append(Structure.from_file(elph_poscar))
+
+        ionic_steps = (
+            vasprun.ionic_steps
+            if store_trajectory == StoreTrajectoryOption.NO
+            else None
+        )
+        num_elec_steps = None
+        if ionic_steps is not None:
+            num_elec_steps = [
+                len(ionic_step.get("electronic_steps", []) or [])
+                for ionic_step in ionic_steps
+            ]
+
         return cls(
             structure=structure,
             energy=vasprun.final_energy,
@@ -587,11 +612,8 @@ 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=ionic_steps,
+            num_electronic_steps=num_elec_steps,
             locpot=locpot_avg,
             outcar=outcar_dict,
             run_stats=RunStatistics.from_outcar(outcar) if outcar else None,