More MPRester Updates

.coveragerc

@@ -4,4 +4,5 @@ omit =
     */resources.py
     */models.py
     */models/*
+    */_consumer/client.py
 

setup.py

@@ -30,6 +30,7 @@
         "requests>=2.23.0",
         "monty",
         "emmet-core",
+        "maggma",
     ],
     extras_require={
         "server": [

src/mp_api/__init__.py

@@ -6,7 +6,7 @@
 
 try:
     __version__ = get_distribution(__name__).version
-except DistributionNotFound:
+except DistributionNotFound:  # pragma: no cover
     # package is not installed
     __version__ = None  # type: ignore
 
@@ -15,7 +15,7 @@
 
 settings = MAPISettings()
 
-try:
+try:  # pragma: no cover
     if Path(settings.app_path).exists():
         mapi = loadfn(settings.app_path)
         app = mapi.app
@@ -24,7 +24,7 @@
         if settings.debug:
             print(f"Failed loading App at {settings.app_path}")
 
-except Exception as e:
+except Exception as e:  # pragma: no cover
     # Something went wrong with loading default app
     if settings.debug:
         print("Failed loading App")

src/mp_api/core/client.py

@@ -25,7 +25,7 @@
 
 try:
     from pymatgen.core import __version__ as pmg_version  # type: ignore
-except ImportError:
+except ImportError:  # pragma: no cover
     # fallback to root-level import for older pymatgen versions
     from pymatgen import __version__ as pmg_version  # type: ignore
 
@@ -112,13 +112,13 @@ def _create_session(api_key, include_user_agent):
             )
         return session
 
-    def __enter__(self):
+    def __enter__(self):  # pragma: no cover
         """
         Support for "with" context.
         """
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb):  # pragma: no cover
         """
         Support for "with" context.
         """
@@ -330,7 +330,7 @@ def get_document_by_id(
         else:
             criteria = {"limit": 1}
 
-        if isinstance(fields, str):
+        if isinstance(fields, str):  # pragma: no cover
             fields = (fields,)
 
         results = []
@@ -368,7 +368,7 @@ def get_document_by_id(
 
         if not results:
             raise MPRestError(f"No result for record {document_id}.")
-        elif len(results) > 1:
+        elif len(results) > 1:  # pragma: no cover
             raise ValueError(
                 f"Multiple records for {document_id}, this shouldn't happen. Please report as a bug."
             )
@@ -475,7 +475,7 @@ def _get_all_documents(
 
         return all_results
 
-    def query_by_task_id(self, *args, **kwargs):  # pragma: ignore
+    def query_by_task_id(self, *args, **kwargs):  # pragma: no cover
         print(
             "query_by_task_id has been renamed to get_document_by_id to be more general"
         )
@@ -496,19 +496,19 @@ def count(self, criteria: Optional[Dict] = None) -> Union[int, str]:
                 criteria=criteria, monty_decode=False
             )  # do not waste cycles Monty decoding
             return results["meta"]["total_doc"]
-        except Exception:
-            return "unknown"
+        except Exception:  # pragma: no cover
+            return "Problem getting count"
 
     @property
     def available_fields(self) -> List[str]:
         if self.document_model is None:
             return ["Unknown fields."]
         return list(self.document_model.schema()["properties"].keys())  # type: ignore
 
-    def __repr__(self):  # pragma: ignore
+    def __repr__(self):  # pragma: no cover
         return f"<{self.__class__.__name__} {self.endpoint}>"
 
-    def __str__(self):  # pragma: ignore
+    def __str__(self):  # pragma: no cover
         if self.document_model is None:
             return self.__repr__()
         return (

src/mp_api/matproj.py

@@ -1,12 +1,14 @@
 from os import environ
 import warnings
-from collections import defaultdict
+from typing import Optional, Tuple, List
 from enum import Enum, unique
 
 from pymatgen.core import Structure
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from pymatgen.core.surface import get_symmetrically_equivalent_miller_indices
-from pymatgen.util.sequence import get_chunks
+from pymatgen.analysis.magnetism import Ordering
+from emmet.core.mpid import MPID
+from emmet.core.symmetry import CrystalSystem
 
 from mp_api.core.client import BaseRester
 from mp_api.routes import *
@@ -205,13 +207,13 @@ def get_materials_id_from_task_id(self, task_id):
             materials_id (MPID)
         """
         docs = self.materials.search(task_ids=[task_id], fields=["material_id"])
-        if len(docs) == 1:
+        if len(docs) == 1:  # pragma: no cover
             return str(docs[0].material_id)
-        elif len(docs) > 1:
+        elif len(docs) > 1:  # pragma: no cover
             raise ValueError(
                 f"Multiple documents return for {task_id}, this should not happen, please report it!"
             )
-        else:
+        else:  # pragma: no cover
             warnings.warn(
                 f"No material found containing task {task_id}. Please report it if you suspect a task has gone missing."
             )
@@ -370,65 +372,176 @@ def get_phonon_bandstructure_by_material_id(self, material_id):
 
     def query(
         self,
-        criteria,
-        properties,
-        sort_field=None,
-        ascending=None,
-        num_chunks=None,
-        chunk_size=1000,
-        all_fields=True,
-        fields=None,
+        material_ids: Optional[List[MPID]] = None,
+        chemsys_formula: Optional[str] = None,
+        nsites: Optional[Tuple[int, int]] = None,
+        volume: Optional[Tuple[float, float]] = None,
+        density: Optional[Tuple[float, float]] = None,
+        crystal_system: Optional[CrystalSystem] = None,
+        spacegroup_number: Optional[int] = None,
+        spacegroup_symbol: Optional[str] = None,
+        deprecated: Optional[bool] = None,
+        total_energy: Optional[Tuple[float, float]] = None,
+        formation_energy: Optional[Tuple[float, float]] = None,
+        energy_above_hull: Optional[Tuple[float, float]] = None,
+        equillibrium_reaction_energy: Optional[Tuple[float, float]] = None,
+        uncorrected_energy: Optional[Tuple[float, float]] = None,
+        is_stable: Optional[bool] = None,
+        band_gap: Optional[Tuple[float, float]] = None,
+        efermi: Optional[Tuple[float, float]] = None,
+        is_gap_direct: Optional[bool] = None,
+        is_metal: Optional[bool] = None,
+        magnetic_ordering: Optional[Ordering] = None,
+        total_magnetization: Optional[Tuple[float, float]] = None,
+        total_magnetization_normalized_vol: Optional[Tuple[float, float]] = None,
+        total_magnetization_normalized_formula_units: Optional[
+            Tuple[float, float]
+        ] = None,
+        k_voigt: Optional[Tuple[float, float]] = None,
+        k_reuss: Optional[Tuple[float, float]] = None,
+        k_vrh: Optional[Tuple[float, float]] = None,
+        g_voigt: Optional[Tuple[float, float]] = None,
+        g_reuss: Optional[Tuple[float, float]] = None,
+        g_vrh: Optional[Tuple[float, float]] = None,
+        elastic_anisotropy: Optional[Tuple[float, float]] = None,
+        poisson_ratio: Optional[Tuple[float, float]] = None,
+        e_total: Optional[Tuple[float, float]] = None,
+        e_ionic: Optional[Tuple[float, float]] = None,
+        e_static: Optional[Tuple[float, float]] = None,
+        n: Optional[Tuple[float, float]] = None,
+        piezoelectric_modulus: Optional[Tuple[float, float]] = None,
+        weighted_surface_energy: Optional[Tuple[float, float]] = None,
+        weighted_work_function: Optional[Tuple[float, float]] = None,
+        surface_anisotropy: Optional[Tuple[float, float]] = None,
+        shape_factor: Optional[Tuple[float, float]] = None,
+        has_props: Optional[List[str]] = None,
+        theoretical: Optional[bool] = None,
+        sort_field: Optional[str] = None,
+        ascending: Optional[bool] = None,
+        num_chunks: Optional[int] = None,
+        chunk_size: int = 1000,
+        all_fields: bool = True,
+        fields: Optional[List[str]] = None,
     ):
-        r"""
-
-        Performs an advanced query using MongoDB-like syntax for directly
-        querying the Materials Project database. This allows one to perform
-        queries which are otherwise too cumbersome to perform using the standard
-        convenience methods.
-
-        Please consult the Materials API documentation at
-        https://github.com/materialsproject/mapidoc, which provides a
-        comprehensive explanation of the document schema used in the Materials
-        Project (supported criteria and properties) and guidance on how best to
-        query for the relevant information you need.
-
-        For queries that request data on more than CHUNK_SIZE materials at once,
-        this method will chunk a query by first retrieving a list of material
-        IDs that satisfy CRITERIA, and then merging the criteria with a
-        restriction to one chunk of materials at a time of size CHUNK_SIZE. You
-        can opt out of this behavior by setting CHUNK_SIZE=0. To guard against
-        intermittent server errors in the case of many chunks per query,
-        possibly-transient server errors will result in re-trying a give chunk
-        up to MAX_TRIES_PER_CHUNK times.
+        """
+        Query core data using a variety of search criteria.
 
-        Args:
-            criteria (str/dict): Criteria of the query as a dictionary.
-                For example, {"elements":{"$in":["Li",
-                "Na", "K"], "$all": ["O"]}, "nelements":2} selects all Li, Na
-                and K oxides. {"band_gap": {"$gt": 1}} selects all materials
-                with band gaps greater than 1 eV.
-            properties (list): Properties to request for as a list. For
-                example, ["formula_pretty", "formation_energy_per_atom"] returns
-                the formula and formation energy per atom.
-            chunk_size (int): Number of materials for which to fetch data at a
-                time. More data-intensive properties may require smaller chunk
-                sizes. Use chunk_size=0 to force no chunking -- this is useful
-                when fetching only properties such as 'material_id'.
-            max_tries_per_chunk (int): How many times to re-try fetching a given
-                chunk when the server gives a 5xx error (e.g. a timeout error).
-            mp_decode (bool): Whether to do a decoding to a Pymatgen object
-                where possible. In some cases, it might be useful to just get
-                the raw python dict, i.e., set to False.
+        Arguments:
+            material_ids (List[MPID]): List of Materials Project IDs to return data for.
+            chemsys_formula (str): A chemical system (e.g., Li-Fe-O),
+                or formula including anonomyzed formula
+                or wild cards (e.g., Fe2O3, ABO3, Si*).
+            crystal_system (CrystalSystem): Crystal system of material.
+            spacegroup_number (int): Space group number of material.
+            spacegroup_symbol (str): Space group symbol of the material in international short symbol notation.
+            nsites (Tuple[int,int]): Minimum and maximum number of sites to consider.
+            volume (Tuple[float,float]): Minimum and maximum volume to consider.
+            density (Tuple[float,float]): Minimum and maximum density to consider.
+            deprecated (bool): Whether the material is tagged as deprecated.
+            total_energy (Tuple[int,int]): Minimum and maximum corrected total energy in eV/atom to consider.
+            formation_energy (Tuple[int,int]): Minimum and maximum formation energy in eV/atom to consider.
+            energy_above_hull (Tuple[int,int]): Minimum and maximum energy above the hull in eV/atom to consider.
+            uncorrected_energy (Tuple[int,int]): Minimum and maximum uncorrected total energy in eV/atom to consider.
+            band_gap (Tuple[float,float]): Minimum and maximum band gap in eV to consider.
+            efermi (Tuple[float,float]): Minimum and maximum fermi energy in eV to consider.
+            is_gap_direct (bool): Whether the material has a direct band gap.
+            is_metal (bool): Whether the material is considered a metal.
+            magnetic_ordering (Ordering): Magnetic ordering of the material.
+            total_magnetization (Tuple[float,float]): Minimum and maximum total magnetization values to consider.
+            total_magnetization_normalized_vol (Tuple[float,float]): Minimum and maximum total magnetization values
+                normalized by volume to consider.
+            total_magnetization_normalized_formula_units (Tuple[float,float]): Minimum and maximum total magnetization
+                values normalized by formula units to consider.
+            k_voigt (Tuple[float,float]): Minimum and maximum value in GPa to consider for
+                the Voigt average of the bulk modulus.
+            k_reuss (Tuple[float,float]): Minimum and maximum value in GPa to consider for
+                the Reuss average of the bulk modulus.
+            k_vrh (Tuple[float,float]): Minimum and maximum value in GPa to consider for
+                the Voigt-Reuss-Hill average of the bulk modulus.
+            g_voigt (Tuple[float,float]): Minimum and maximum value in GPa to consider for
+                the Voigt average of the shear modulus.
+            g_reuss (Tuple[float,float]): Minimum and maximum value in GPa to consider for
+                the Reuss average of the shear modulus.
+            g_vrh (Tuple[float,float]): Minimum and maximum value in GPa to consider for
+                the Voigt-Reuss-Hill average of the shear modulus.
+            elastic_anisotropy (Tuple[float,float]): Minimum and maximum value to consider for
+                the elastic anisotropy.
+            poisson_ratio (Tuple[float,float]): Minimum and maximum value to consider for
+                Poisson's ratio.
+            e_total (Tuple[float,float]): Minimum and maximum total dielectric constant to consider.
+            e_ionic (Tuple[float,float]): Minimum and maximum ionic dielectric constant to consider.
+            e_static (Tuple[float,float]): Minimum and maximum electronic dielectric constant to consider.
+            n (Tuple[float,float]): Minimum and maximum refractive index to consider.
+            piezoelectric_modulus (Tuple[float,float]): Minimum and maximum piezoelectric modulus to consider.
+            weighted_surface_energy (Tuple[float,float]): Minimum and maximum weighted surface energy in J/m² to
+                consider.
+            weighted_work_function (Tuple[float,float]): Minimum and maximum weighted work function in eV to consider.
+            surface_energy_anisotropy (Tuple[float,float]): Minimum and maximum surface energy anisotropy values to
+                consider.
+            shape_factor (Tuple[float,float]): Minimum and maximum shape factor values to consider.
+            has_props: (List[str]): The calculated properties available for the material.
+            theoretical: (bool): Whether the material is theoretical.
+            sort_field (str): Field used to sort results.
+            ascending (bool): Whether sorting should be in ascending order.
+            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.
 
         Returns:
-            List of results. E.g.,
-            [{u'composition': {u'O': 1, u'Li': 2.0}},
-            {u'composition': {u'Na': 2.0, u'O': 2.0}},
-            {u'composition': {u'K': 1, u'O': 3.0}},
-            ...]
-        """
-        # TODO: discuss
-        raise NotImplementedError
+            ([SearchDoc]) List of SearchDoc documents
+        """
+        return self.search.search_docs(  # type: ignore
+            material_ids=material_ids,
+            chemsys_formula=chemsys_formula,
+            nsites=nsites,
+            volume=volume,
+            density=density,
+            crystal_system=crystal_system,
+            spacegroup_number=spacegroup_number,
+            spacegroup_symbol=spacegroup_symbol,
+            deprecated=deprecated,
+            total_energy=total_energy,
+            formation_energy=formation_energy,
+            energy_above_hull=energy_above_hull,
+            equillibrium_reaction_energy=equillibrium_reaction_energy,
+            uncorrected_energy=uncorrected_energy,
+            is_stable=is_stable,
+            band_gap=band_gap,
+            efermi=efermi,
+            is_gap_direct=is_gap_direct,
+            is_metal=is_metal,
+            magnetic_ordering=magnetic_ordering,
+            total_magnetization=total_magnetization,
+            total_magnetization_normalized_vol=total_magnetization_normalized_vol,
+            total_magnetization_normalized_formula_units=total_magnetization_normalized_formula_units,
+            k_voigt=k_voigt,
+            k_reuss=k_reuss,
+            k_vrh=k_vrh,
+            g_voigt=g_voigt,
+            g_reuss=g_reuss,
+            g_vrh=g_vrh,
+            elastic_anisotropy=elastic_anisotropy,
+            poisson_ratio=poisson_ratio,
+            e_total=e_total,
+            e_ionic=e_ionic,
+            e_static=e_static,
+            n=n,
+            piezoelectric_modulus=piezoelectric_modulus,
+            weighted_surface_energy=weighted_surface_energy,
+            weighted_work_function=weighted_work_function,
+            surface_anisotropy=surface_anisotropy,
+            shape_factor=shape_factor,
+            has_props=has_props,
+            theoretical=theoretical,
+            sort_field=sort_field,
+            ascending=ascending,
+            num_chunks=num_chunks,
+            chunk_size=chunk_size,
+            all_fields=all_fields,
+            fields=fields,
+        )
 
     def submit_structures(self, structures, public_name, public_email):
         """