From a15dce08f61c39d22e520701bdb21c503e7e3cc8 Mon Sep 17 00:00:00 2001
From: Jason Munro <jmunro@lbl.gov>
Date: Thu, 27 May 2021 18:15:37 -0700
Subject: [PATCH] Thermo now uses numeric query op

---
 src/mp_api/routes/thermo/query_operators.py | 87 +++------------------
 src/mp_api/routes/thermo/resources.py       |  8 +-
 2 files changed, 17 insertions(+), 78 deletions(-)

diff --git a/src/mp_api/routes/thermo/query_operators.py b/src/mp_api/routes/thermo/query_operators.py
index 38ec1a91..0e9917d8 100644
--- a/src/mp_api/routes/thermo/query_operators.py
+++ b/src/mp_api/routes/thermo/query_operators.py
@@ -13,11 +13,16 @@ class ThermoChemicalQuery(QueryOperator):
 
     def query(
         self,
-        chemsys: Optional[str] = Query(None, description="Dash-delimited list of elements in the material.",),
+        chemsys: Optional[str] = Query(
+            None, description="Dash-delimited list of elements in the material.",
+        ),
         elements: Optional[str] = Query(
-            None, description="Elements in the material composition as a comma-separated list",
+            None,
+            description="Elements in the material composition as a comma-separated list",
+        ),
+        nelements: Optional[int] = Query(
+            None, description="Number of elements in the material",
         ),
-        nelements: Optional[int] = Query(None, description="Number of elements in the material",),
     ):
 
         crit = {}  # type: dict
@@ -48,7 +53,10 @@ class IsStableQuery(QueryOperator):
     """
 
     def query(
-        self, is_stable: Optional[bool] = Query(None, description="Whether the material is stable."),
+        self,
+        is_stable: Optional[bool] = Query(
+            None, description="Whether the material is stable."
+        ),
     ):
 
         crit = {}
@@ -61,74 +69,3 @@ def query(
     def ensure_indexes(self):
         keys = self._keys_from_query()
         return [(key, False) for key in keys]
-
-
-class ThermoEnergyQuery(QueryOperator):
-    """
-    Method to generate a query for ranges of thermo energy data
-    """
-
-    def query(
-        self,
-        energy_per_atom_max: Optional[float] = Query(
-            None, description="Maximum value for the corrected total energy in eV/atom.",
-        ),
-        energy_per_atom_min: Optional[float] = Query(
-            None, description="Minimum value for the corrected total energy in eV/atom.",
-        ),
-        formation_energy_per_atom_max: Optional[float] = Query(
-            None, description="Maximum value for the formation energy in eV/atom.",
-        ),
-        formation_energy_per_atom_min: Optional[float] = Query(
-            None, description="Minimum value for the formation energy in eV/atom.",
-        ),
-        energy_above_hull_max: Optional[float] = Query(
-            None, description="Maximum value for the energy above the hull in eV/atom.",
-        ),
-        energy_above_hull_min: Optional[float] = Query(
-            None, description="Minimum value for the energy above the hull in eV/atom.",
-        ),
-        equillibrium_reaction_energy_per_atom_max: Optional[float] = Query(
-            None, description="Maximum value for the equilibrium reaction energy in eV/atom.",
-        ),
-        equillibrium_reaction_energy_per_atom_min: Optional[float] = Query(
-            None, description="Minimum value for the equilibrium reaction energy in eV/atom.",
-        ),
-        uncorrected_energy_per_atom_max: Optional[float] = Query(
-            None, description="Maximum value for the uncorrected total energy in eV.",
-        ),
-        uncorrected_energy_per_atom_min: Optional[float] = Query(
-            None, description="Minimum value for the uncorrected total energy in eV.",
-        ),
-    ) -> STORE_PARAMS:
-
-        crit = defaultdict(dict)  # type: dict
-
-        d = {
-            "energy_per_atom": [energy_per_atom_min, energy_per_atom_max],
-            "formation_energy_per_atom": [formation_energy_per_atom_min, formation_energy_per_atom_max,],
-            "energy_above_hull": [energy_above_hull_min, energy_above_hull_max],
-            "equillibrium_reaction_energy_per_atom": [
-                equillibrium_reaction_energy_per_atom_min,
-                equillibrium_reaction_energy_per_atom_max,
-            ],
-            "uncorrected_energy": [uncorrected_energy_per_atom_min, uncorrected_energy_per_atom_max,],
-        }
-
-        for entry in d:
-            if d[entry][0] is not None:
-                crit[entry]["$gte"] = d[entry][0]
-
-            if d[entry][1] is not None:
-                crit[entry]["$lte"] = d[entry][1]
-
-        return {"criteria": crit}
-
-    def ensure_indexes(self):
-        keys = self._keys_from_query()
-        indexes = []
-        for key in keys:
-            if "_min" in key:
-                key = key.replace("_min", "")
-                indexes.append((key, False))
-        return indexes
diff --git a/src/mp_api/routes/thermo/resources.py b/src/mp_api/routes/thermo/resources.py
index 2eb71ed7..85b23805 100644
--- a/src/mp_api/routes/thermo/resources.py
+++ b/src/mp_api/routes/thermo/resources.py
@@ -1,3 +1,4 @@
+from maggma.api.query_operator.dynamic import NumericQuery
 from maggma.api.resource import ReadOnlyResource
 from emmet.core.thermo import ThermoDoc
 
@@ -9,7 +10,6 @@
 )
 from mp_api.routes.thermo.query_operators import (
     ThermoChemicalQuery,
-    ThermoEnergyQuery,
     IsStableQuery,
 )
 
@@ -27,10 +27,12 @@ def thermo_resource(thermo_store):
             MultiMaterialIDQuery(),
             ThermoChemicalQuery(),
             IsStableQuery(),
-            ThermoEnergyQuery(),
+            NumericQuery(model=ThermoDoc),
             SortQuery(),
             PaginationQuery(),
-            SparseFieldsQuery(ThermoDoc, default_fields=["material_id", "last_updated"]),
+            SparseFieldsQuery(
+                ThermoDoc, default_fields=["material_id", "last_updated"]
+            ),
         ],
         tags=["Thermo"],
     )