from __future__ import annotations
from collections import defaultdict
from emmet.core.molecules.summary import HasProps, MoleculeSummaryDoc
from emmet.core.mpid import MPculeID
from mp_api.client.core import BaseRester
[docs]
class MoleculesSummaryRester(BaseRester):
suffix = "molecules/summary"
document_model = MoleculeSummaryDoc # type: ignore
primary_key = "molecule_id"
[docs]
def search(
self,
charge: int | None = None,
spin_multiplicity: int | None = None,
nelements: tuple[int, int] | None = None,
chemsys: str | list[str] | None = None,
elements: list[str] | None = None,
exclude_elements: list[str] | None = None,
formula: str | list[str] | None = None,
has_props: list[HasProps] | None = None,
molecule_ids: str | list[str | MPculeID] | None = None,
num_chunks: int | None = None,
chunk_size: int = 1000,
all_fields: bool = True,
fields: list[str] | None = None,
_page: int | None = None,
_sort_fields: str | None = None,
):
"""Query core data using a variety of search criteria.
Arguments:
charge (int): Minimum and maximum charge for the molecule.
spin_multiplicity (int): Minimum and maximum spin for the molecule.
nelements (Tuple[int, int]): Minimum and maximum number of elements
chemsys (str, List[str]): A chemical system, list of chemical systems
(e.g., Li-C-O, [C-O-H-N, Li-N]).
#deprecated (bool): Whether the material is tagged as deprecated.
elements (List[str]): A list of elements.
exclude_elements (List(str)): List of elements to exclude.
formula (str, List[str]): An alphabetical formula or list of formulas
(e.g. "C2 Li2 O4", ["C2 H4", "C2 H6"]).
has_props: (List[HasProps]): The calculated properties available for the material.
molecule_ids (str or MPculeID, or list[str | MPculeID]):
(List of) Materials Project Molecule IDs (MPculeIDs) to return data for.
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.
all_fields (bool): Whether to return all fields in the document. Defaults to True.
fields (List[str]): List of fields in SearchDoc to return data for.
Default is material_id if all_fields is False.
_page (int or None) : Page of the results to skip to.
_sort_fields (str or None) : Field to sort on. Including a leading "-" sign will reverse sort order.
Returns:
([MoleculeSummaryDoc]) List of molecules summary documents
"""
query_params: dict = defaultdict(dict)
min_max = [
"nelements",
"ionization_energy",
"electron_affinity",
"reduction_free_energy",
"oxidation_free_energy",
]
for param, value in locals().items():
if param in min_max and value:
if isinstance(value, (int, float)):
value = (value, value)
query_params.update(
{
f"{param}_min": value[0],
f"{param}_max": value[1],
}
)
if molecule_ids:
if isinstance(molecule_ids, str | MPculeID):
molecule_ids = [molecule_ids]
query_params.update({"molecule_ids": ",".join(molecule_ids)})
_locals = locals()
for k in ("charge", "spin_multiplicity", "_page", "_sort_fields"):
if (v := _locals.get(k)) is not None:
query_params[k] = v
for k in ("formula", "chemsys", "elements", "exclude_elements"):
if (v := _locals.get(k)) is not None:
query_params[k] = ",".join([v] if isinstance(v, str) else v)
if has_props:
query_params.update({"has_props": ",".join([i.value for i in has_props])})
query_params = {
entry: query_params[entry]
for entry in query_params
if query_params[entry] is not None
}
return super()._search(
num_chunks=num_chunks,
chunk_size=chunk_size,
all_fields=all_fields,
fields=fields,
**query_params,
)
