revise find structure method to push more work to user

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -1617,9 +1617,7 @@ def __getattr__(self, v: str):
                     force_renew=self.force_renew,
                 )
             return self.sub_resters[v]
-        raise AttributeError(
-            f"{self.__class__} has no attribute {v}"
-        )
+        raise AttributeError(f"{self.__class__} has no attribute {v}")
 
     def __dir__(self):
         return dir(self.__class__) + list(self._sub_resters)

mp_api/client/routes/materials/materials.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 from emmet.core.symmetry import CrystalSystem
@@ -172,11 +173,11 @@ def search(
 
     def find_structure(
         self,
-        filename_or_structure,
+        filename_or_structure: str | Path | Structure,
         ltol=MAPI_CLIENT_SETTINGS.LTOL,
         stol=MAPI_CLIENT_SETTINGS.STOL,
         angle_tol=MAPI_CLIENT_SETTINGS.ANGLE_TOL,
-        allow_multiple_results=False,
+        allow_multiple_results: bool | int = False,
     ) -> list[str] | str:
         """Finds matching structures from the Materials Project database.
 
@@ -186,48 +187,72 @@ def find_structure(
         default tolerances.
 
         Args:
-            filename_or_structure: filename or Structure object
+            filename_or_structure: filename as a str or Path, or a Structure object
             ltol: fractional length tolerance
             stol: site tolerance
             angle_tol: angle tolerance in degrees
-            allow_multiple_results: changes return type for either
-            a single material_id or list of material_ids
+            allow_multiple_results (bool or int): changes return type for either
+                a single material_id or list of material_ids.
+                If a bool, returns either all matches (True) or one match at most (False).
+                If an int, returns that many matches at most.
+
         Returns:
             A matching material_id if one is found or list of results if allow_multiple_results
             is True
         Raises:
             MPRestError
         """
-        params = {"ltol": ltol, "stol": stol, "angle_tol": angle_tol, "_limit": 1}
+        from pymatgen.analysis.structure_matcher import (
+            ElementComparator,
+            StructureMatcher,
+        )
 
-        if isinstance(filename_or_structure, str):
+        if (
+            isinstance(filename_or_structure, str | Path)
+            and Path(filename_or_structure).exists()
+        ):
             s = Structure.from_file(filename_or_structure)
         elif isinstance(filename_or_structure, Structure):
             s = filename_or_structure
         else:
             raise MPRestError("Provide filename or Structure object.")
 
-        results = self._post_resource(
-            body=s.as_dict(),
-            params=params,
-            suburl="find_structure",
-            use_document_model=False,
-        ).get("data")
-
-        if not results:
+        mat_docs = self.search(
+            formula=s.reduced_formula, fields=["material_id", "structure"]
+        )
+        if not mat_docs:
             return []
 
-        material_ids = validate_ids([doc["material_id"] for doc in results])
+        if isinstance(allow_multiple_results, bool):
+            max_matches: int = len(mat_docs) if allow_multiple_results else 1
+        elif isinstance(allow_multiple_results, int):
+            max_matches = allow_multiple_results
+        else:
+            raise MPRestError(
+                f"`allow_multiple_results` must be a bool or int, not {type(allow_multiple_results)}"
+            )
+
+        matcher = StructureMatcher(
+            ltol=ltol,
+            stol=stol,
+            angle_tol=angle_tol,
+            primitive_cell=True,
+            scale=True,
+            attempt_supercell=False,
+            comparator=ElementComparator(),
+        )
 
-        if len(material_ids) > 1:  # type: ignore
-            if not allow_multiple_results:
-                raise ValueError(
-                    "Multiple matches found for this combination of tolerances, but "
-                    "`allow_multiple_results` set to False."
-                )
-            return material_ids  # type: ignore
+        matches = []
+        for doc in mat_docs:
+            if matcher.fit(s, doc.structure):
+                matches.append(doc.material_id)
+                if len(matches) >= max_matches:
+                    break
 
-        return material_ids[0]
+        if not matches:
+            return []
+        material_ids = validate_ids(matches)
+        return material_ids if allow_multiple_results else material_ids[0]
 
     def get_blessed_entries(
         self,

mp_api/client/routes/molecules/molecules.py

@@ -7,16 +7,15 @@
 
 from __future__ import annotations
 
-from emmet.core.mpid import MPculeID
 from emmet.core.qchem.molecule import MoleculeDoc
-from pymatgen.core.structure import Molecule
 
-from mp_api.client.core.client import CoreRester, MPRestError
-from mp_api.client.core.utils import validate_ids
+from mp_api.client.core.client import CoreRester
 from mp_api.client.routes.molecules import MOLECULES_RESTERS
 
+
 class MoleculeRester(CoreRester):
     """Define molecules stub for accessing JCESR and summary data."""
+
     document_model = MoleculeDoc
     primary_key = "molecule_id"
     suffix = "molecules/core"

mp_api/client/routes/molecules/summary.py

@@ -43,7 +43,7 @@ def search(
             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]): 
+            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.
@@ -76,7 +76,7 @@ def search(
                 )
 
         if molecule_ids:
-            if isinstance(molecule_ids,str):
+            if isinstance(molecule_ids, str):
                 molecule_ids = [molecule_ids]
             query_params.update({"molecule_ids": ",".join(molecule_ids)})
 
@@ -119,4 +119,4 @@ def search(
             all_fields=all_fields,
             fields=fields,
             **query_params,
-        )
+        )

tests/client/molecules/test_summary.py

@@ -25,7 +25,7 @@
     "molecule_ids": [
         "042b6da7a6eb790fd5038f3729ef715c-C5H8O3-m1-2",
         "6fef3b0ef7f484b8c2dd0a0012bf7c70-C2H2O3P1-0-2",
-        "a3572e662965c190b2a085c6a67a8a03-C9H12-1-2"
+        "a3572e662965c190b2a085c6a67a8a03-C9H12-1-2",
     ],
     "formula": "C2 H4",
     "chemsys": "C-H",

tests/client/test_core_client.py

@@ -20,12 +20,6 @@ def mpr():
     rester.session.close()
 
 
-@requires_api_key
-@pytest.mark.xfail
-def test_post_fail(rester):
-    rester._post_resource({}, suburl="materials/find_structure")
-
-
 @requires_api_key
 def test_pagination(mpr):
     mpids = mpr.materials.search(

tests/client/test_mprester.py

@@ -131,20 +131,39 @@ def test_find_structure(self, mpr):
  _atom_site_occupancy
   Ne  Ne0  1  0.00000000  0.00000000  -0.00000000  1
 """
+        struct_from_cif = CifParser.from_str(cif_str).parse_structures(primitive=True)[
+            0
+        ]
         temp_file = NamedTemporaryFile(suffix=".cif")
         with open(temp_file.name, "wt") as f:
             f.write(cif_str)
             f.seek(0)
 
         for struct_or_path in (
             temp_file.name,
-            CifParser.from_str(cif_str).parse_structures(primitive=True)[0],
+            struct_from_cif,
         ):
             data = mpr.find_structure(struct_or_path)
             assert isinstance(data, str) and data == "mp-111"
 
         f.close()
 
+        with pytest.raises(MPRestError, match="Provide filename or Structure object."):
+            mpr.find_structure(struct_from_cif.as_dict())
+
+        with pytest.raises(MPRestError, match="`allow_multiple_results` must be a"):
+            mpr.find_structure(struct_from_cif, allow_multiple_results=1.0)
+
+        assert (
+            len(
+                mpr.find_structure(
+                    struct_from_cif.copy().replace_species({"Ne": "K"}),
+                    allow_multiple_results=2,
+                )
+            )
+            <= 2
+        )
+
     def test_get_bandstructure_by_material_id(self, mpr):
         bs = mpr.get_bandstructure_by_material_id("mp-149")
         assert isinstance(bs, BandStructureSymmLine)