Grain boundary now uses numeric query op

materialsproject · May 28, 2021 · 32f9bd9 · 32f9bd9
1 parent 17bb254
commit 32f9bd9
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 54 deletions.
diff --git a/src/mp_api/routes/grain_boundary/query_operators.py b/src/mp_api/routes/grain_boundary/query_operators.py
@@ -8,64 +8,18 @@
 from mp_api.routes.grain_boundary.models import GBTypeEnum
 
 
-class GBEnergyQuery(QueryOperator):
-    """
-    Method to generate a query for energy values associated with grain boundary data
-    """
-
-    def query(
-        self,
-        gb_energy_max: Optional[float] = Query(
-            None, description="Maximum value for the grain boundary energy in J/m^2.",
-        ),
-        gb_energy_min: Optional[float] = Query(
-            None, description="Minimum value for the grain boundary energy in J/m^2.",
-        ),
-        w_sep_energy_max: Optional[float] = Query(
-            None, description="Maximum value for the work of separation energy in J/m^2.",
-        ),
-        w_sep_energy_min: Optional[float] = Query(
-            None, description="Minimum value for work of separation energy in J/m^2.",
-        ),
-    ) -> STORE_PARAMS:
-
-        crit = defaultdict(dict)  # type: dict
-
-        d = {
-            "gb_energy": [gb_energy_min, gb_energy_max],
-            "w_sep": [w_sep_energy_min, w_sep_energy_max],
-        }
-
-        for entry in d:
-            if d[entry][0]:
-                crit[entry]["$gte"] = d[entry][0]
-
-            if d[entry][1]:
-                crit[entry]["$lte"] = d[entry][1]
-
-        return {"criteria": crit}
-
-    def ensure_indexes(self):
-        keys = ["gb_energy", "w_sep"]
-        return [(key, False) for key in keys]
-
-
 class GBStructureQuery(QueryOperator):
     """
     Method to generate a query for structure related data associated with grain boundary entries
     """
 
     def query(
         self,
-        rotation_angle_max: Optional[float] = Query(
-            None, description="Maximum value for the rotation angle in degrees.",
-        ),
-        rotation_angle_min: Optional[float] = Query(
-            None, description="Minimum value for the rotation angle in degrees.",
-        ),
         sigma: Optional[float] = Query(None, description="Value of sigma.",),
         type: Optional[GBTypeEnum] = Query(None, description="Grain boundary type.",),
-        chemsys: Optional[str] = Query(None, description="Dash-delimited string of elements in the material.",),
+        chemsys: Optional[str] = Query(
+            None, description="Dash-delimited string of elements in the material.",
+        ),
     ) -> STORE_PARAMS:
 
         crit = defaultdict(dict)  # type: dict
@@ -106,7 +60,8 @@ class GBTaskIDQuery(QueryOperator):
     def query(
         self,
         task_ids: Optional[str] = Query(
-            None, description="Comma-separated list of Materials Project IDs to query on.",
+            None,
+            description="Comma-separated list of Materials Project IDs to query on.",
         ),
     ) -> STORE_PARAMS:
 

diff --git a/src/mp_api/routes/grain_boundary/resources.py b/src/mp_api/routes/grain_boundary/resources.py
@@ -1,8 +1,16 @@
 from maggma.api.resource import ReadOnlyResource
 from mp_api.routes.grain_boundary.models import GrainBoundaryDoc
 
-from mp_api.routes.grain_boundary.query_operators import GBEnergyQuery, GBStructureQuery, GBTaskIDQuery
-from maggma.api.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
+from mp_api.routes.grain_boundary.query_operators import (
+    GBStructureQuery,
+    GBTaskIDQuery,
+)
+from maggma.api.query_operator import (
+    PaginationQuery,
+    SortQuery,
+    SparseFieldsQuery,
+    NumericQuery,
+)
 
 
 def gb_resource(gb_store):
@@ -11,11 +19,15 @@ def gb_resource(gb_store):
         GrainBoundaryDoc,
         query_operators=[
             GBTaskIDQuery(),
-            GBEnergyQuery(),
+            NumericQuery(
+                model=GrainBoundaryDoc, excluded_fields=["rotation_axis", "gb_plane"]
+            ),
             GBStructureQuery(),
             SortQuery(),
             PaginationQuery(),
-            SparseFieldsQuery(GrainBoundaryDoc, default_fields=["task_id", "last_updated"]),
+            SparseFieldsQuery(
+                GrainBoundaryDoc, default_fields=["task_id", "last_updated"]
+            ),
         ],
         tags=["Grain Boundaries"],
         enable_get_by_key=False,