Add documentation a notebook for aggregating dynamically sized inputs

[agoscinski]

The notebook parallel.ipynb explains this concept using aiida's context but this is not very comprehensible. There is another solution which is still a bit experimental but much more understandable by increasing the link limit

from aiida_workgraph import task
@task.calcfunction
def struc_generator(seed: orm.Int):
    np.random.seed(seed.value)
    return {"result": {f"{j}": {f"{i}": orm.StructureData(cell=np.random.rand(3,3)) for i in range(3)} for j in range(3)}}


@task.calcfunction
def aggregate_strucs(**collected_strucs):
    for key,value in collected_strucs.items():
        print(key,value)
    # we have to clone because aiida does not allow to return identity
    return {"result": orm.Int(0)}


wg = WorkGraph("aggregate")

aggregate_strucs_task = wg.add_task(aggregate_strucs, name="aggregate_strucs_task")

# we have to increase the link limit because by default workgraph only supports one link per input socket
# this is still an experimental feature that is why
aggregate_strucs_task.inputs["collected_strucs"].link_limit = 50

for i in range(2):
    struc_generator_task = wg.add_task(
        struc_generator, name=f"struc_generator{i}", seed=orm.Int(i)
    )
    # We create a `link` between the output of the `query_diag_mean` task and the input of the aggregation and plot task
    wg.add_link(
        struc_generator_task.outputs["result"],
        aggregate_strucs_task.inputs["collected_strucs"],
    )
    
#display(wg.to_html())
wg.run()
aggregate_strucs_task.outputs["result"].value

Passing multiple dynamically sized inputs

We can also include in this notebook the cases when you want to pass two dynamically inputs as discussed here https://aiida.discourse.group/t/aiida-workgraph-how-to-handle-dynamic-inputs-and-outputs-of-tasks/457/4

Passing nested dynamically sized inputs

from aiida.engine import calcfunction
@calcfunction
def struc_generator(seed: orm.Int):
    np.random.seed(seed.value)
    return {f"a{j}": {f"b{i}": orm.StructureData(cell=np.random.rand(3,3)) for i in range(3)} for j in range(3)}

@calcfunction
def aggregate_strucs(**collected_strucs):
    for key,value in collected_strucs.items():
        print(key,value)
    # we have to clone because aiida does not allow to return identity
    return {"result": orm.Int(0)}
outputs = struc_generator(orm.Int(1))
print(outputs)
aggregate_strucs(**outputs)

Passing dynamically sized inputs to filter elements using If workgraph

This I am not yet sure how to do, but the idea is that we use the If workgraph to filter out elements of a list. This is a bit tricky when I tried it because the aggregate task is not executed if the if case is not fulfilled. One could pass a None type and then filter it but I think it will not be presented in the provenance properly.