add phonon tests

Author: esoteric-ephemera

mp_api/client/routes/materials/phonon.py

@@ -1,13 +1,18 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from typing import TYPE_CHECKING
 
 import numpy as np
 from emmet.core.phonon import PhononBS, PhononBSDOSDoc, PhononDOS
 
 from mp_api.client.core import BaseRester, MPRestError
 from mp_api.client.core.utils import validate_ids
 
+if TYPE_CHECKING:
+    from typing import Any
+    from emmet.core.math import Matrix3D
+
 
 class PhononRester(BaseRester):
     suffix = "materials/phonon"
@@ -63,7 +68,7 @@ def search(
             **query_params,
         )
 
-    def get_bandstructure_from_material_id(self, material_id: str, phonon_method: str):
+    def get_bandstructure_from_material_id(self, material_id: str, phonon_method: str) -> PhononBS | dict[str,Any]:
         """Get the phonon band structure pymatgen object associated with a given material ID and phonon method.
 
         Arguments:
@@ -73,10 +78,13 @@ def get_bandstructure_from_material_id(self, material_id: str, phonon_method: st
         Returns:
             bandstructure (PhononBS): PhononBS object
         """
-        result = self._query_open_data(
-            bucket="materialsproject-parsed",
-            key=f"ph-bandstructures/{phonon_method}/{material_id}.json.gz",
-        )[0]
+        try:
+            result = self._query_open_data(
+                bucket="materialsproject-parsed",
+                key=f"ph-bandstructures/{phonon_method}/{material_id}.json.gz",
+            )[0]
+        except OSError:
+            result = None
 
         if not result or not result[0]:
             raise MPRestError("No object found")
@@ -86,7 +94,7 @@ def get_bandstructure_from_material_id(self, material_id: str, phonon_method: st
 
         return result[0]
 
-    def get_dos_from_material_id(self, material_id: str, phonon_method: str):
+    def get_dos_from_material_id(self, material_id: str, phonon_method: str) -> PhononDOS | dict[str,Any]:
         """Get the phonon dos pymatgen object associated with a given material ID and phonon method.
 
         Arguments:
@@ -96,10 +104,13 @@ def get_dos_from_material_id(self, material_id: str, phonon_method: str):
         Returns:
             dos (PhononDOS): PhononDOS object
         """
-        result = self._query_open_data(
-            bucket="materialsproject-parsed",
-            key=f"ph-dos/{phonon_method}/{material_id}.json.gz",
-        )[0]
+        try:
+            result = self._query_open_data(
+                bucket="materialsproject-parsed",
+                key=f"ph-dos/{phonon_method}/{material_id}.json.gz",
+            )[0]
+        except OSError:
+            result = None
 
         if not result or not result[0]:
             raise MPRestError("No object found")
@@ -109,7 +120,7 @@ def get_dos_from_material_id(self, material_id: str, phonon_method: str):
 
         return result[0]
 
-    def get_forceconstants_from_material_id(self, material_id: str):
+    def get_forceconstants_from_material_id(self, material_id: str) -> list[list[Matrix3D]]:
         """Get the force constants associated with a given material ID.
 
         Arguments:
@@ -118,10 +129,13 @@ def get_forceconstants_from_material_id(self, material_id: str):
         Returns:
             force constants (list[list[Matrix3D]]): PhononDOS object
         """
-        result = self._query_open_data(
-            bucket="materialsproject-parsed",
-            key=f"ph-force-constants/{material_id}.json.gz",
-        )[0]
+        try:
+            result = self._query_open_data(
+                bucket="materialsproject-parsed",
+                key=f"ph-force-constants/{material_id}.json.gz",
+            )[0]
+        except OSError:
+            result = None
 
         if not result or not result[0]:
             raise MPRestError("No object found")

tests/materials/test_phonon.py

@@ -0,0 +1,79 @@
+
+import os
+
+import numpy as np
+import pytest
+
+from emmet.core.phonon import PhononBS, PhononDOS
+
+from mp_api.client.core import MPRestError
+from mp_api.client.routes.materials.phonon import PhononRester
+
+from core_function import client_search_testing
+
+@pytest.mark.skipif(os.getenv("MP_API_KEY") is None, reason="No API key found.")
+def test_phonon_search():
+    client_search_testing(
+        search_method=PhononRester().search,
+        excluded_params=[
+            "num_chunks",
+            "chunk_size",
+            "all_fields",
+            "fields",
+        ],
+        alt_name_dict={
+            "material_ids": "material_id",
+        },
+        custom_field_tests = {
+            "material_ids": ["mp-149","mp-13"],
+            "material_ids": "mp-149",
+            "phonon_method": "dfpt",
+        },
+        sub_doc_fields=[],
+    )
+
+@pytest.mark.skipif(os.getenv("MP_API_KEY") is None, reason="No API key found.")
+@pytest.mark.parametrize("use_document_model",[True,False])
+def test_phonon_get_methods(use_document_model):
+    
+    rester = PhononRester(use_document_model=use_document_model)
+
+    # TODO: update when there is force constant data
+    for func_name, schema in {
+        "bandstructure" : PhononBS,
+        "dos": PhononDOS,
+        # "forceconstants": list
+    }.items():
+        search_method = getattr(
+            rester,
+            f"get_{func_name}_from_material_id",
+        )
+        assert isinstance(
+            search_method("mp-149","dfpt"),
+            schema if use_document_model else dict
+        )
+
+        with pytest.raises(MPRestError,match="No object found"):
+            _ = search_method("mp-0","dfpt")
+
+@pytest.mark.skipif(os.getenv("MP_API_KEY") is None, reason="No API key found.")
+@pytest.mark.parametrize("use_document_model",[True,False])
+def test_phonon_thermo(use_document_model):
+
+    with pytest.raises(MPRestError,match="No phonon document found"):
+        _ = PhononRester(
+            use_document_model=use_document_model
+        ).compute_thermo_quantities("mp-0","dfpt")
+
+    thermo_props = PhononRester(
+        use_document_model=use_document_model
+    ).compute_thermo_quantities("mp-149","dfpt")
+    
+    # Default set in the method
+    num_vals = 100
+    
+    assert all(
+        isinstance(v, np.ndarray if k == "temperature" else list)
+        and len(v) == num_vals
+        for k, v in thermo_props.items()
+    )

tests/materials/test_similarity.py

@@ -28,7 +28,7 @@ def test_struct():
     return Structure.from_str(poscar,fmt="poscar")
 
 @pytest.mark.skipif(os.getenv("MP_API_KEY") is None, reason="No API key found.")
-def test_client():
+def test_similarity_search():
     client_search_testing(
         search_method=SimilarityRester().search,
         excluded_params=[