Fix edge and spectrum type in xas client (#619)

materialsproject · Jun 6, 2022 · 3ea2f4b · 3ea2f4b
1 parent 250f51c
commit 3ea2f4b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/src/mp_api/routes/xas.py b/src/mp_api/routes/xas.py
@@ -1,6 +1,6 @@
 from typing import List, Optional, Union
 
-from emmet.core.xas import Edge, XASDoc
+from emmet.core.xas import Edge, XASDoc, Type
 from mp_api.core.client import BaseRester
 from mp_api.core.utils import validate_ids
 from pymatgen.core.periodic_table import Element
@@ -36,6 +36,7 @@ def search(
         chemsys: Optional[Union[str, List[str]]] = None,
         elements: Optional[List[str]] = None,
         material_ids: Optional[List[str]] = None,
+        spectrum_type: Optional[Type] = None,
         sort_fields: Optional[List[str]] = None,
         num_chunks: Optional[int] = None,
         chunk_size: int = 1000,
@@ -47,12 +48,14 @@ def search(
 
         Arguments:
             edge (Edge): The absorption edge (e.g. K, L2, L3, L2,3).
+            absorbing_element (Element): The absorbing element.
             formula (str): A formula including anonomyzed formula
                 or wild cards (e.g., Fe2O3, ABO3, Si*).
             chemsys (str, List[str]): A chemical system or list of chemical systems
                 (e.g., Li-Fe-O, Si-*, [Si-O, Li-Fe-P]).
             elements (List[str]): A list of elements.
             material_ids (List[str]): List of Materials Project IDs to return data for.
+            spectrum_type (Type): Spectrum type (e.g. EXAFS, XAFS, or XANES).
             sort_fields (List[str]): Fields used to sort results. Prefix with '-' to sort in descending order.
             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.
@@ -66,10 +69,19 @@ def search(
         query_params = {}
 
         if edge:
-            query_params.update({"edge": str(edge.value)})
+            query_params.update({"edge": edge})
 
         if absorbing_element:
-            query_params.update({"absorbing_element": str(absorbing_element.symbol)})
+            query_params.update(
+                {
+                    "absorbing_element": str(absorbing_element.symbol)
+                    if type(absorbing_element) == Element
+                    else absorbing_element
+                }
+            )
+
+        if spectrum_type:
+            query_params.update({"spectrum_type": spectrum_type})
 
         if formula:
             query_params.update({"formula": formula})

diff --git a/tests/test_xas.py b/tests/test_xas.py
@@ -1,7 +1,7 @@
 import os
 import pytest
 from mp_api.routes.xas import XASRester
-from emmet.core.xas import Edge
+from emmet.core.xas import Edge, Type
 from pymatgen.core.periodic_table import Element
 
 import typing
@@ -32,7 +32,8 @@ def rester():
 }  # type: dict
 
 custom_field_tests = {
-    "edge": Edge.K,
+    "edge": Edge.L2_3,
+    "spectrum_type": Type.EXAFS,
     "absorbing_element": Element("Ce"),
     "required_elements": [Element("Ce")],
     "formula": "Ce(WO4)2",