Add numeric query op to materials

materialsproject · May 27, 2021 · 5b3cdd5 · 5b3cdd5
1 parent fabb1a4
commit 5b3cdd5
Showing 1 changed file with 59 additions and 16 deletions.
diff --git a/src/mp_api/routes/materials/resources.py b/src/mp_api/routes/materials/resources.py
@@ -1,11 +1,18 @@
 from fastapi import HTTPException
+from maggma.api.query_operator.dynamic import StringQueryOperator
 
 from maggma.api.resource.read_resource import ReadOnlyResource
 
 
 from emmet.core.material import MaterialsDoc
 
-from maggma.api.query_operator import PaginationQuery, SparseFieldsQuery, SortQuery, VersionQuery
+from maggma.api.query_operator import (
+    PaginationQuery,
+    SparseFieldsQuery,
+    SortQuery,
+    VersionQuery,
+    NumericQuery,
+)
 
 from mp_api.core.settings import MAPISettings
 
@@ -50,7 +57,12 @@ async def get_versions():
 
             col_names = db.list_collection_names()
 
-            d = [name.replace("_", ".")[15:] for name in col_names if "materials" in name if name != "materials.core"]
+            d = [
+                name.replace("_", ".")[15:]
+                for name in col_names
+                if "materials" in name
+                if name != "materials.core"
+            ]
 
             response = {"data": d}
 
@@ -67,16 +79,23 @@ def custom_findstructure_prep(self):
         model_name = self.model.__name__
 
         async def find_structure(
-            structure: Structure = Body(..., title="Pymatgen structure object to query with",),
-            ltol: float = Query(0.2, title="Fractional length tolerance. Default is 0.2.",),
+            structure: Structure = Body(
+                ..., title="Pymatgen structure object to query with",
+            ),
+            ltol: float = Query(
+                0.2, title="Fractional length tolerance. Default is 0.2.",
+            ),
             stol: float = Query(
                 0.3,
                 title="Site tolerance. Defined as the fraction of the average free \
                     length per atom := ( V / Nsites ) ** (1/3). Default is 0.3.",
             ),
-            angle_tol: float = Query(5, title="Angle tolerance in degrees. Default is 5 degrees.",),
+            angle_tol: float = Query(
+                5, title="Angle tolerance in degrees. Default is 5 degrees.",
+            ),
             limit: int = Query(
-                1, title="Maximum number of matches to show. Defaults to 1, only showing the best match.",
+                1,
+                title="Maximum number of matches to show. Defaults to 1, only showing the best match.",
             ),
         ):
             """
@@ -90,7 +109,8 @@ async def find_structure(
                 s = Structure.from_dict(structure.dict())
             except Exception:
                 raise HTTPException(
-                    status_code=404, detail="Body cannot be converted to a pymatgen structure object.",
+                    status_code=404,
+                    detail="Body cannot be converted to a pymatgen structure object.",
                 )
 
             m = StructureMatcher(
@@ -109,7 +129,9 @@ async def find_structure(
 
             matches = []
 
-            for r in self.store.query(criteria=crit, properties=["structure", "task_id"]):
+            for r in self.store.query(
+                criteria=crit, properties=["structure", "task_id"]
+            ):
 
                 s2 = Structure.from_dict(r["structure"])
                 matched = m.fit(s, s2)
@@ -127,7 +149,11 @@ async def find_structure(
 
             response = {
                 "data": sorted(
-                    matches[:limit], key=lambda x: (x["normalized_rms_displacement"], x["max_distance_paired_sites"],),
+                    matches[:limit],
+                    key=lambda x: (
+                        x["normalized_rms_displacement"],
+                        x["max_distance_paired_sites"],
+                    ),
                 )
             }
 
@@ -142,8 +168,12 @@ async def find_structure(
 
     def custom_autocomplete_prep(self):
         async def formula_autocomplete(
-            text: str = Query(..., description="Text to run against formula autocomplete",),
-            limit: int = Query(10, description="Maximum number of matches to show. Defaults to 10",),
+            text: str = Query(
+                ..., description="Text to run against formula autocomplete",
+            ),
+            limit: int = Query(
+                10, description="Maximum number of matches to show. Defaults to 10",
+            ),
         ):
             store = formula_autocomplete_store
 
@@ -191,7 +221,12 @@ async def formula_autocomplete(
                             "length": {"$strLenCP": "$formula_pretty"},
                         }
                     },
-                    {"$match": {"length": {"$gte": len(final_terms[0])}, "elements": {"$all": eles}}},
+                    {
+                        "$match": {
+                            "length": {"$gte": len(final_terms[0])},
+                            "elements": {"$all": eles},
+                        }
+                    },
                     {"$limit": limit},
                     {"$sort": {"length": 1}},
                     {"$project": {"elements": 0, "length": 0}},
@@ -205,7 +240,8 @@ async def formula_autocomplete(
 
             except Exception:
                 raise HTTPException(
-                    status_code=404, detail="Cannot autocomplete with provided formula.",
+                    status_code=404,
+                    detail="Cannot autocomplete with provided formula.",
                 )
 
             return response
@@ -227,13 +263,20 @@ async def formula_autocomplete(
             MultiTaskIDQuery(),
             SymmetryQuery(),
             DeprecationQuery(),
-            MinMaxQuery(),
+            NumericQuery(model=MaterialsDoc),
             SortQuery(),
             PaginationQuery(),
-            SparseFieldsQuery(MaterialsDoc, default_fields=["material_id", "formula_pretty", "last_updated"],),
+            SparseFieldsQuery(
+                MaterialsDoc,
+                default_fields=["material_id", "formula_pretty", "last_updated"],
+            ),
         ],
         tags=["Materials"],
-        custom_endpoint_funcs=[custom_version_prep, custom_findstructure_prep, custom_autocomplete_prep],
+        custom_endpoint_funcs=[
+            custom_version_prep,
+            custom_findstructure_prep,
+            custom_autocomplete_prep,
+        ],
     )
 
     return resource