Use while loop

Introduction

In this tutorial, you’ll learn how to implement a while loop in WorkGraph, allowing you to execute a set of tasks repeatedly as long as a specified condition remains true.

WorkGraph provides a active_while_zone context manager that allows you to create a while loop in a WorkGraph. Here’s a simple example to use active_while_zone context manager to a WorkGraph:

with active_while_zone(task1.output.result, max_iterations=100):
    wg.add_task(task2, name="task2", a=1, b=2),

Example

Suppose we want to calculate the following workflow, the tasks for each step are shown:

from aiida_workgraph import WorkGraph, task
from aiida import load_profile

load_profile()

@task.calcfunction()
def compare(x, y):
    return x < y

@task.calcfunction()
def add(x, y):
    return x + y

@task.calcfunction()
def multiply(x, y):
    return x*y

#-------------------------------------------------------------------
# start while block
n = add._func(1, 1)            # task add1
# start while loop
while compare._func(n, 50):    # task compare
    n = add._func(n, 1)        # task add2
    n = multiply._func(n, 2)   # task multiply1
# end while block
z = add._func(n, 1)            # task add3
#-------------------------------------------------------------------

print("Result: ", z)

Result:  uuid: 61ab43b2-8e18-4973-b7a2-cd2918d46f6e (pk: 54684) value: 63

Create the workflow

Now, let’use create the workflow using the active_while_zone.

from aiida_workgraph import task, WorkGraph, active_while_zone

with WorkGraph("while_task") as wg:
    # set a context variable before running.
    wg.ctx = {"should_run": True}
    result1 = add(x=1, y=1)
    wg.ctx.n = result1
    #---------------------------------------------------------------------
    # Create the while zone
    with active_while_zone(wg.ctx.n < 50):
        result2 = add(x=wg.ctx.n, y=1)
        result2._waiting_on.add(result1)
        result3 = multiply(x=result2, y=2)
        wg.ctx.n = result3
    #---------------------------------------------------------------------
    result4 = add(x=result3, y=1)
wg.to_html()
# comment out the following line to visualize the graph in the notebook
# wg

GUI Representation

In the graphical user interface, the while_ instruction is depicted as a While Zone, containing all its child tasks. This zone simplifies the visualization of the loop structure:

• Linking Tasks: Unlike regular tasks, the While Zone lacks data input and output sockets. Tasks outside the zone can directly link to those inside, facilitating a seamless workflow integration.

Submit the WorkGraph and check the results

wg.submit(wait=True)
print("State of WorkGraph:   {}".format(wg.state))
print('Result            :   {}'.format(result4.value))

WorkGraph process created, PK: 54685
Process 54685 finished with state: FINISHED
State of WorkGraph:   FINISHED
Result            :   uuid: 7c0eed5d-af77-44bd-bf62-c193b433ea0b (pk: 54721) value: 63

Generate node graph from the AiiDA process,and we can see that when compare1 node outputs False, the workgraph stops.

from aiida_workgraph.utils import generate_node_graph
generate_node_graph(wg.pk)

Using the While Task

Internally, the active_while_zone instruction is used to implement the While task. The While task in WorkGraph is designed to run a set of tasks multiple times within a loop. You can add a While task to your WorkGraph and configure it with specific conditions and tasks to be executed repeatedly.

• max_iterations: Specifies the maximum number of iterations the loop should execute to prevent infinite loops.

• conditions: The loop will terminate when any condition evaluates to False. This input can be the outputs of other tasks, or a list of boolean expressions derived from the context variables.

Here’s a simple example to add a While task to a WorkGraph:

while_task = wg.add_task("While", name="while_loop",
            max_iterations=100)

Adding tasks to the While loop

We can add tasks to the While zone using the add_task method.

# Add a new task to the while zone
while_zone1.add_task(add, name="task1", a=1, b=2)

from aiida_workgraph import task, WorkGraph

wg = WorkGraph("while_task")
# set a context variable before running.
wg.ctx = {"should_run": True}
add1 = wg.add_task(add, name="add1", x=1, y=1)
wg.update_ctx({"n": add1.outputs.result})
#---------------------------------------------------------------------
# Create the while zone
compare1 = wg.add_task(compare, name="compare1", x=wg.ctx.n, y=50)
while_zone1 = wg.add_task("While", max_iterations=100, conditions=compare1.outputs["result"])
while_zone1.waiting_on.add("add1")
# Create the tasks in the while zone.
add2 = while_zone1.add_task(add, name="add2", x=wg.ctx.n, y=1)
multiply1 = while_zone1.add_task(multiply, name="multiply1",
                        x=add2.outputs["result"],
                        y=2)
# update the context variable
wg.update_ctx({"n": multiply1.outputs.result})
#---------------------------------------------------------------------
add3 = wg.add_task(add, name="add3", x=1, y=1)
wg.add_link(multiply1.outputs["result"], add3.inputs["x"])
wg.to_html()
# comment out the following line to visualize the graph in the notebook
# wg

Summary

This tutorial introduces while loops in WorkGraph using active_while_zone context manager.