Arbitrarily nested dynamic namespace

[ahkole]

I am exploring the new type annotation system/API for declaring the inputs/outputs of tasks (https://aiida-workgraph.readthedocs.io/en/latest/howto/autogen/annotate_inputs_outputs.html#). I was wondering if it's possible to declare a dynamic namespace that can contain an arbitrary number of nested (dynamic) namespaces. There is an example of having a dynamic namespace that contains an arbitrary number of static namespaces here (https://aiida-workgraph.readthedocs.io/en/latest/howto/autogen/annotate_inputs_outputs.html#nested-namespaces) but it's not clear to me if it's possible to have an arbirary nested dynamic namespace. For example, say that you have this task,

from aiida_workgraph import task, dynamic
import typing as t

@task.graph
def some_workgraph(
    nested_namespace: t.Annotated[
        dict,
        dynamic(<what to put here?>),
    ],
) -> <some_output_port_type>:
... some code ...

And then you would be able to call it with an arbitrarily nested dict so that it contains some arbitrarily nested namespace:

wg = some_workgraph.build(
    nested_namespace = {
        'ns1': {'i1': 1, 'f1': 5.0},
        'ns2': {'ns3': {'f2': 10.0}, 'i2': 3,},
        'i3': 20,
    }
)

So that it creates the following input ports:

nested_namespace.ns1.i1
nested_namespace.ns1.f1
nested_namespace.ns2.ns3.f2
nested_namespace.ns2.i2
nested_namespace.i3

Is this possible with the current type annotation system?

[superstar54]

Hi @ahkole , this is not possible with the current type annotation system. In aiida-core, one can declare the entire input as dynamic and then pass any nested data into it. This is convenient, but there will be no validation on the inputs.

What's the use case for this? Do you really want to pass an arbitrarily nested dict as input? In this case, WorkGraph will not know how to serialize the input, so every item in the dict will be stored as an AiiDA node.

[ahkole]

Hi @superstar54 thank you for the quick reply. My main use case was iterating over inputs of a DFT calculation, basically supplying a full namespace for each calculation that I want to iterate over, i.e.

@task.graph
def iterate_dft(
    scf_inputs: t.Annotated[
        dict,
        dynamic(<fully arbitrary dynamic>),
    ],
):
    scf_inputs_dict = namespace_to_dict(scf_inputs)
    for inp in scf_inputs_dict:
        outputs = task(DFTBaseWorkChain)(**inp)

wg = iterate_dft.build(
    scf_inputs={
        'dft1': dict_inputs1,
        'dft2': dict_inputs2,
        'dft3': dict_inputs3,
    }
)

Since some of the DFT input ports are namespaces, using dynamic(t.Any) doesn't work. I can use the data linkage functionality described here (https://aiida-workgraph.readthedocs.io/en/latest/howto/autogen/annotate_inputs_outputs.html#data-linkage) to do something like this,

@task.graph
def iterate_dft(
    scf_inputs: t.Annotated[
        dict,
        dynamic(task(DFTBaseWorkChain).inputs),
    ],
):
    scf_inputs_dict = namespace_to_dict(scf_inputs)
    for inp in scf_inputs_dict:
        outputs = task(DFTBaseWorkChain)(**inp)

That might work but I wanted some flexibility in sometimes leaving out some of the required ports of the DFTBaseWorkChain because they are provided elsewhere in the WorkGraph. But maybe it would be better to learn how to use the meta functionality to mark those ports as optional. How exactly does this work? The example in the docs here (https://aiida-workgraph.readthedocs.io/en/latest/howto/autogen/annotate_inputs_outputs.html#modifying-specification-metadata) only shows how to mark the entire dynamic namespace as optional.

[superstar54]

leaving out some of the required ports of the DFTBaseWorkChain because they are provided elsewhere in the WorkGraph.

In this case, you can use select to exclude the ports that are provided from elsewhere.

Here is another example, where the structure is excluded from the inputs.

I will also make a PR soon to support an arbitrarily nested dynamic namespace.

[ahkole]

leaving out some of the required ports of the DFTBaseWorkChain because they are provided elsewhere in the WorkGraph.

In this case, you can use select to exclude the ports that are provided from elsewhere.

Here is another example, where the structure is excluded from the inputs.

Thanks for the suggestion, I'll give it a try!

I will also make a PR soon to support an arbitrarily nested dynamic namespace.

Looking forward to it!

[superstar54]

Hi @ahkole , aiida-workgraph 0.7.3 is released to support arbitrarily nested dynamic namespace.

If no item type is specified (i.e., just dynamic()), the namespace becomes fully dynamic, enabling arbitrary nested structure (dictionary). Each key-value pair will then be recursively treated as a leaf socket.

[ahkole]

Thanks for the update and the very fast new release!