Fix: Support 3rd and higher order elastic tensor in schema

Author: hongyi-zhao

src/atomate2/common/schemas/elastic.py

@@ -2,7 +2,6 @@
 
 from copy import deepcopy
 from enum import Enum
-from typing import Optional
 
 import numpy as np
 from emmet.core.math import Matrix3D, MatrixVoigt
@@ -21,66 +20,67 @@
 from typing_extensions import Self
 
 from atomate2 import SETTINGS
+from atomate2.common.utils import _recursive_to_list
 
 
 class DerivedProperties(BaseModel):
     """Properties derived from an elastic tensor."""
 
-    k_voigt: Optional[float] = Field(
+    k_voigt: float | None = Field(
         None, description="Voigt average of the bulk modulus."
     )
-    k_reuss: Optional[float] = Field(
+    k_reuss: float | None = Field(
         None, description="Reuss average of the bulk modulus."
     )
-    k_vrh: Optional[float] = Field(
+    k_vrh: float | None = Field(
         None, description="Voigt-Reuss-Hill average of the bulk modulus."
     )
-    g_voigt: Optional[float] = Field(
+    g_voigt: float | None = Field(
         None, description="Voigt average of the shear modulus."
     )
-    g_reuss: Optional[float] = Field(
+    g_reuss: float | None = Field(
         None, description="Reuss average of the shear modulus."
     )
-    g_vrh: Optional[float] = Field(
+    g_vrh: float | None = Field(
         None, description="Voigt-Reuss-Hill average of the shear modulus."
     )
-    universal_anisotropy: Optional[float] = Field(
+    universal_anisotropy: float | None = Field(
         None, description="Universal elastic anisotropy."
     )
-    homogeneous_poisson: Optional[float] = Field(
+    homogeneous_poisson: float | None = Field(
         None, description="Homogeneous poisson ratio."
     )
-    y_mod: Optional[float] = Field(
+    y_mod: float | None = Field(
         None,
         description="Young's modulus (SI units) from the Voight-Reuss-Hill averages of "
         "the bulk and shear moduli.",
     )
-    trans_v: Optional[float] = Field(
+    trans_v: float | None = Field(
         None,
         description="Transverse sound velocity (SI units) obtained from the "
         "Voigt-Reuss-Hill average bulk modulus.",
     )
-    long_v: Optional[float] = Field(
+    long_v: float | None = Field(
         None,
         description="Longitudinal sound velocity (SI units) obtained from the "
         "Voigt-Reuss-Hill average bulk modulus.",
     )
-    snyder_ac: Optional[float] = Field(
+    snyder_ac: float | None = Field(
         None, description="Synder's acoustic sound velocity (SI units)."
     )
-    snyder_opt: Optional[float] = Field(
+    snyder_opt: float | None = Field(
         None, description="Synder's optical sound velocity (SI units)."
     )
-    snyder_total: Optional[float] = Field(
+    snyder_total: float | None = Field(
         None, description="Synder's total sound velocity (SI units)."
     )
-    clark_thermalcond: Optional[float] = Field(
+    clark_thermalcond: float | None = Field(
         None, description="Clarke's thermal conductivity (SI units)."
     )
-    cahill_thermalcond: Optional[float] = Field(
+    cahill_thermalcond: float | None = Field(
         None, description="Cahill's thermal conductivity (SI units)."
     )
-    debye_temperature: Optional[float] = Field(
+    debye_temperature: float | None = Field(
         None,
         description="Debye temperature from longitudinal and transverse sound "
         "velocities (SI units).",
@@ -90,34 +90,34 @@ class DerivedProperties(BaseModel):
 class FittingData(BaseModel):
     """Data used to fit elastic tensors."""
 
-    cauchy_stresses: Optional[list[Matrix3D]] = Field(
+    cauchy_stresses: list[Matrix3D] | None = Field(
         None, description="The Cauchy stresses used to fit the elastic tensor."
     )
-    strains: Optional[list[Matrix3D]] = Field(
+    strains: list[Matrix3D] | None = Field(
         None, description="The strains used to fit the elastic tensor."
     )
-    pk_stresses: Optional[list[Matrix3D]] = Field(
+    pk_stresses: list[Matrix3D] | None = Field(
         None, description="The Piola-Kirchoff stresses used to fit the elastic tensor."
     )
-    deformations: Optional[list[Matrix3D]] = Field(
+    deformations: list[Matrix3D] | None = Field(
         None, description="The deformations corresponding to each strain state."
     )
-    uuids: Optional[list[str]] = Field(
+    uuids: list[str] | None = Field(
         None, description="The uuids of the deformation jobs."
     )
-    job_dirs: Optional[list[Optional[str]]] = Field(
+    job_dirs: list[str | None] | None = Field(
         None, description="The directories where the deformation jobs were run."
     )
-    failed_uuids: Optional[list[str]] = Field(
+    failed_uuids: list[str] | None = Field(
         None, description="The uuids of perturbations that were not completed"
     )
 
 
 class ElasticTensorDocument(BaseModel):
     """Raw and standardized elastic tensors."""
 
-    raw: Optional[MatrixVoigt] = Field(None, description="Raw elastic tensor.")
-    ieee_format: Optional[MatrixVoigt] = Field(
+    raw: MatrixVoigt | list | None = Field(None, description="Raw elastic tensor.")
+    ieee_format: MatrixVoigt | list | None = Field(
         None, description="Elastic tensor in IEEE format."
     )
 
@@ -131,28 +131,28 @@ class ElasticWarnings(Enum):
 class ElasticDocument(StructureMetadata):
     """Document containing elastic tensor information and related properties."""
 
-    structure: Optional[Structure] = Field(
+    structure: Structure | None = Field(
         None, description="The structure for which the elastic data is calculated."
     )
-    elastic_tensor: Optional[ElasticTensorDocument] = Field(
+    elastic_tensor: ElasticTensorDocument | None = Field(
         None, description="Fitted elastic tensor."
     )
-    eq_stress: Optional[Matrix3D] = Field(
+    eq_stress: Matrix3D | None = Field(
         None, description="The equilibrium stress of the structure."
     )
-    derived_properties: Optional[DerivedProperties] = Field(
+    derived_properties: DerivedProperties | None = Field(
         None, description="Properties derived from the elastic tensor."
     )
-    fitting_data: Optional[FittingData] = Field(
+    fitting_data: FittingData | None = Field(
         None, description="Data used to fit the elastic tensor."
     )
-    fitting_method: Optional[str] = Field(
+    fitting_method: str | None = Field(
         None, description="Method used to fit the elastic tensor."
     )
-    order: Optional[int] = Field(
+    order: int | None = Field(
         None, description="Order of the expansion of the elastic tensor."
     )
-    warnings: Optional[list[str]] = Field(None, description="Warnings.")
+    warnings: list[str] | None = Field(None, description="Warnings.")
 
     @classmethod
     def from_stresses(
@@ -163,8 +163,8 @@ def from_stresses(
         uuids: list[str],
         job_dirs: list[str],
         fitting_method: str = SETTINGS.ELASTIC_FITTING_METHOD,
-        order: Optional[int] = None,
-        equilibrium_stress: Optional[Matrix3D] = None,
+        order: int | None = None,
+        equilibrium_stress: Matrix3D | None = None,
         symprec: float = SETTINGS.SYMPREC,
         allow_elastically_unstable_structs: bool = True,
         failed_uuids: list[str] = None,
@@ -264,7 +264,8 @@ def from_stresses(
             fitting_method=fitting_method,
             order=order,
             elastic_tensor=ElasticTensorDocument(
-                raw=result.voigt.tolist(), ieee_format=ieee.voigt.tolist()
+                raw=_recursive_to_list(result.voigt),
+                ieee_format=_recursive_to_list(ieee.voigt),
             ),
             fitting_data=FittingData(
                 cauchy_stresses=[s.tolist() for s in stresses],

src/atomate2/common/utils.py

@@ -209,3 +209,15 @@ def _recursive_get_dir_names(jobs: list, dir_names: list) -> None:
             _recursive_get_dir_names(sub_jobs, dir_names)
         else:
             dir_names.append(a_job.output.dir_name)
+
+
+def _recursive_to_list(voigt_data: Any) -> Any:
+    """Recursively convert tensor-like data to nested lists.
+
+    Useful for converting numpy or torch arrays to lists.
+    """
+    if isinstance(voigt_data, list):
+        return [_recursive_to_list(item) for item in voigt_data]
+    if hasattr(voigt_data, "tolist"):
+        return voigt_data.tolist()
+    return voigt_data

tests/common/test_utils.py

@@ -0,0 +1,25 @@
+"""Test utility functions.
+
+TODO: there should be tests for the other utility functions.
+Not high priority but good for long term.
+"""
+
+import numpy as np
+
+from atomate2.common.utils import _recursive_to_list
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+def test_to_list():
+    as_list = [[1, 2, 3, 4, 5], [10, 9, 8, 7, 6], [3, 5, 7, 9, 11]]
+    arr = np.array(as_list)
+
+    assert _recursive_to_list(arr) == as_list
+    if torch is not None:
+        assert _recursive_to_list(torch.from_numpy(arr)) == as_list
+    for obj in (None, 1, 1.5):
+        assert _recursive_to_list(obj) == obj