Use CalcJob and WorkChain insdie WorkGraph

One can use CalcJob, WorkChain and other AiiDA components direclty in the WorkGraph. The inputs and outputs of the task is automatically generated based on the input and output port of the AiiDA component.

from aiida_workgraph import WorkGraph
from aiida.calculations.arithmetic.add import ArithmeticAddCalculation

wg = WorkGraph("test_use_calcjob")
task1 = wg.add_task(ArithmeticAddCalculation, name="add1")
task2 = wg.add_task(ArithmeticAddCalculation, name="add2", x=wg.tasks["add1"].outputs["sum"])
wg.to_html()

Set inputs

One can set the inputs when adding the task, or using the set method of the Task object.

from aiida import load_profile
from aiida.orm import Int

load_profile()

# use set method
task1.set({"x": Int(1), "y": Int(2)})
# set the inputs when adding the task
task3 = wg.add_task(ArithmeticAddCalculation, name="add3", x=Int(3), y=Int(4))

Use process builder

One can also set the inputs of the task using the process builder.

from aiida.calculations.arithmetic.add import ArithmeticAddCalculation
from aiida.orm import Int, load_code


code = load_code("add@localhost")
builder = ArithmeticAddCalculation.get_builder()
builder.code = code
builder.x = Int(2)
builder.y = Int(3)

wg = WorkGraph("test_set_inputs_from_builder")
add1 = wg.add_task(ArithmeticAddCalculation, name="add1")
add1.set_from_builder(builder)

Dynamic namespace

In AiiDA, one can define a dynamic namespace for the process, which allows the user to pass any nested dictionary with AiiDA data nodes as values. However, in the WorkGraph, we need to define the input and output sockets explicitly, so that one can make a link between tasks. To address this discrepancy, and still allow user to pass any nested dictionary with AiiDA data nodes, as well as the output sockets of other tasks, we automatically create the input for each item in the dictionary if the input is not defined. Besides, if the value of the item is a socket, we will link the socket to the task, and remove the item from the dictionary.

For example, the WorkChainWithDynamicNamespace has a dynamic namespace dynamic_port, and the user can pass any nested dictionary as the input.

from aiida.engine import WorkChain

class WorkChainWithDynamicNamespace(WorkChain):
    """WorkChain with dynamic namespace."""

    @classmethod
    def define(cls, spec):
        """Specify inputs and outputs."""
        super().define(spec)
        spec.input_namespace("dynamic_port", dynamic=True)

wg = WorkGraph("test_dynamic_namespace")
task1 = wg.add_task(ArithmeticAddCalculation, name="add1")
task2 = wg.add_task(
        WorkChainWithDynamicNamespace,
        dynamic_port={
            "input1": None,
            "input2": Int(2),
            "input3": task1.outputs["sum"],
        },
    )
wg.to_html()

Failed to inspect function WorkChainWithDynamicNamespace: source code not available

Summary

In this example, we will show how to use CalcJob and WorkChain inside the WorkGraph. One can also use WorkGraph inside a WorkChain, please refer to the Calling WorkGraph within a WorkChain for more details.