diff --git a/src/mp_api/routes/electronic_structure/query_operators.py b/src/mp_api/routes/electronic_structure/query_operators.py index 22fb3621..baceba80 100644 --- a/src/mp_api/routes/electronic_structure/query_operators.py +++ b/src/mp_api/routes/electronic_structure/query_operators.py @@ -18,29 +18,19 @@ class ESSummaryDataQuery(QueryOperator): def query( self, - band_gap_max: Optional[float] = Query(None, description="Maximum value for the band gap energy in eV."), - band_gap_min: Optional[float] = Query(None, description="Minimum value for the band gap energy in eV."), - efermi_max: Optional[float] = Query(None, description="Maximum value for the fermi energy in eV."), - efermi_min: Optional[float] = Query(None, description="Minimum value for the fermi energy in eV."), - magnetic_ordering: Optional[Ordering] = Query(None, description="Magnetic ordering associated with the data."), - is_gap_direct: Optional[bool] = Query(None, description="Whether a band gap is direct or not."), - is_metal: Optional[bool] = Query(None, description="Whether the material is considered a metal."), + magnetic_ordering: Optional[Ordering] = Query( + None, description="Magnetic ordering associated with the data." + ), + is_gap_direct: Optional[bool] = Query( + None, description="Whether a band gap is direct or not." + ), + is_metal: Optional[bool] = Query( + None, description="Whether the material is considered a metal." + ), ) -> STORE_PARAMS: crit = defaultdict(dict) # type: dict - d = { - "band_gap": [band_gap_min, band_gap_max], - "efermi": [efermi_min, efermi_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] - if magnetic_ordering: crit["magnetic_ordering"] = magnetic_ordering.value @@ -69,13 +59,27 @@ def query( path_type: Optional[BSPathType] = Query( None, description="k-path selection convention for the band structure.", ), - band_gap_max: Optional[float] = Query(None, description="Maximum value for the band gap energy in eV."), - band_gap_min: Optional[float] = Query(None, description="Minimum value for the band gap energy in eV."), - efermi_max: Optional[float] = Query(None, description="Maximum value for the fermi energy in eV."), - efermi_min: Optional[float] = Query(None, description="Minimum value for the fermi energy in eV."), - magnetic_ordering: Optional[Ordering] = Query(None, description="Magnetic ordering associated with the data."), - is_gap_direct: Optional[bool] = Query(None, description="Whether a band gap is direct or not."), - is_metal: Optional[bool] = Query(None, description="Whether the material is considered a metal."), + band_gap_max: Optional[float] = Query( + None, description="Maximum value for the band gap energy in eV." + ), + band_gap_min: Optional[float] = Query( + None, description="Minimum value for the band gap energy in eV." + ), + efermi_max: Optional[float] = Query( + None, description="Maximum value for the fermi energy in eV." + ), + efermi_min: Optional[float] = Query( + None, description="Minimum value for the fermi energy in eV." + ), + magnetic_ordering: Optional[Ordering] = Query( + None, description="Magnetic ordering associated with the data." + ), + is_gap_direct: Optional[bool] = Query( + None, description="Whether a band gap is direct or not." + ), + is_metal: Optional[bool] = Query( + None, description="Whether the material is considered a metal." + ), ) -> STORE_PARAMS: crit = defaultdict(dict) # type: dict @@ -83,7 +87,10 @@ def query( if path_type is not None: d = { - f"bandstructure.{path_type.value}.band_gap": [band_gap_min, band_gap_max], + f"bandstructure.{path_type.value}.band_gap": [ + band_gap_min, + band_gap_max, + ], f"bandstructure.{path_type.value}.efermi": [efermi_min, efermi_max], } @@ -95,7 +102,9 @@ def query( crit[entry]["$lte"] = d[entry][1] if magnetic_ordering: - crit[f"bandstructure.{path_type.value}.magnetic_ordering"] = magnetic_ordering.value + crit[ + f"bandstructure.{path_type.value}.magnetic_ordering" + ] = magnetic_ordering.value if is_gap_direct is not None: crit[f"bandstructure.{path_type.value}.is_gap_direct"] = is_gap_direct @@ -127,15 +136,30 @@ def query( None, description="Projection type for the density of states data.", ), spin: Optional[Union[Literal["1", "-1"], Spin]] = Query( - None, description="Spin channel for density of states data. '1' corresponds to spin up.", - ), - element: Optional[Element] = Query(None, description="Element type for projected density of states data.",), - orbital: Optional[OrbitalType] = Query(None, description="Orbital type for projected density of states data.",), - band_gap_max: Optional[float] = Query(None, description="Maximum value for the band gap energy in eV."), - band_gap_min: Optional[float] = Query(None, description="Minimum value for the band gap energy in eV."), - efermi_max: Optional[float] = Query(None, description="Maximum value for the fermi energy in eV."), - efermi_min: Optional[float] = Query(None, description="Minimum value for the fermi energy in eV."), - magnetic_ordering: Optional[Ordering] = Query(None, description="Magnetic ordering associated with the data."), + None, + description="Spin channel for density of states data. '1' corresponds to spin up.", + ), + element: Optional[Element] = Query( + None, description="Element type for projected density of states data.", + ), + orbital: Optional[OrbitalType] = Query( + None, description="Orbital type for projected density of states data.", + ), + band_gap_max: Optional[float] = Query( + None, description="Maximum value for the band gap energy in eV." + ), + band_gap_min: Optional[float] = Query( + None, description="Minimum value for the band gap energy in eV." + ), + efermi_max: Optional[float] = Query( + None, description="Maximum value for the fermi energy in eV." + ), + efermi_min: Optional[float] = Query( + None, description="Minimum value for the fermi energy in eV." + ), + magnetic_ordering: Optional[Ordering] = Query( + None, description="Magnetic ordering associated with the data." + ), ) -> STORE_PARAMS: crit = defaultdict(dict) # type: dict @@ -146,7 +170,8 @@ def query( if projection_type is not None: if spin is None: raise HTTPException( - status_code=400, detail="Must specify a spin channel for querying dos summary data.", + status_code=400, + detail="Must specify a spin channel for querying dos summary data.", ) else: @@ -182,7 +207,9 @@ def query( key_prefix = f"element.{str(element.value)}.{str(orbital.name)}.{str(spin.value)}" else: - key_prefix = f"element.{str(element.value)}.total.{str(spin.value)}" + key_prefix = ( + f"element.{str(element.value)}.total.{str(spin.value)}" + ) key = f"dos.{key_prefix}.{entry}" diff --git a/src/mp_api/routes/electronic_structure/resources.py b/src/mp_api/routes/electronic_structure/resources.py index dd1712cc..7792bc2a 100644 --- a/src/mp_api/routes/electronic_structure/resources.py +++ b/src/mp_api/routes/electronic_structure/resources.py @@ -1,3 +1,4 @@ +from maggma.api.query_operator.dynamic import NumericQuery from maggma.api.resource import ReadOnlyResource from emmet.core.electronic_structure import ElectronicStructureDoc @@ -30,10 +31,13 @@ def es_resource(es_store): ESSummaryDataQuery(), FormulaQuery(), ElementsQuery(), - MinMaxQuery(), + # MinMaxQuery(), + NumericQuery(model=ElectronicStructureDoc), SortQuery(), PaginationQuery(), - SparseFieldsQuery(ElectronicStructureDoc, default_fields=["material_id", "last_updated"]), + SparseFieldsQuery( + ElectronicStructureDoc, default_fields=["material_id", "last_updated"] + ), ], tags=["Electronic Structure"], ) @@ -50,7 +54,9 @@ def custom_bs_endpoint_prep(self): key_name = "task_id" async def get_object( - task_id: str = Query(..., alias=key_name, title=f"The {key_name} of the {model_name} to get",), + task_id: str = Query( + ..., alias=key_name, title=f"The {key_name} of the {model_name} to get", + ), ): f""" Get's a document by the primary key in the store @@ -71,12 +77,14 @@ async def get_object( if not bs_object_doc: raise HTTPException( - status_code=404, detail=f"Band structure with task_id = {task_id} not found", + status_code=404, + detail=f"Band structure with task_id = {task_id} not found", ) except ValueError: raise HTTPException( - status_code=404, detail=f"Band structure with task_id = {task_id} not found", + status_code=404, + detail=f"Band structure with task_id = {task_id} not found", ) return {"data": [bs_object_doc]} @@ -96,7 +104,10 @@ async def get_object( BSDataQuery(), SortQuery(), PaginationQuery(), - SparseFieldsQuery(ElectronicStructureDoc, default_fields=["material_id", "last_updated", "bandstructure"],), + SparseFieldsQuery( + ElectronicStructureDoc, + default_fields=["material_id", "last_updated", "bandstructure"], + ), ], tags=["Electronic Structure"], enable_get_by_key=False, @@ -115,7 +126,9 @@ def custom_dos_endpoint_prep(self): key_name = "task_id" async def get_object( - task_id: str = Query(..., alias=key_name, title=f"The {key_name} of the {model_name} to get",), + task_id: str = Query( + ..., alias=key_name, title=f"The {key_name} of the {model_name} to get", + ), ): f""" Get's a document by the primary key in the store @@ -136,12 +149,14 @@ async def get_object( if not dos_object_doc: raise HTTPException( - status_code=404, detail=f"Density of states with task_id = {task_id} not found", + status_code=404, + detail=f"Density of states with task_id = {task_id} not found", ) except ValueError: raise HTTPException( - status_code=404, detail=f"Density of states with task_id = {task_id} not found", + status_code=404, + detail=f"Density of states with task_id = {task_id} not found", ) return {"data": [dos_object_doc]} @@ -161,7 +176,10 @@ async def get_object( DOSDataQuery(), SortQuery(), PaginationQuery(), - SparseFieldsQuery(ElectronicStructureDoc, default_fields=["material_id", "last_updated", "dos"],), + SparseFieldsQuery( + ElectronicStructureDoc, + default_fields=["material_id", "last_updated", "dos"], + ), ], tags=["Electronic Structure"], custom_endpoint_funcs=[custom_dos_endpoint_prep],