Skip to content

Commit 7913e24

Browse files
merge conflicts
2 parents ede36d0 + 1d25d67 commit 7913e24

5 files changed

Lines changed: 108 additions & 49 deletions

File tree

mp_api/client/core/client.py

Lines changed: 39 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -659,7 +659,7 @@ def _query_resource(
659659
use_document_model: if None, will defer to the self.use_document_model attribute
660660
num_chunks: Maximum number of chunks of data to yield. None will yield all possible.
661661
chunk_size: Number of data entries per chunk.
662-
timeout : Time in seconds to wait until a request timeout error is thrown
662+
timeout (float or None): Time in seconds to wait until a request timeout error is thrown
663663
664664
Returns:
665665
A Resource, a dict with two keys, "data" containing a list of documents, and
@@ -805,27 +805,32 @@ def _query_resource(
805805
except RequestException as ex:
806806
raise MPRestError(str(ex))
807807

808-
def _submit_requests( # noqa
808+
def _submit_requests(
809809
self,
810-
url,
811-
criteria,
812-
use_document_model,
813-
chunk_size,
814-
num_chunks=None,
815-
timeout=None,
810+
url: str,
811+
criteria: dict[str, Any],
812+
use_document_model: bool,
813+
chunk_size: int | None,
814+
num_chunks: int | None = None,
815+
timeout: int | None = None,
816+
max_batch_size: int = 100,
817+
norecur: bool = False,
816818
) -> dict:
817819
"""Handle submitting requests sequentially with pagination.
818820
819821
If criteria contains comma-separated parameters (except those that are naturally comma-separated),
820822
split them into multiple sequential requests and combine results.
821823
822824
Arguments:
823-
criteria: dictionary of criteria to filter down
824-
url: url used to make request
825-
use_document_model: if None, will defer to the self.use_document_model attribute
826-
num_chunks: Maximum number of chunks of data to yield. None will yield all possible.
827-
chunk_size: Number of data entries per chunk.
828-
timeout: Time in seconds to wait until a request timeout error is thrown
825+
url (str): url used to make request
826+
criteria (dict of str): dictionary of criteria to filter down
827+
use_document_model (bool): whether to use the document model
828+
num_chunks (int or None): Maximum number of chunks of data to yield. None will yield all possible.
829+
chunk_size (int or None): Number of data entries per chunk.
830+
timeout (int or None): Time in seconds to wait until a request timeout error is thrown
831+
max_batch_size (int) : Maximum size of a batch when retrieving batches in parallel
832+
norecur (bool) : Whether to forbid recursive splitting of a query field
833+
when a direct query fails
829834
830835
Returns:
831836
Dictionary containing data and metadata
@@ -884,14 +889,6 @@ def _submit_requests( # noqa
884889
timeout=timeout,
885890
)
886891

887-
# Check if we got 0 results - some parameters are silently ignored by the API
888-
# when passed as comma-separated values, so we need to split them anyway
889-
if total_num_docs == 0 and len(split_values) > 1:
890-
# Treat this the same as a 422 error - split into batches
891-
raise MPRestError(
892-
"Got 0 results for comma-separated parameter, will try splitting"
893-
)
894-
895892
# If successful, continue with normal pagination
896893
data_chunks = [data["data"]]
897894
total_data: dict[str, Any] = {"data": []}
@@ -903,18 +900,26 @@ def _submit_requests( # noqa
903900
# Continue with pagination if needed (handled below)
904901

905902
except MPRestError as e:
906-
# If we get 422 or 414 error, or 0 results for comma-separated params, split into batches
907-
if any(trace in str(e) for trace in ("422", "414", "Got 0 results")):
903+
# If we get 422 or 414 error, split into batches
904+
if not norecur and any(
905+
trace in str(e)
906+
for trace in (
907+
"422",
908+
"414",
909+
)
910+
):
908911
total_data = {"data": []}
909912
total_num_docs = 0
910913
data_chunks = []
911914

912915
# Batch the split values to reduce number of requests
913916
# Use batches of up to 100 values to balance URL length and request count
914-
batch_size = min(100, max(1, len(split_values) // 10))
917+
num_batches = min(
918+
max_batch_size, max(1, len(split_values) // max_batch_size)
919+
)
920+
batch_size = min(len(split_values), max_batch_size)
915921

916922
# Setup progress bar for split parameter requests
917-
num_batches = ceil(len(split_values) / batch_size)
918923
pbar_message = f"Retrieving {len(split_values)} {split_param} values in {num_batches} batches"
919924
pbar = (
920925
tqdm(
@@ -938,6 +943,7 @@ def _submit_requests( # noqa
938943
chunk_size=chunk_size,
939944
num_chunks=num_chunks,
940945
timeout=timeout,
946+
norecur=len(batch) <= max_batch_size,
941947
)
942948

943949
data_chunks.append(result["data"])
@@ -983,6 +989,12 @@ def _submit_requests( # noqa
983989
if "meta" in data:
984990
total_data["meta"] = data["meta"]
985991

992+
# otherwise, paginate sequentially
993+
if chunk_size is None or chunk_size < 1:
994+
raise ValueError(
995+
"A positive chunk size must be provided to enable pagination"
996+
)
997+
986998
# Get max number of response pages
987999
max_pages = (
9881000
num_chunks if num_chunks is not None else ceil(total_num_docs / chunk_size)
@@ -1002,7 +1014,7 @@ def _submit_requests( # noqa
10021014
desc=pbar_message,
10031015
total=num_docs_needed,
10041016
)
1005-
if not self.mute_progress_bars
1017+
if not self.mute_progress_bars and total_num_docs > 0
10061018
else None
10071019
)
10081020

@@ -1022,10 +1034,6 @@ def _submit_requests( # noqa
10221034
pbar.close()
10231035
return new_total_data
10241036

1025-
# otherwise, paginate sequentially
1026-
if chunk_size is None:
1027-
raise ValueError("A chunk size must be provided to enable pagination")
1028-
10291037
# Warning to select specific fields only for many results
10301038
if criteria.get("_all_fields", False) and (total_num_docs / chunk_size > 10):
10311039
warnings.warn(

mp_api/client/mprester.py

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@
5858

5959
from mp_api.client.core.client import _DictLikeAccess
6060

61-
DEFAULT_THERMOTYPE_CRITERIA = {"thermo_types": ["GGA_GGA+U"]}
61+
DEFAULT_THERMOTYPE_CRITERIA = {"thermo_types": ["GGA_GGA+U_R2SCAN"]}
6262

6363
RESTER_LAYOUT = {
6464
"molecules/core": LazyImport(
@@ -1032,7 +1032,7 @@ def get_entries_in_chemsys(
10321032
compatible_only: bool = True,
10331033
property_data: list[str] | None = None,
10341034
conventional_unit_cell: bool = False,
1035-
additional_criteria: dict = DEFAULT_THERMOTYPE_CRITERIA,
1035+
additional_criteria: dict | None = None,
10361036
**kwargs,
10371037
) -> list[ComputedStructureEntry] | list[GibbsComputedStructureEntry]:
10381038
"""Helper method to get a list of ComputedEntries in a chemical system.
@@ -1091,6 +1091,17 @@ def get_entries_in_chemsys(
10911091
for els in itertools.combinations(elements_set, i + 1)
10921092
]
10931093

1094+
if additional_criteria is None:
1095+
warnings.warn(
1096+
"The default thermo type when retrieving entries has been changed from "
1097+
"the mixed/corrected PBE GGA and GGA+U hull (`thermo_type = GGA_GGA+U`) "
1098+
"to the joint PBE GGA / GGA+U / r2SCAN hull (`thermo_type = GGA_GGA+U_R2SCAN`). "
1099+
"To use the older behavior, call `get_entries_in_chemsys` with "
1100+
'`additional_criteria = {"thermo_types": ["GGA_GGA+U"]}`',
1101+
category=MPRestWarning,
1102+
stacklevel=2,
1103+
)
1104+
10941105
entries = self.get_entries(
10951106
all_chemsyses,
10961107
compatible_only=compatible_only,
@@ -1348,23 +1359,17 @@ def get_download_info(
13481359
return meta, urls
13491360

13501361
def _check_get_download_info_url_by_task_id(self, prefix, task_ids) -> list[str]:
1351-
nomad_exist_task_ids: list[str] = []
13521362
prefix = prefix.replace("/raw/query", "/repo/")
1353-
for task_id in task_ids:
1354-
url = prefix + task_id
1355-
if self._check_nomad_exist(url):
1356-
nomad_exist_task_ids.append(task_id)
1357-
return nomad_exist_task_ids
1363+
return [
1364+
task_id for task_id in task_ids if self._check_nomad_exist(prefix + task_id)
1365+
]
13581366

13591367
@staticmethod
13601368
def _check_nomad_exist(url) -> bool:
13611369
response = get(url=url)
13621370
if response.status_code != 200:
13631371
return False
1364-
content = load_json(response.text)
1365-
if content["pagination"]["total"] == 0:
1366-
return False
1367-
return True
1372+
return load_json(response.text)["pagination"]["total"] != 0
13681373

13691374
@staticmethod
13701375
def _print_help_message(nomad_exist_task_ids, task_ids, file_patterns, calc_types):

mp_api/client/routes/materials/summary.py

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import warnings
44
from collections import defaultdict
5+
from itertools import chain, product
56

67
from emmet.core.summary import HasProps, SummaryDoc
78
from emmet.core.symmetry import CrystalSystem
@@ -324,10 +325,11 @@ def _csrc(x):
324325
"spacegroup_number": 230,
325326
"spacegroup_symbol": 230,
326327
}
328+
batched_symm_query = {}
327329
for k, cardinality in symm_cardinality.items():
328-
if isinstance(symm_vals := locals().get(k), list | tuple | set):
330+
if isinstance(symm_vals := _locals.get(k), list | tuple | set):
329331
if len(symm_vals) < cardinality // 2:
330-
query_params.update({k: ",".join(str(v) for v in symm_vals)})
332+
batched_symm_query[k] = symm_vals
331333
else:
332334
raise MPRestError(
333335
f"Querying `{k}` by a list of values is only "
@@ -370,6 +372,24 @@ def _csrc(x):
370372
if query_params[entry] is not None
371373
}
372374

375+
if batched_symm_query:
376+
ordered_symm_key = sorted(batched_symm_query)
377+
return list(
378+
chain.from_iterable(
379+
self._search( # type: ignore[return-value]
380+
num_chunks=num_chunks,
381+
chunk_size=chunk_size,
382+
all_fields=all_fields,
383+
fields=fields,
384+
**query_params,
385+
**{sk: symm_params[i] for i, sk in enumerate(ordered_symm_key)},
386+
)
387+
for symm_params in product(
388+
*[batched_symm_query[k] for k in ordered_symm_key]
389+
)
390+
)
391+
)
392+
373393
return super()._search( # type: ignore[return-value]
374394
num_chunks=num_chunks,
375395
chunk_size=chunk_size,

tests/client/test_core_client.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,3 +94,21 @@ def test_warnings_exceptions():
9494

9595
with pytest.raises(MPRestError, match="Number of chunks must be greater than zero"):
9696
MaterialsRester()._get_all_documents({}, num_chunks=-1)
97+
98+
99+
def test_regression_batching():
100+
# See https://github.com/materialsproject/api/pull/1077
101+
# This test ensures that queries with batched input
102+
# that return no results do not infinitely recur
103+
104+
# This test should simply run if there is no regression
105+
106+
num_idxs = 100
107+
assert (
108+
len(
109+
MaterialsRester().search(
110+
material_ids=[f"mp-{idx}" for idx in range(num_idxs)], deprecated=True
111+
)
112+
)
113+
<= num_idxs
114+
)

tests/client/test_mprester.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,10 @@ def test_get_entries(self, mpr):
261261
def test_get_entries_in_chemsys(self, mpr):
262262
syms = ["Li", "Fe", "O"]
263263
syms2 = "Li-Fe-O"
264-
entries = mpr.get_entries_in_chemsys(syms)
264+
with pytest.warns(
265+
MPRestWarning, match="The default thermo type when retrieving entries"
266+
):
267+
entries = mpr.get_entries_in_chemsys(syms)
265268
entries2 = mpr.get_entries_in_chemsys(syms2)
266269
elements = {Element(sym) for sym in syms}
267270
for e in entries:
@@ -326,7 +329,9 @@ def test_get_pourbaix_entries(self, mpr):
326329
reason="`pip install mpcontribs-client` to use pourbaix functionality.",
327330
)
328331
def test_get_ion_entries(self, mpr):
329-
entries = mpr.get_entries_in_chemsys("Ti-O-H")
332+
entries = mpr.get_entries_in_chemsys(
333+
"Ti-O-H", additional_criteria={"thermo_types": ["GGA_GGA+U"]}
334+
)
330335
pd = PhaseDiagram(entries)
331336
ion_entry_data = mpr.get_ion_reference_data_for_chemsys("Ti-O-H")
332337
ion_entries = mpr.get_ion_entries(pd, ion_entry_data)
@@ -338,7 +343,9 @@ def test_get_ion_entries(self, mpr):
338343
assert len(bi_v_entry_data) == len(bi_data + v_data)
339344

340345
# test an incomplete phase diagram
341-
entries = mpr.get_entries_in_chemsys("Ti-O")
346+
entries = mpr.get_entries_in_chemsys(
347+
"Ti-O", additional_criteria={"thermo_types": ["GGA_GGA+U"]}
348+
)
342349
pd = PhaseDiagram(entries)
343350
with pytest.raises(ValueError, match="The phase diagram chemical system"):
344351
mpr.get_ion_entries(pd)
@@ -352,7 +359,8 @@ def test_get_ion_entries(self, mpr):
352359
itertools.chain.from_iterable(i.elements for i in ion_ref_comps)
353360
)
354361
ion_ref_entries = mpr.get_entries_in_chemsys(
355-
[*map(str, ion_ref_elts), "O", "H"]
362+
[*map(str, ion_ref_elts), "O", "H"],
363+
additional_criteria={"thermo_types": ["GGA_GGA+U"]},
356364
)
357365
mpc = MaterialsProjectAqueousCompatibility()
358366
ion_ref_entries = mpc.process_entries(ion_ref_entries)

0 commit comments

Comments
 (0)