re-enable and tweak electronic structure tests

Author: esoteric-ephemera

mp_api/client/routes/materials/electronic_structure.py

@@ -75,7 +75,7 @@ def search(
             num_chunks (int): Maximum number of chunks of data to yield. None will yield all possible.
             chunk_size (int): Number of data entries per chunk.
             all_fields (bool): Whether to return all fields in the document. Defaults to True.
-            fields (List[str]): List of fields in EOSDoc to return data for.
+            fields (List[str]): List of fields in ElectronicStructureDoc to return data for.
                 Default is material_id and last_updated if all_fields is False.
 
         Returns:
@@ -186,8 +186,8 @@ def search(
         efermi: tuple[float, float] | None = None,
         is_gap_direct: bool | None = None,
         is_metal: bool | None = None,
-        magnetic_ordering: Ordering | None = None,
-        path_type: BSPathType = BSPathType.setyawan_curtarolo,
+        magnetic_ordering: Ordering | str | None = None,
+        path_type: BSPathType | str = BSPathType.setyawan_curtarolo,
         num_chunks: int | None = None,
         chunk_size: int = 1000,
         all_fields: bool = True,
@@ -200,8 +200,8 @@ def search(
             efermi (Tuple[float,float]): Minimum and maximum fermi energy in eV to consider.
             is_gap_direct (bool): Whether the material has a direct band gap.
             is_metal (bool): Whether the material is considered a metal.
-            magnetic_ordering (Ordering): Magnetic ordering of the material.
-            path_type (BSPathType): k-path selection convention for the band structure.
+            magnetic_ordering (Ordering or str): Magnetic ordering of the material.
+            path_type (BSPathType or str): k-path selection convention for the band structure.
             num_chunks (int): Maximum number of chunks of data to yield. None will yield all possible.
             chunk_size (int): Number of data entries per chunk.
             all_fields (bool): Whether to return all fields in the document. Defaults to True.
@@ -213,7 +213,9 @@ def search(
         """
         query_params = defaultdict(dict)  # type: dict
 
-        query_params["path_type"] = path_type.value
+        query_params["path_type"] = (
+            BSPathType[path_type] if isinstance(path_type, str) else path_type
+        ).value
 
         if band_gap:
             query_params.update(
@@ -224,7 +226,15 @@ def search(
             query_params.update({"efermi_min": efermi[0], "efermi_max": efermi[1]})
 
         if magnetic_ordering:
-            query_params.update({"magnetic_ordering": magnetic_ordering.value})
+            query_params.update(
+                {
+                    "magnetic_ordering": (
+                        Ordering(magnetic_ordering)
+                        if isinstance(magnetic_ordering, str)
+                        else magnetic_ordering
+                    ).value
+                }
+            )
 
         if is_gap_direct is not None:
             query_params.update({"is_gap_direct": is_gap_direct})
@@ -351,11 +361,11 @@ def search(
         self,
         band_gap: tuple[float, float] | None = None,
         efermi: tuple[float, float] | None = None,
-        element: Element | None = None,
-        magnetic_ordering: Ordering | None = None,
-        orbital: OrbitalType | None = None,
-        projection_type: DOSProjectionType = DOSProjectionType.total,
-        spin: Spin = Spin.up,
+        element: Element | str | None = None,
+        magnetic_ordering: Ordering | str | None = None,
+        orbital: OrbitalType | str | None = None,
+        projection_type: DOSProjectionType | str = DOSProjectionType.total,
+        spin: Spin | str = Spin.up,
         num_chunks: int | None = None,
         chunk_size: int = 1000,
         all_fields: bool = True,
@@ -366,30 +376,54 @@ def search(
         Arguments:
             band_gap (Tuple[float,float]): Minimum and maximum band gap in eV to consider.
             efermi (Tuple[float,float]): Minimum and maximum fermi energy in eV to consider.
-            element (Element): Element for element-projected dos data.
-            magnetic_ordering (Ordering): Magnetic ordering of the material.
-            orbital (OrbitalType): Orbital for orbital-projected dos data.
-            projection_type (DOSProjectionType): Projection type of dos data. Default is the total dos.
-            spin (Spin): Spin channel of dos data. If non spin-polarized data is stored in Spin.up
+            element (Element or str): Element for element-projected dos data.
+            magnetic_ordering (Ordering or str): Magnetic ordering of the material.
+            orbital (OrbitalType or str): Orbital for orbital-projected dos data.
+            projection_type (DOSProjectionType or str): Projection type of dos data. Default is the total dos.
+            spin (Spin or str): Spin channel of dos data. If non spin-polarized data is stored in Spin.up
             num_chunks (int): Maximum number of chunks of data to yield. None will yield all possible.
             chunk_size (int): Number of data entries per chunk.
             all_fields (bool): Whether to return all fields in the document. Defaults to True.
-            fields (List[str]): List of fields in EOSDoc to return data for.
+            fields (List[str]): List of fields in ElectronicStructureDoc to return data for.
                 Default is material_id and last_updated if all_fields is False.
 
         Returns:
             ([ElectronicStructureDoc]) List of electronic structure documents
         """
         query_params = defaultdict(dict)  # type: dict
 
-        query_params["projection_type"] = projection_type.value
-        query_params["spin"] = spin.value
+        query_params["projection_type"] = (
+            DOSProjectionType[projection_type]
+            if isinstance(projection_type, str)
+            else projection_type
+        ).value
+        query_params["spin"] = (Spin[spin] if isinstance(spin, str) else spin).value
+
+        if (
+            query_params["projection_type"] == DOSProjectionType.elemental.value
+            and element is None
+        ):
+            raise MPRestError(
+                "To query element-projected DOS, you must also specify the `element` onto which the DOS is projected."
+            )
+
+        if (
+            query_params["projection_type"] == DOSProjectionType.orbital.value
+            and orbital is None
+        ):
+            raise MPRestError(
+                "To query orbital-projected DOS, you must also specify the `orbital` character onto which the DOS is projected."
+            )
 
         if element:
-            query_params["element"] = element.value
+            query_params["element"] = (
+                Element[element] if isinstance(element, str) else element
+            ).value
 
         if orbital:
-            query_params["orbital"] = orbital.value
+            query_params["orbital"] = (
+                OrbitalType[orbital] if isinstance(orbital, str) else orbital
+            ).value
 
         if band_gap:
             query_params.update(
@@ -400,7 +434,15 @@ def search(
             query_params.update({"efermi_min": efermi[0], "efermi_max": efermi[1]})
 
         if magnetic_ordering:
-            query_params.update({"magnetic_ordering": magnetic_ordering.value})
+            query_params.update(
+                {
+                    "magnetic_ordering": (
+                        Ordering[magnetic_ordering]
+                        if isinstance(magnetic_ordering, str)
+                        else magnetic_ordering
+                    ).value
+                }
+            )
 
         query_params = {
             entry: query_params[entry]

tests/materials/test_electronic_structure.py

@@ -47,7 +47,6 @@ def es_rester():
 
 
 @requires_api_key
-@pytest.mark.skip(reason="magnetic ordering fields not built correctly")
 def test_es_client(es_rester):
     search_method = es_rester.search
 
@@ -66,51 +65,51 @@ def test_es_client(es_rester):
     "is_gap_direct": True,
     "efermi": (0, 100),
     "band_gap": (0, 5),
+    "path_type": "hinuma",
 }
 
 bs_sub_doc_fields = ["bandstructure"]
 
 bs_alt_name_dict = {}  # type: dict
 
 
-@pytest.fixture
-def bs_rester():
-    rester = BandStructureRester()
-    yield rester
-    rester.session.close()
-
-
 @requires_api_key
-@pytest.mark.skip(reason="magnetic ordering fields not built correctly")
-def test_bs_client(bs_rester):
+def test_bs_client():
     # Get specific search method
-    search_method = bs_rester.search
 
-    # Query fields
-    for param in bs_custom_field_tests:
-        project_field = bs_alt_name_dict.get(param, None)
-        q = {
-            param: bs_custom_field_tests[param],
-            "chunk_size": 1,
-            "num_chunks": 1,
-        }
-        doc = search_method(**q)[0].model_dump()
+    with BandStructureRester() as bs_rester:
+        # Query fields
+        for param in bs_custom_field_tests:
+            project_field = bs_alt_name_dict.get(param, None)
+            q = {
+                param: bs_custom_field_tests[param],
+                "chunk_size": 1,
+                "num_chunks": 1,
+            }
+            doc = bs_rester.search(**q)[0].model_dump()
 
-        for sub_field in bs_sub_doc_fields:
-            if sub_field in doc:
-                doc = doc[sub_field]
+            for sub_field in bs_sub_doc_fields:
+                if sub_field in doc:
+                    doc = doc[sub_field]
 
-        if param != "path_type":
-            doc = doc["setyawan_curtarolo"]
+            if param != "path_type":
+                doc = doc["setyawan_curtarolo"]
 
-        assert doc[project_field if project_field is not None else param] is not None
+                assert (
+                    doc[project_field if project_field is not None else param]
+                    is not None
+                )
 
 
-dos_custom_field_tests = {
-    "magnetic_ordering": Ordering.FM,
-    "efermi": (0, 100),
-    "band_gap": (0, 5),
-}
+dos_custom_field_tests = [
+    {"magnetic_ordering": Ordering.FM},
+    {"efermi": (1, 1.1)},
+    {"band_gap": (8.0, 9.0)},
+    {"projection_type": "elemental", "element": "As"},
+    {
+        "magnetic_ordering": "FM",
+    },
+]
 
 dos_excluded_params = ["orbital", "element"]
 
@@ -119,35 +118,28 @@ def test_bs_client(bs_rester):
 dos_alt_name_dict = {}  # type: dict
 
 
-@pytest.fixture
-def dos_rester():
-    rester = DosRester()
-    yield rester
-    rester.session.close()
-
-
 @requires_api_key
-@pytest.mark.skip(reason="magnetic ordering fields not built correctly")
-def test_dos_client(dos_rester):
-    search_method = dos_rester.search
-
-    # Query fields
-    for param in dos_custom_field_tests:
-        if param not in dos_excluded_params:
-            project_field = dos_alt_name_dict.get(param, None)
+def test_dos_client():
+    with DosRester() as dos_rester:
+        # Query fields
+        for params in dos_custom_field_tests:
+            if any(param in dos_excluded_params for param in params):
+                continue
             q = {
-                param: dos_custom_field_tests[param],
+                **params,
                 "chunk_size": 1,
                 "num_chunks": 1,
             }
-            doc = search_method(**q)[0].model_dump()
+            doc = dos_rester.search(**q)[0].model_dump()
             for sub_field in dos_sub_doc_fields:
                 if sub_field in doc:
                     doc = doc[sub_field]
 
-            if param != "projection_type" and param != "magnetic_ordering":
+            if not any(
+                param in params for param in {"projection_type", "magnetic_ordering"}
+            ):
                 doc = doc["total"]["1"]
 
-            assert (
-                doc[project_field if project_field is not None else param] is not None
+            assert all(
+                doc[dos_alt_name_dict.get(param, param)] is not None for param in params
             )