Use annotations to control data provenance

Introduction

In aiida-workgraph, a critical feature is tracking task inputs and outputs to ensure data provenance and reproducibility. To achieve this, task functions must be annotated with input and output specifications. This tells the WorkGraph how to handle, serialize, and store data as individual AiiDA nodes.

This process addresses two key aspects of provenance:

• Data creation: How should data be created and stored? For example, if a task returns a nested dictionary, should it be stored as a single entity or unpacked into separate nodes?

• Data lineage: Where does the data come from? How does it flow between tasks in the workflow?

This guide will walk you through the various ways to annotate your tasks—using both the native annotation helpers and Pydantic models.

import typing as t

from aiida import load_profile

from aiida_workgraph import dynamic, namespace, task
from aiida_workgraph.utils import get_process_summary

load_profile()

Profile<uuid='2c6af82ef2b241408a4c27c7c32d5490' name='presto'>

Data creation

Data creation is controlled at the Calculation level. In AiiDA, a Calculation is a process that performs a computation and creates new data nodes.

Static namespaces for outputs

Let’s define two tasks that perform the same calculation but have different output annotations.

The first task, add_multiply1, has no output specification. Consequently, the returned dictionary will be stored as a single AiiDA Dict node. The second task, add_multiply2, uses an output namespace to specify that each key-value pair in the dictionary should be stored as a separate AiiDA node.

@task
def add_multiply1(x, y):
    """Return a dictionary, which will be stored as a single Dict node."""
    return {'sum': x + y, 'product': x * y}


@task
def add_multiply2(
    x: int,
    y: int,
) -> t.Annotated[
    dict,
    namespace(sum=int, product=int),
]:
    """Return a dictionary, but its elements are stored as separate Int nodes."""
    return {'sum': x + y, 'product': x * y}


@task.graph
def AddMultiply(x: int, y: int):
    """A graph to run both versions of the add_multiply task."""
    add_multiply1(x=x, y=y)
    add_multiply2(x=x, y=y)


wg = AddMultiply.build(x=1, y=2)
wg.run()

{}

Note

In add_multiply2, we also annotated the input types (x: int, y: int). This adds a layer of validation to ensure only integers are passed to the task. This feature is experimental - its API and behavior may change in future releases.

Let’s visualize the data provenance of our executed workflow:

wg.generate_provenance_graph()

As the provenance graph shows, add_multiply1 has a single output node (result), while add_multiply2 has two separate output nodes (sum and product), as defined in its namespace.

Static namespaces for inputs

Similarly, we can annotate inputs to unpack a dictionary into distinct data nodes. We use Python’s standard typing.Annotated to attach the aiida-workgraph namespace metadata to the input type.

@task
def add_multiply3(
    data: t.Annotated[
        dict,
        namespace(x=int, y=int),
    ],
) -> t.Annotated[
    dict,
    namespace(sum=int, product=int),
]:
    """Take a dictionary as input but treat 'x' and 'y' as separate nodes."""
    return {'sum': data['x'] + data['y'], 'product': data['x'] * data['y']}


@task.graph
def AddMultiplyInputs(x: int, y: int):
    add_multiply3(data={'x': x, 'y': y})


wg = AddMultiplyInputs.build(x=1, y=2)
wg

Note how the x input is passed to data.x (and similarly for y). This is due to the namespace specifications.

wg.run()

{}

Finally, we can inspect the provenance graph for this workflow:

wg.generate_provenance_graph()

We can see that even though we passed the inputs as a single dictionary, they were serialized as two separate Int nodes, x and y, before being passed to the task.

Dynamic namespaces

Sometimes, the number and names of outputs are not known until the task runs. Dynamic namespaces are designed for this scenario.

This is particularly useful for tasks that generate a variable number of outputs based on their inputs.

@task
def generate_square_numbers(
    n: int,
) -> t.Annotated[
    dict,
    dynamic(t.Any),
]:
    """Generate a dict of square numbers. The number of outputs depends on 'n'."""
    return {f'square_{i}': i**2 for i in range(n)}


@task.graph
def SquareNumbersGenerator(n: int):
    generate_square_numbers(n=n)


wg = SquareNumbersGenerator.build(n=5)
wg.run()

{}

Let’s examine the provenance of this dynamic workflow:

wg.generate_provenance_graph()

The graph shows that the generate_square_numbers task has multiple output nodes, one for each entry in the dynamically generated dictionary. The dynamic(typing.Any) specification instructs the workgraph to treat each value in the returned dictionary as a separate output node of any type.

Note

If no item type is specified (i.e., just dynamic()), the namespace becomes fully dynamic, enabling arbitrary nested structures (e.g., dictionaries). Each value will then be recursively treated as a leaf node.

Nested namespaces

Namespaces can be nested to represent complex, structured data. Let’s define a task that returns a nested dictionary.

@task
def generate_nested_dict(
    x: int,
    y: int,
) -> t.Annotated[
    dict,
    namespace(sum=int, nested=namespace(diff=int, product=int)),
]:
    """Returns a nested dictionary with a corresponding nested namespace."""
    return {'sum': x + y, 'nested': {'diff': x - y, 'product': x * y}}


@task.graph
def NestedDictGenerator(x: int, y: int):
    generate_nested_dict(x=x, y=y)


wg = NestedDictGenerator.build(x=1, y=2)
wg.run()

{}

Instead of visualizing the full graph, let’s inspect the outputs of the task using a summary utility.

print(get_process_summary(wg.tasks[-1].pk))

Outputs      PK    Type
-----------  ----  ------
nested
    diff     502   Int
    product  503   Int
sum          501   Int

The summary confirms that the output is correctly structured with a top-level sum and a nested nested dictionary, just as defined in the namespace.

We can also combine dynamic and nested namespaces.

@task
def generate_dynamic_nested_dict(
    n: int,
) -> t.Annotated[
    dict,
    dynamic(namespace(square=int, cube=int)),
]:
    """Generate a nested dict of square and cube numbers from 0 to n."""
    return {f'data_{i}': {'square': i**2, 'cube': i**3} for i in range(n)}


@task.graph
def DynamicNestedDictGenerator(n: int):
    generate_dynamic_nested_dict(n=n)


wg = DynamicNestedDictGenerator.build(n=3)
wg.run()

{}

Let’s check the output summary for this dynamically generated nested structure:

print(get_process_summary(wg.tasks[-1].pk))

Outputs     PK    Type
----------  ----  ------
data_0
    square  507   Int
    cube    508   Int
data_1
    square  509   Int
    cube    510   Int
data_2
    square  511   Int
    cube    512   Int

The output shows a dictionary with dynamic keys (data_0, data_1, etc.), where each value is itself a dictionary with a fixed square and cube structure, as specified by dynamic(namespace(...)).

Using TypedDict

TypedDict works like a static namespace with one socket per key. Optional keys (total=False) become optional sockets.

from typing import TypedDict


class XYIn(TypedDict):
    x: int
    y: int


class AddMultiplyOut(TypedDict):
    sum: int
    product: int


@task
def add_multiply_typed_dict(data: XYIn) -> AddMultiplyOut:
    return {'sum': data['x'] + data['y'], 'product': data['x'] * data['y']}


@task.graph
def AddMultiplyTypedDict():
    add_multiply_typed_dict(data={'x': 2, 'y': 5})


wg = AddMultiplyTypedDict.build()
wg.run()
wg.generate_provenance_graph()

Note

You can also expand **kwargs with Unpack[YourTypedDict] to make the inputs explicit in the function signature (same socket expansion). Optional keys can be expressed with total=False or NotRequired. NotRequired requires Python 3.11+ (or typing_extensions).

Using Pydantic models

You can use Pydantic models in annotations as a more structured, reusable way to define namespaces. By default, a BaseModel expands to a static namespace with one socket per field.

If you want a dynamic namespace, set model_config = {"extra": "allow"} and (optionally) "item_type" for the type of each dynamic value. If you want to treat a model as a single leaf (blob), set model_config = {"leaf": True} or use Leaf[YourModel] in the annotation.

from pydantic import BaseModel
from aiida_workgraph.socket_spec import Leaf


class InputsModel(BaseModel):
    x: int
    y: int


class OutputsModel(BaseModel):
    sum: int
    product: int


@task
def add_multiply_pydantic_in_out(data: InputsModel) -> OutputsModel:
    return OutputsModel(sum=data.x + data.y, product=data.x * data.y)


@task.graph
def AddMultiplyPydantic(data: InputsModel) -> OutputsModel:
    return add_multiply_pydantic_in_out(data=data)


wg = AddMultiplyPydantic.build(data=InputsModel(x=3, y=4))
wg.run()
wg.generate_provenance_graph()

Dynamic Pydantic models

Mark a model as dynamic with extra='allow'. Add item_type to specify the per-key value type. Fixed fields still appear as normal sockets alongside your dynamic keys.

class DynamicOut(BaseModel):
    model_config = {'extra': 'allow', 'item_type': int}

    header: int = 42  # fixed (non-dynamic) field


@task
def make_dynamic_with_model(n: int) -> DynamicOut:
    # fixed field + dynamic keys with int values
    payload = {'header': 100}
    payload.update({f'k{i}': i * i for i in range(n)})
    return DynamicOut(**payload)


@task.graph
def GraphDynamicOut(n: int):
    make_dynamic_with_model(n=n)


wg = GraphDynamicOut.build(n=4)
wg.run()
wg.generate_provenance_graph()

Leaf Pydantic models (single blob)

Sometimes you want to validate with a Pydantic model but store it as a single node instead of expanding fields. For leaf models, WorkGraph treats the value as a blob. If a serializer is registered for the class, it is used; otherwise JsonableData stores the model as a JSON-friendly dict.

Note: Pydantic annotations define the schema for WorkGraph sockets. Outputs are stored as typed AiiDA nodes per field, so runtime results are dicts of nodes (not Pydantic instances). Use those nodes for linking/provenance; rebuild a Pydantic model only for convenience. There are two ways:

1. Mark the model: model_config = {"leaf": True}

2. Per-use override: annotate with Leaf[YourModel]

class BlobModel(BaseModel):
    model_config = {'leaf': True}  # always a leaf blob

    a: int
    b: int


@task
def consume_blob(m: BlobModel) -> dict:
    # 'm' is validated by Pydantic but stored/treated as one leaf node
    return {'sum': m.a + m.b}


# Per-use override without modifying the model:
class AnotherModel(BaseModel):
    a: int
    b: int


@task
def consume_blob_per_use(m: Leaf[AnotherModel]) -> dict:
    return {'sum': m.a + m.b}


@task.graph
def BlobExamples():
    consume_blob(m=BlobModel(a=1, b=2))
    consume_blob_per_use(m=AnotherModel(a=3, b=4))


wg = BlobExamples.build()
wg.run()
wg.generate_provenance_graph()

Nested Pydantic models

Pydantic models can be nested to represent complex data structures. You can mix static, dynamic, and leaf models as needed.

class InnerModel(BaseModel):
    value: int


class OuterModel(BaseModel):
    name: str
    inner: InnerModel

Using dataclasses

Dataclasses work just like Pydantic models for annotations:

• Plain dataclass -> expanded static namespace (one socket per field)

• model_config={‘extra’: ‘allow’, ‘item_type’: T} -> dynamic namespace

• model_config={‘leaf’: True} or Leaf[YourDataclass] -> single leaf (blob)

from dataclasses import dataclass


@dataclass
class DCOutputs:
    sum: int
    product: int


@task
def add_multiply_dc_in_out(x, y) -> DCOutputs:
    return {'sum': x + y, 'product': x * y}


@task.graph
def AddMultiplyDataclass():
    # IMPORTANT: pass a plain dict, not DCInputs(...)
    add_multiply_dc_in_out(x=2, y=5)


wg = AddMultiplyDataclass.build()
wg.run()
wg.generate_provenance_graph()

Important

Structured models (Pydantic or dataclasses) are supported as runtime values. You may pass instances to tasks/graphs and return them from tasks:

• Instances are expanded to plain dicts when assigned to namespace sockets, so WorkGraph can wire provenance edges precisely (e.g., data.x –> task.data.x).

• Graph inputs can still be collected from task outputs as a dict of AiiDA ORM nodes, preserving AiiDA links between nodes.

• Validation still happens via the WorkGraph spec (derived from your annotations).

Data linkage

Data linkage tracks the flow of data between tasks. At the workflow level, a task.graph can define its own inputs and outputs, providing a clean interface to a complex chain of tasks. aiida-workgraph validates data against these graph-level specifications and automatically links graph inputs to the appropriate task inputs.

In this final example, we will build a graph that reuses the input and output specifications from the tasks it contains. This is a powerful feature for building complex, modular, and self-consistent workflows.

@task
def add_multiply(
    data: t.Annotated[
        dict,
        namespace(x=int, y=int),
    ],
) -> t.Annotated[
    dict,
    namespace(sum=int, product=int),
]:
    """A reusable task with well-defined I/O specifications."""
    return {'sum': data['x'] + data['y'], 'product': data['x'] * data['y']}


@task.graph
def AddMultiplyFinal(
    n: int,
    data: t.Annotated[
        dict,
        namespace(
            add_multiply1=add_multiply.inputs,
            add_multiply2=add_multiply.inputs,
        ),
    ],
) -> t.Annotated[
    dict,
    namespace(
        square=generate_square_numbers.outputs,
        add_multiply1=add_multiply.outputs,
        add_multiply2=add_multiply.outputs,
    ),
]:
    """A complex graph demonstrating I/O reuse and data linkage."""
    square_numbers = generate_square_numbers(n)

    # Unpack nested inputs and pass them to the respective tasks
    out1 = add_multiply(data=data['add_multiply1']['data'])
    out2 = add_multiply(data=data['add_multiply2']['data'])

    # Gather task outputs into the graph-level output structure
    return {'square': square_numbers, 'add_multiply1': out1, 'add_multiply2': out2}


wg = AddMultiplyFinal.build(
    n=3,
    data={
        'add_multiply1': {'data': {'x': 1, 'y': 2}},
        'add_multiply2': {'data': {'x': 3, 'y': 4}},
    },
)
wg

In the example above:

• Graph outputs: The outputs are annotated with a nested namespace that defines the shape of the final result. Here we reuse generate_square_numbers.outputs and add_multiply.outputs to ensure the graph’s output signature is consistent with the tasks it contains.

• Graph inputs: The data input is annotated with a nested namespace that reuses add_multiply.inputs. This allows aiida-workgraph to validate the complex input dictionary and create the correct data links.

In the GUI representation of the WorkGraph, you will see how the nested inputs are correctly wired. For instance, there is a direct link from the graph input socket data.add_multiply1.data.x to the task input socket add_multiply_task_1.data.x, guaranteeing perfect data lineage.

Tip

If a graph only exposes the outputs of a single task, this can be simplified as

@task.graph
def SomeGraph(...) - t.Annotated[dict, some_task.outputs]:
    return some_task(...)

We can see similar linkage in the provenance graph. Let’s run the graph and visualize its provenance.

wg.run()
wg.generate_provenance_graph()

Note how the outputs of the various tasks are exposed (linked) to the graph, making accessible via the graph node.

Reshaping specifications with select

Often, you’ll want to reuse parts of a specification while modifying it. For example, you might want to exclude a field that is provided by another source or rename fields for clarity. This can be done declaratively using select.

The main parameters are:

• include=... / exclude=...: Keep or drop fields (supports dotted paths for nested fields).

• include_prefix=... / exclude_prefix=...: Filter top-level fields by name prefix.

• rename={old:new}: Rename top-level fields.

• prefix="p_": Add a prefix to all top-level fields.

Important

To apply a select, you must place it in the metadata list of t.Annotated[<type>, ...] alongside the specification you are modifying.

Let’s see an example where we build a graph that runs a task twice. We want a shared structure input at the graph level, so we must exclude it from the nested input specifications for each task call.

from aiida_workgraph.socket_spec import meta, select


@task
def consume_complex(
    data: t.Annotated[
        dict,
        namespace(
            pw=namespace(structure=int, kpoints=int, parameters=int),
            metadata=dict,
        ),
    ],
) -> dict:
    return {'seen': list(sorted(data.keys()))}


# Exclude a nested field (drop data.pw.structure)
# because it's provided separately as a graph input
# and shared between the two task calls.
@task.graph
def UseExclude(
    structure,
    inputs: t.Annotated[
        dict,
        namespace(
            consume_complex1=t.Annotated[dict, consume_complex.inputs, select(exclude='data.pw.structure')],
            consume_complex2=t.Annotated[dict, consume_complex.inputs, select(exclude='data.pw.structure')],
        ),
    ],
):
    # Manually add the shared 'structure' to each task's inputs
    consume_complex_input1 = inputs['consume_complex1']
    consume_complex_input1['data']['pw']['structure'] = structure
    consume_complex(**consume_complex_input1)

    consume_complex_input2 = inputs['consume_complex2']
    consume_complex_input2['data']['pw']['structure'] = structure
    consume_complex(**consume_complex_input2)


wg = UseExclude.build(
    structure=1,
    inputs={
        'consume_complex1': {'data': {'pw': {'kpoints': 2, 'parameters': 3}, 'metadata': {}}},
        'consume_complex2': {'data': {'pw': {'kpoints': 4, 'parameters': 5}, 'metadata': {}}},
    },
)
wg

In the GUI representation, you can see that the graph has a top-level structure input, and the nested inputs inputs.consume_complex1.data.pw and inputs.consume_complex2.data.pw are missing the structure socket, just as we specified.

Modifying specification metadata

You can also set namespace-level metadata declaratively using meta, for example, to mark a reused specification as optional.

• meta(required=False): Makes the input optional.

• meta(is_metadata=True): Marks the input as metadata-only.

It is attached alongside select in the same Annotated metadata list.

@task.graph
def UseMeta(
    data: t.Annotated[
        dict,
        consume_complex.inputs,  # Reuse the original spec
        meta(required=False),  # Make the entire 'data' input optional
    ],
):
    if data:
        return consume_complex(data=data)

Attaching extra info to data nodes

Sometimes you need to attach extra information to the data a task creates, like a physical unit or a descriptive label. You can do this directly in the output annotation by passing a dictionary to meta(extras=...). This populates the .extras attribute of the resulting AiiDA node, perfect for any custom, JSON-serializable metadata. Let’s see an example where we calculate the potential energy of a crystal structure and attach its unit to the output node.

from ase import Atoms
from ase.build import bulk


@task()
def calc_energy(atoms: Atoms) -> t.Annotated[float, meta(extras={'unit': 'eV'})]:
    """Calculates the potential energy and attaches metadata to the output."""
    from ase.calculators.emt import EMT

    atoms.set_calculator(EMT())
    return atoms.get_potential_energy()


@task.graph()
def MyWorkflow(atoms: Atoms):
    """A simple workflow to run the energy calculation."""
    return calc_energy(atoms).result


wg = MyWorkflow.build(atoms=bulk('Cu'))
wg.run()

/home/docs/checkouts/readthedocs.org/user_builds/aiida-workgraph/checkouts/latest/docs/gallery/howto/autogen/annotate_inputs_outputs.py:680: FutureWarning: Please use atoms.calc = calc
  atoms.set_calculator(EMT())

{'result': <Float: uuid: 8b9da1a3-d98d-48de-a989-2e4aaebb9b7d (pk: 565) value: -0.0056815113585813>}

Now, let’s inspect the output AiiDA node to confirm that our metadata was correctly attached.

output_node = wg.outputs.result.value
energy_unit = output_node.base.extras.get('unit')
print(f'Calculated energy: {output_node.value} {energy_unit}')

Calculated energy: -0.0056815113585813 eV

Conclusion

You now know how to annotate task and graph inputs and outputs in aiida-workgraph. By leveraging static (namespace), dynamic (dynamic), nested namespaces, and structured models, you can precisely control data serialization and create transparent data lineages.

The key takeaways are:

• Annotate task/graph outputs to unpack results into individual AiiDA nodes.

• Annotate inputs to specify input structures.

• Employ dynamic (or Pydantic models with extra='allow') for tasks with a variable number of outputs.

• Use Pydantic model, dataclass or TypedDict for reusable, validated schemas:

  • Plain BaseModel, dataclass or TypedDict –> expanded namespace

  • model_config={'extra':'allow', 'item_type': T} –> dynamic namespace

  • model_config={'leaf': True} or Leaf[Model] –> single leaf (blob)

• Reuse .inputs and .outputs specifications at the graph level to build modular and robust workflows.

• Use select inside Annotated to reshape reused specifications (e.g., include/exclude with dotted paths, rename, prefix).

• Use meta to modify metadata of a specification, such as making it optional.

These tools are fundamental to building reproducible and verifiable scientific workflows with complete data provenance.