Skip to content

Commit

Permalink
Task API endpoint enhancement (#531)
Browse files Browse the repository at this point in the history
* Fix trajectory frame prop generation

* Add post-process to remove task fields

* Update task hint scheme

* Disable validation on task endpoints

* Add default return value to task doc value pop
  • Loading branch information
munrojm committed Sep 10, 2022
1 parent 6c5361a commit 247a4bf
Show file tree
Hide file tree
Showing 4 changed files with 40 additions and 17 deletions.
13 changes: 9 additions & 4 deletions emmet-api/emmet/api/routes/tasks/hint_scheme.py
Expand Up @@ -8,9 +8,14 @@ class TasksHintScheme(HintScheme):

def generate_hints(self, query):

hints = {"hint": {}}

if query["criteria"] == {}:
hints["hint"]["_id"] = 1
return {"hint": {"_id": 1}}

for param in query["criteria"]:

if "composition_reduced" in param:
return {"hint": {"composition_reduced.$**": 1}}
elif "nelements" in param:
return {"hint": {"nelements": 1}}

return hints
return {"hint": {}}
12 changes: 12 additions & 0 deletions emmet-api/emmet/api/routes/tasks/query_operators.py
Expand Up @@ -31,6 +31,18 @@ def query(

return {"criteria": crit}

def post_process(self, docs, query):
"""
Post processing to remove unwanted fields from all task queries
"""

for doc in docs:
doc.pop("tags", None)
doc.pop("sbxn", None)
doc.pop("dir_name", None)

return docs


class TrajectoryQuery(QueryOperator):
"""
Expand Down
12 changes: 8 additions & 4 deletions emmet-api/emmet/api/routes/tasks/resources.py
@@ -1,3 +1,4 @@
from this import d
from maggma.api.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
from maggma.api.resource import ReadOnlyResource

Expand Down Expand Up @@ -38,7 +39,8 @@ def task_resource(task_store):
header_processor=GlobalHeaderProcessor(),
hint_scheme=TasksHintScheme(),
tags=["Tasks"],
timeout=timeout
timeout=timeout,
disable_validation=True,
)

return resource
Expand All @@ -54,7 +56,7 @@ def task_deprecation_resource(materials_store):
enable_default_search=True,
sub_path="/deprecation/",
header_processor=GlobalHeaderProcessor(),
timeout=timeout
timeout=timeout,
)

return resource
Expand All @@ -69,7 +71,8 @@ def trajectory_resource(task_store):
tags=["Tasks"],
sub_path="/trajectory/",
header_processor=GlobalHeaderProcessor(),
timeout=timeout
timeout=timeout,
disable_validation=True,
)

return resource
Expand All @@ -92,7 +95,8 @@ def entries_resource(task_store):
tags=["Tasks"],
sub_path="/entries/",
header_processor=GlobalHeaderProcessor(),
timeout=timeout
timeout=timeout,
disable_validation=True,
)

return resource
20 changes: 11 additions & 9 deletions emmet-api/emmet/api/routes/tasks/utils.py
Expand Up @@ -22,7 +22,7 @@ def calcs_reversed_to_trajectory(calcs_reversed: List[dict]):

for calculation in calcs_reversed:
structures = []
frame_props = defaultdict(list) # type: dict
frame_props = []

steps = calculation.get("output", {}).get("ionic_steps", None)

Expand All @@ -33,21 +33,23 @@ def calcs_reversed_to_trajectory(calcs_reversed: List[dict]):
else:
for step in steps:

step_dict = {}

structure_dict = step.get("structure", None)

if structure_dict is not None:
structure = Structure.from_dict(structure_dict)

structures.append(structure)

frame_props["e_fr_energy"].append(step.get("e_fr_energy", None))
frame_props["e_wo_entrp"].append(step.get("e_wo_entrp", None))
frame_props["e_0_energy"].append(step.get("e_0_energy", None))
frame_props["forces"].append(step.get("forces", None))
frame_props["stresses"].append(step.get("stress", None))
frame_props["electronic_steps"].append(
step.get("electronic_steps", None)
)
step_dict["e_fr_energy"] = step.get("e_fr_energy", None)
step_dict["e_wo_entrp"] = step.get("e_wo_entrp", None)
step_dict["e_0_energy"] = step.get("e_0_energy", None)
step_dict["forces"] = step.get("forces", None)
step_dict["stresses"] = step.get("stress", None)
step_dict["electronic_steps"] = step.get("electronic_steps", None)

frame_props.append(step_dict)

traj = Trajectory.from_structures(
structures, frame_properties=frame_props, time_step=None
Expand Down

0 comments on commit 247a4bf

Please sign in to comment.