Merge pull request #514 from gpetretto/devel

Add StructureMetadata as baseclass for output documents

src/atomate2/common/schemas/elastic.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 from emmet.core.math import Matrix3D, MatrixVoigt
+from emmet.core.structure import StructureMetadata
 from pydantic import BaseModel, Field
 from pymatgen.analysis.elasticity import (
     Deformation,
@@ -108,7 +109,7 @@ class ElasticTensorDocument(BaseModel):
     ieee_format: MatrixVoigt = Field(None, description="Elastic tensor in IEEE format.")
 
 
-class ElasticDocument(BaseModel):
+class ElasticDocument(StructureMetadata):
     """Document containing elastic tensor information and related properties."""
 
     structure: Structure = Field(
@@ -123,10 +124,6 @@ class ElasticDocument(BaseModel):
     derived_properties: DerivedProperties = Field(
         None, description="Properties derived from the elastic tensor."
     )
-    formula_pretty: str = Field(
-        None,
-        description="Cleaned representation of the formula",
-    )
     fitting_data: FittingData = Field(
         None, description="Data used to fit the elastic tensor."
     )
@@ -231,11 +228,11 @@ def from_stresses(
 
         eq_stress = eq_stress.tolist() if eq_stress is not None else eq_stress
 
-        return cls(
+        return cls.from_structure(
             structure=structure,
+            meta_structure=structure,
             eq_stress=eq_stress,
             derived_properties=derived_properties,
-            formula_pretty=structure.composition.reduced_formula,
             fitting_method=fitting_method,
             order=order,
             elastic_tensor=ElasticTensorDocument(

src/atomate2/common/schemas/phonons.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 from emmet.core.math import Matrix3D
+from emmet.core.structure import StructureMetadata
 from monty.json import MSONable
 from phonopy import Phonopy
 from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections
@@ -105,7 +106,7 @@ class PhononJobDirs(BaseModel):
     )
 
 
-class PhononBSDOSDoc(BaseModel):
+class PhononBSDOSDoc(StructureMetadata):
     """Collection of all data produced by the phonon workflow."""
 
     structure: Structure = Field(
@@ -442,8 +443,9 @@ def from_forces_born(
             total_dft_energy / formula_units if total_dft_energy is not None else None
         )
 
-        return cls(
+        return cls.from_structure(
             structure=structure,
+            meta_structure=structure,
             phonon_bandstructure=bs_symm_line,
             phonon_dos=dos,
             free_energies=free_energies,

src/atomate2/lobster/schemas.py

@@ -6,6 +6,7 @@
 from typing import Any, List, Optional, Union
 
 import numpy as np
+from emmet.core.structure import StructureMetadata
 from monty.dev import requires
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
@@ -334,7 +335,7 @@ class StrongestBonds(BaseModel):
     )
 
 
-class LobsterTaskDocument(BaseModel):
+class LobsterTaskDocument(StructureMetadata):
     """Definition of LOBSTER task document."""
 
     structure: Structure = Field(None, description="The structure used in this task")
@@ -609,8 +610,9 @@ def from_directory(
             for spin, value in band_overlaps_obj.bandoverlapsdict.items():
                 band_overlaps[str(spin.value)] = value
 
-        doc = cls(
+        doc = cls.from_structure(
             structure=struct,
+            meta_structure=struct,
             dir_name=dir_name,
             lobsterin=lobster_in,
             lobsterout=lobster_out,

src/atomate2/vasp/schemas/elph.py

@@ -4,6 +4,7 @@
 from typing import Dict, List, Tuple
 
 import numpy as np
+from emmet.core.structure import StructureMetadata
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
 from pymatgen.electronic_structure.bandstructure import BandStructure
@@ -71,18 +72,14 @@ class RawElectronicData(BaseModel):
     )
 
 
-class ElectronPhononRenormalisationDoc(BaseModel):
+class ElectronPhononRenormalisationDoc(StructureMetadata):
     """Electron-phonon band gap renormalisation document."""
 
     structure: Structure = Field(
         None,
         description="The primitive structure for which the electron-phonon was"
         " calculated",
     )
-    formula_pretty: str = Field(
-        None,
-        description="Cleaned representation of the formula",
-    )
     temperatures: List[float] = Field(
         None, description="Temperatures at which electron-phonon coupling was obtained"
     )
@@ -214,9 +211,9 @@ def from_band_structures(
         band_gaps = cbms - vbms
         delta_band_gaps = band_gaps - bulk_band_gap
 
-        return cls(
+        return cls.from_structure(
             structure=original_structure,
-            formula_pretty=original_structure.composition.reduced_formula,
+            meta_structure=original_structure,
             temperatures=temperatures,
             band_gaps=band_gaps.tolist(),
             vbms=vbms.tolist(),

tests/vasp/flows/test_elastic.py

@@ -63,3 +63,4 @@ def test_elastic(mock_vasp, clean_dir, si_structure):
         ],
         atol=1e-3,
     )
+    assert elastic_output.chemsys == "Si"

tests/vasp/flows/test_elph.py

@@ -69,3 +69,4 @@ def test_elph_renormalisation(mock_vasp, clean_dir, si_structure):
         -0.488900000000001,
         -0.48850000000000104,
     }
+    assert renorm_output.chemsys == "Si"

tests/vasp/flows/test_phonons.py

@@ -145,6 +145,7 @@ def test_phonon_wf_only_displacements3(mock_vasp, clean_dir):
             11255.660261586278,
         ],
     )
+    assert responses[job.jobs[-1].uuid][1].output.chemsys == "Si"
 
 
 # structure will be kept in the format that was transferred

tests/vasp/lobster/schemas/test_lobster.py

@@ -113,6 +113,8 @@ def test_LobsterTaskDocument(lobster_test_dir):
     )
     assert len(doc.band_overlaps["1"]) + len(doc.band_overlaps["-1"]) == 12
 
+    assert doc.chemsys == "As-Ga"
+
     doc2 = LobsterTaskDocument.from_directory(
         dir_name=lobster_test_dir / "lobsteroutputs/mp-754354", save_cohp_plots=False
     )