I’m migrating a project from aiida-workgraph==0.5.2 and need some help finding the modern equivalent for the active_map_zone functionality.
Previously, my workflow would dynamically generate an unknown number of slab structures and then use active_map_zone to run a parallel relaxation for each one. The pattern looked like this:
# 1. A task generates a dictionary of slab structures (number unknown beforehand)
generate_slabs_task = wg.add_task(
get_slabs,
name="generate_slabs",
relaxed_structure=output_struct_node,
# ... other parameters
)
slab_structures = generate_slabs_task.outputs.structures
# 2. active_map_zone takes this output and creates parallel tasks
with active_map_zone(slab_structures) as map_zone:
slab_task = map_zone.add_task(
dft_workchain,
name="slab_relaxation",
structure=map_zone.item, # Placeholder for each slab
)
slab_task.set_from_builder(builder_slab)
This worked perfectly, creating a dynamic number of parallel calculations based on the output of generate_slabs_task.
However, the active_map_zone context manager has been removed in the latest version of aiida-workgraph. My attempts to use the new Map or Zone context managers to replicate this behavior have failed — I tried several different approaches without success.
What is the correct way to achieve this “generate-then-map” pattern in the current version of aiida-workgraph? Specifically, how can I:
Generate a variable number of inputs dynamically within a workflow, and
Map a sub-workflow/task over those inputs in parallel?
Any pointers, code snippets, or example workflows would be greatly appreciated!
The online docs correspond to the latest code of the aiida-workgraph repository, and the Map zone was updated recently, which could be the reason why it does not work. I suggest you install the latest pre-stable release.
Thanks for the suggestions and quick answer! I tried today to make the Zone approach working, but unfortunately it fails when trying to map over a task’s output socket. Maybe its due to my lack of experience, but I will try to explain what I am getting.
I will put here a minimal example to illustrate the problem.
Minimal Example
from aiida_workgraph import WorkGraph, task, Map
from aiida.orm import Dict, StructureData
@task.calcfunction(outputs=["structures"])
def generate_slabs():
"""Returns a dictionary of structures."""
return {'structures': {'s_0': StructureData(), 's_1': StructureData()}}
@task
def process_structure(structure):
return structure
# This works fine:
with WorkGraph('static_map') as wg:
static_dict = {'s_0': StructureData(), 's_1': StructureData()}
with Map(static_dict) as map_zone:
result = process_structure(structure=map_zone.item).result
# This fails with the error above:
with WorkGraph('dynamic_map') as wg:
gen_task = wg.add_task(generate_slabs, name="generate")
dynamic_dict = gen_task.outputs.structures
with Map(dynamic_dict) as map_zone: # ❌ Fails here
result = process_structure(structure=map_zone.item).result
The error occurs when trying to pass map_zone.item to a task parameter inside the Map context, specifically when the Map is iterating over a task output socket rather than a concrete dictionary.
The previous version ‘active_map_zone’ worked for what I needed, so I think that I’ll stick with version 0.5.2 of aiida-workgraph for the moment.
If you think that this is indeed a problem and not just me that’s having a problem with this, I’ll open an issue in Github.
In the latest version, map_zone.item has key and value outputs. For your example code, you only need to change map_zone.item to map_zone.item.value in both cases (static and dynamic).
result = process_structure(structure=map_zone.item.value).result
important: you need to install the pre-stable release 1.0.0b3,