6
6
from pymatgen .core .periodic_table import Element
7
7
8
8
from mp_api .client .core import BaseRester
9
+ from mp_api .client .core .utils import validate_ids
9
10
10
11
11
12
class JcesrMoleculesRester (BaseRester [MoleculesDoc ]):
@@ -15,6 +16,7 @@ class JcesrMoleculesRester(BaseRester[MoleculesDoc]):
15
16
16
17
def search (
17
18
self ,
19
+ task_ids : str | list [str ] | None = None ,
18
20
charge : tuple [float , float ] | None = None ,
19
21
elements : list [Element ] | None = None ,
20
22
EA : tuple [float , float ] | None = None ,
@@ -30,6 +32,8 @@ def search(
30
32
"""Query equations of state docs using a variety of search criteria.
31
33
32
34
Arguments:
35
+ task_ids (str, List[str]): A single molecule task ID string or list of strings.
36
+ (e.g., mol-45004, [mol-45004, mol-45228]).
33
37
charge (Tuple[float,float]): Minimum and maximum value of the charge in +e to consider.
34
38
elements (List[Element]): A list of elements.
35
39
film_orientation (List[Elements]): List of elements that are in the molecule.
@@ -49,6 +53,12 @@ def search(
49
53
"""
50
54
query_params = defaultdict (dict ) # type: dict
51
55
56
+ if task_ids :
57
+ if isinstance (task_ids , str ):
58
+ task_ids = [task_ids ]
59
+
60
+ query_params .update ({"task_ids" : "," .join (validate_ids (task_ids ))})
61
+
52
62
if elements :
53
63
query_params .update ({"elements" : "," .join ([str (ele ) for ele in elements ])})
54
64
0 commit comments