Skip to content

Commit e3a22b7

Browse files
GeigerJ2GeigerJ2
committed
Fix aiida_to_jobflow_simple 
1 parent a5b5098 commit e3a22b7

File tree

4 files changed

+613
-57
lines changed

4 files changed

+613
-57
lines changed

‎.github/workflows/aiida.yml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ jobs:
2929
papermill universal_simple_to_aiida.ipynb universal_simple_to_aiida_out.ipynb -k "python3"
3030
papermill universal_qe_to_aiida.ipynb universal_qe_to_aiida_out.ipynb -k "python3"
3131
32-
papermill aiida_to_pyiron_base_simple.ipynb aiida_to_pyiron_base_simple_out.ipynb -k "python3"
33-
papermill aiida_to_pyiron_base_qe.ipynb pyiron_base_to_jobflow_qe_out.ipynb -k "python3"
32+
: # papermill aiida_to_pyiron_base_simple.ipynb aiida_to_pyiron_base_simple_out.ipynb -k "python3"
33+
: # papermill aiida_to_pyiron_base_qe.ipynb pyiron_base_to_jobflow_qe_out.ipynb -k "python3"
3434
3535
papermill aiida_to_jobflow_simple.ipynb aiida_to_jobflow_simple_out.ipynb -k "python3"
36-
papermill aiida_to_jobflow_qe.ipynb aiida_to_jobflow_qe_out.ipynb -k "python3"
36+
: # papermill aiida_to_jobflow_qe.ipynb aiida_to_jobflow_qe_out.ipynb -k "python3"

‎aiida_to_jobflow_simple.ipynb‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 7,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"from python_workflow_definition.aiida import write_workflow_json\n",
10+
"from python_workflow_definition.jobflow import load_workflow_json\n",
11+
"from aiida_workgraph import task, WorkGraph\n",
12+
"from jobflow.managers.local import run_locally\n",
13+
"\n",
14+
"from aiida import load_profile\n",
15+
"load_profile()\n",
16+
"\n",
17+
"workflow_json_filename = \"workflow_simple_aiida.json\""
18+
]
19+
},
20+
{
21+
"cell_type": "code",
22+
"execution_count": 8,
23+
"metadata": {},
24+
"outputs": [],
25+
"source": [
26+
"from simple_workflow import (\n",
27+
" add_x_and_y as _add_x_and_y, \n",
28+
" add_x_and_y_and_z as _add_x_and_y_and_z,\n",
29+
")"
30+
]
31+
},
32+
{
33+
"cell_type": "code",
34+
"execution_count": 9,
35+
"metadata": {},
36+
"outputs": [
37+
{
38+
"name": "stdout",
39+
"output_type": "stream",
40+
"text": [
41+
"outputs [{'identifier': 'workgraph.any', 'name': 'result'}]\n"
42+
]
43+
},
44+
{
45+
"data": {
46+
"application/vnd.jupyter.widget-view+json": {
47+
"model_id": "826872d2011e4daaaf3f173b38bf49dc",
48+
"version_major": 2,
49+
"version_minor": 1
50+
},
51+
"text/plain": [
52+
"NodeGraphWidget(settings={'minimap': True}, style={'width': '90%', 'height': '600px'}, value={'name': 'wg-simp…"
53+
]
54+
},
55+
"execution_count": 9,
56+
"metadata": {},
57+
"output_type": "execute_result"
58+
}
59+
],
60+
"source": [
61+
"def pickle_node(value):\n",
62+
" \"\"\"Handle data nodes\"\"\"\n",
63+
" return value\n",
64+
"\n",
65+
"helper_1 = task.pythonjob()(pickle_node)\n",
66+
"helper_2 = task.pythonjob()(pickle_node)\n",
67+
"\n",
68+
"add_x_and_y = task.pythonjob(outputs=[\"x\", \"y\", \"z\"])(_add_x_and_y)\n",
69+
"add_x_and_y_and_z = task.pythonjob()(_add_x_and_y_and_z)\n",
70+
"\n",
71+
"# workgraph = write_workflow_json(filename='workflow_simple.json')\n",
72+
"\n",
73+
"# TODO: Create inputs rather than tasks out of data nodes\n",
74+
"wg = WorkGraph('wg-simple')\n",
75+
"\n",
76+
"helper_task1 = wg.add_task(\n",
77+
" helper_1,\n",
78+
" name=\"x\",\n",
79+
" value=1\n",
80+
")\n",
81+
"\n",
82+
"helper_task2 = wg.add_task(\n",
83+
" helper_2,\n",
84+
" name=\"y\",\n",
85+
" value=2\n",
86+
")\n",
87+
"\n",
88+
"add_x_and_y_task = wg.add_task(\n",
89+
" add_x_and_y,\n",
90+
" name='add_x_and_y',\n",
91+
" x=helper_task1.outputs.result,\n",
92+
" y=helper_task2.outputs.result,\n",
93+
")\n",
94+
"\n",
95+
"add_x_and_y_and_z_task = wg.add_task(\n",
96+
" add_x_and_y_and_z,\n",
97+
" name='add_x_and_y_and_z',\n",
98+
" x=add_x_and_y_task.outputs.x,\n",
99+
" y=add_x_and_y_task.outputs.y,\n",
100+
" z=add_x_and_y_task.outputs.z,\n",
101+
")\n",
102+
"\n",
103+
"wg"
104+
]
105+
},
106+
{
107+
"cell_type": "code",
108+
"execution_count": 10,
109+
"metadata": {},
110+
"outputs": [
111+
{
112+
"data": {
113+
"text/plain": [
114+
"{'nodes': {0: 1,\n",
115+
" 1: 2,\n",
116+
" 2: 'simple_workflow.add_x_and_y',\n",
117+
" 3: 'simple_workflow.add_x_and_y_and_z'},\n",
118+
" 'edges': [{'source': 0,\n",
119+
" 'target': 2,\n",
120+
" 'sourceHandle': None,\n",
121+
" 'targetHandle': 'x'},\n",
122+
" {'source': 1, 'target': 2, 'sourceHandle': None, 'targetHandle': 'y'},\n",
123+
" {'source': 2, 'target': 3, 'sourceHandle': 'x', 'targetHandle': 'x'},\n",
124+
" {'source': 2, 'target': 3, 'sourceHandle': 'y', 'targetHandle': 'y'},\n",
125+
" {'source': 2, 'target': 3, 'sourceHandle': 'z', 'targetHandle': 'z'}]}"
126+
]
127+
},
128+
"execution_count": 10,
129+
"metadata": {},
130+
"output_type": "execute_result"
131+
}
132+
],
133+
"source": [
134+
"write_workflow_json(wg=wg, file_name=workflow_json_filename)\n"
135+
]
136+
},
137+
{
138+
"cell_type": "code",
139+
"execution_count": 11,
140+
"metadata": {},
141+
"outputs": [],
142+
"source": [
143+
"\n",
144+
"flow = load_workflow_json(file_name=workflow_json_filename)"
145+
]
146+
},
147+
{
148+
"cell_type": "code",
149+
"execution_count": 12,
150+
"metadata": {},
151+
"outputs": [
152+
{
153+
"name": "stdout",
154+
"output_type": "stream",
155+
"text": [
156+
"2025-03-15 11:14:35,722 INFO Started executing jobs locally\n",
157+
"2025-03-15 11:14:35,874 INFO Starting job - add_x_and_y (0ea991be-405e-4c9f-8ac7-0c34e7b40d03)\n",
158+
"2025-03-15 11:14:35,878 INFO Finished job - add_x_and_y (0ea991be-405e-4c9f-8ac7-0c34e7b40d03)\n",
159+
"2025-03-15 11:14:35,879 INFO Starting job - add_x_and_y_and_z (0a4613e0-432b-49fe-b0cc-48a057667532)\n",
160+
"2025-03-15 11:14:35,881 INFO Finished job - add_x_and_y_and_z (0a4613e0-432b-49fe-b0cc-48a057667532)\n",
161+
"2025-03-15 11:14:35,882 INFO Finished executing jobs locally\n"
162+
]
163+
},
164+
{
165+
"data": {
166+
"text/plain": [
167+
"{'0ea991be-405e-4c9f-8ac7-0c34e7b40d03': {1: Response(output={'x': 1, 'y': 2, 'z': 3}, detour=None, addition=None, replace=None, stored_data=None, stop_children=False, stop_jobflow=False, job_dir=PosixPath('/home/geiger_j/aiida_projects/adis/git-repos/python-workflow-definition'))},\n",
168+
" '0a4613e0-432b-49fe-b0cc-48a057667532': {1: Response(output=6, detour=None, addition=None, replace=None, stored_data=None, stop_children=False, stop_jobflow=False, job_dir=PosixPath('/home/geiger_j/aiida_projects/adis/git-repos/python-workflow-definition'))}}"
169+
]
170+
},
171+
"execution_count": 12,
172+
"metadata": {},
173+
"output_type": "execute_result"
174+
}
175+
],
176+
"source": [
177+
"result = run_locally(flow)\n",
178+
"result"
179+
]
180+
}
181+
],
182+
"metadata": {
183+
"kernelspec": {
184+
"display_name": "ADIS",
185+
"language": "python",
186+
"name": "python3"
187+
},
188+
"language_info": {
189+
"codemirror_mode": {
190+
"name": "ipython",
191+
"version": 3
192+
},
193+
"file_extension": ".py",
194+
"mimetype": "text/x-python",
195+
"name": "python",
196+
"nbconvert_exporter": "python",
197+
"pygments_lexer": "ipython3",
198+
"version": "3.10.12"
199+
}
200+
},
201+
"nbformat": 4,
202+
"nbformat_minor": 2
203+
}

‎python_workflow_definition/src/python_workflow_definition/aiida.py‎

Lines changed: 57 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from aiida_workgraph import WorkGraph, task
44
import json
55

6+
67
def pickle_node(value):
78
"""Handle data nodes"""
89
return value
@@ -12,51 +13,8 @@ def get_item(data: dict, key: str):
1213
"""Handle get item from the outputs"""
1314
return data[key]
1415

15-
# I defined the function mapping here
16-
# but in principle, this is not needed, because one can import the function from the module path
17-
# func_mapping = {
18-
# "simple_workflow.add_x_and_y": task.pythonjob()(add_x_and_y),
19-
# "simple_workflow.add_x_and_y_and_z": task.pythonjob()(add_x_and_y_and_z),
20-
# "quantum_espresso_workflow.get_bulk_structure": task.pythonjob()(
21-
# get_bulk_structure
22-
# ),
23-
# "quantum_espresso_workflow.calculate_qe": task.pythonjob()(calculate_qe),
24-
# "quantum_espresso_workflow.plot_energy_volume_curve": task.pythonjob()(
25-
# plot_energy_volume_curve
26-
# ),
27-
# "python_workflow_definition.pyiron_base.get_dict": task.pythonjob()(get_dict),
28-
# "python_workflow_definition.pyiron_base.get_list": task.pythonjob()(get_list),
29-
# "quantum_espresso_workflow.generate_structures": task.pythonjob()(
30-
# generate_structures
31-
# ),
32-
# }
33-
34-
35-
def write_workflow_json(wg):
36-
wgdata = wg.to_dict()
37-
data = {"nodes": {}, "edges": []}
38-
node_name_mapping = {}
39-
i = 0
40-
for name, node in wgdata["tasks"].items():
41-
node_name_mapping[name] = i
42-
callable_name = node["executor"]["callable_name"]
43-
data["nodes"][i] = callable_name
44-
if callable_name == "pickle_node":
45-
data["nodes"][i] = node["inputs"]["value"]["property"]["value"].value
46-
i += 1
47-
48-
for link in wgdata["links"]:
49-
if wgdata["tasks"][link["from_node"]]["executor"]["callable_name"] == "pickle_node":
50-
link["from_socket"] = None
51-
link["from_node"] = node_name_mapping[link["from_node"]]
52-
link["to_node"] = node_name_mapping[link["to_node"]]
53-
data["edges"].append(link)
54-
55-
return data
56-
5716

5817
def load_workflow_json(filename):
59-
6018
with open(filename) as f:
6119
data = json.load(f)
6220

@@ -68,7 +26,7 @@ def load_workflow_json(filename):
6826
for name, identifier in data["nodes"].items():
6927
# if isinstance(identifier, str) and identifier in func_mapping:
7028
if isinstance(identifier, str) and "." in identifier:
71-
p, m = identifier.rsplit('.', 1)
29+
p, m = identifier.rsplit(".", 1)
7230
mod = import_module(p)
7331
_func = getattr(mod, m)
7432
func = task.pythonjob()(_func)
@@ -126,8 +84,62 @@ def load_workflow_json(filename):
12684
print("Failed to link", link, "with error:", e)
12785
return wg
12886

87+
12988
def get_list(**kwargs):
13089
return list(kwargs.values())
13190

91+
13292
def get_dict(**kwargs):
133-
return {k: v for k, v in kwargs.items()}
93+
return {k: v for k, v in kwargs.items()}
94+
95+
96+
def write_workflow_json(wg, file_name):
97+
wgdata = wg.to_dict()
98+
data = {"nodes": {}, "edges": []}
99+
node_name_mapping = {}
100+
i = 0
101+
for name, node in wgdata["tasks"].items():
102+
node_name_mapping[name] = i
103+
104+
callable_name = node["executor"]["callable_name"]
105+
106+
if callable_name == "pickle_node":
107+
data["nodes"][i] = node["inputs"]["sockets"]["value"]["property"][
108+
"value"
109+
].value
110+
111+
else:
112+
113+
callable_name = f"{node['executor']['module_path']}.{callable_name}"
114+
115+
data["nodes"][i] = callable_name
116+
117+
i += 1
118+
119+
for link in wgdata["links"]:
120+
if (
121+
wgdata["tasks"][link["from_node"]]["executor"]["callable_name"]
122+
== "pickle_node"
123+
):
124+
link["from_socket"] = None
125+
link["source"] = node_name_mapping[link["from_node"]]
126+
del link['from_node']
127+
link["target"] = node_name_mapping[link["to_node"]]
128+
del link['to_node']
129+
link["sourceHandle"] = link.pop("from_socket")
130+
link["targetHandle"] = link.pop("to_socket")
131+
data["edges"].append(link)
132+
133+
with open(file_name, "w") as f:
134+
# json.dump({"nodes": data[], "edges": edges_new_lst}, f)
135+
json.dump(data, f)
136+
137+
return data
138+
139+
140+
# def construct_wg_simple():
141+
142+
# ...
143+
144+
# def construct_qe_simple():
145+
# ...

0 commit comments

Comments
 (0)