Efficiently design and manage flexible workflows with AiiDA, featuring an interactive GUI, checkpoints, provenance tracking, and remote execution capabilities.
AiiDA-WorkGraph version 0.3.16. Here is an example with nested while task,
from aiida import load_profile
from aiida_workgraph import task, WorkGraph
load_profile()
# we need a compare task for `n<100`,
# it's a normal function instead of a calcfunction
@task()
def compare(x, y):
return x < y
# define multiply task for n*2
@task.calcfunction()
def multiply(x, y):
return x*y
# define add task for n+3
@task.calcfunction()
def add(x, y):
return x + y
n = 1
wg = WorkGraph("test_while")
# set a context variable before running.
wg.context = {"n": 1,
"should_run": True}
add1 = wg.add_task(add, name="add1", x="{{n}}", y=1)
add1.set_context({"result": "n"})
#---------------------------------------------------------------------
# Create a WorkGraph will repeat itself based on the conditions
# then we output the result of from the context (context)
add2 = wg.add_task(add, name="add2", x="{{n}}", y=1)
add2.wait.append("add1")
multiply1 = wg.add_task(multiply, name="multiply1",
x=add2.outputs["result"],
y=2)
# update the context variable
multiply1.set_context({"result": "n"})
compare1 = wg.add_task(compare, name="compare1", x=multiply1.outputs["result"], y=30)
compare1.set_context({"result": "should_run"})
while1 = wg.add_task("While", name = "while1", max_iterations=100,
conditions=["should_run"],
tasks=["add2", "multiply1", "compare1"],)
# the `result` of compare1 taskis used as condition
#---------------------------------------------------------------------
add3 = wg.add_task(add, name="add3", x=1, y=1)
add3.set_context({"result": "n"})
wg.add_link(multiply1.outputs["result"], add3.inputs["x"])
#---------------------------------------------------------------------
add4 = wg.add_task(add, name="add4", x=1, y=1)
add4.set_context({"result": "n"})
compare2 = wg.add_task(compare, name="compare2", x=add3.outputs["result"], y=100)
compare2.set_context({"result": "should_run"})
while2 = wg.add_task("While", name="while2",
conditions=["should_run"],
tasks=["add1", "while1", "add3", "compare2"],)
while2.wait.append("add4")
add5 = wg.add_task(add, name="add5", x=1, y=add3.outputs["result"])
add5.set_context({"result": "n"})
# the `result` of compare1 taskis used as condition
wg.run()
AiiDA-WorkGraph version 0.3.16. Here is an example with nested while task,
The WorkGraph stucks at this state: