Combine workgraphs

Introduction

There are two ways to combine workgraphs:

Use the graph builder inside another workgraph.
Add the workgraph directly to another workgraph.

This tutorial will show an example for the second case. You need to install workgraph_collections to run this tutorial. You can install it by running pip install workgraph_collections.

[1]:

%load_ext aiida
from aiida import load_profile
load_profile()

[1]:

Profile<uuid='57ccbf7d9e2b41b39edb2bfdaf725feb' name='default'>

Relax, Bands and PDOS

Here is an example to combine bands workgraph, pdos workgraph with a relax task to form a new workgraph.

Bands workgraph

[2]:

from workgraph_collections.qe.bands import bands_workgraph
from aiida import load_profile
load_profile()

bands_wg = bands_workgraph(run_relax=False)
bands_wg.to_html()

[2]:

PDOS workgraph

[3]:

from workgraph_collections.qe.pdos import pdos_workgraph
pdos_wg = pdos_workgraph(run_scf=True)
pdos_wg.to_html()

[3]:

Extend workgraph

First, create a workgraph with relax task:

[4]:

from aiida_workgraph import WorkGraph
from aiida_quantumespresso.workflows.pw.relax import PwRelaxWorkChain

wg = WorkGraph('Relax_Bands_PDOS')
relax_task = wg.add_task(PwRelaxWorkChain, name='relax')
wg

[4]:

Now let’s combine bands workgraph, pdos workgraph with a relax task to form a new workgraph.

To add a workgraph into another workgraph, there are two steps:

extend the workgraph
adjust the links, add new links between tasks, and remove old links if needed.

[5]:

# extend a wroktree
wg.extend(bands_wg, prefix='bands_')
wg.extend(pdos_wg, prefix='pdos_')
# adjust the links
wg.add_link(relax_task.outputs['output_structure'], wg.tasks['bands_scf'].inputs['pw.structure'])
wg.add_link(relax_task.outputs['output_structure'], wg.tasks['pdos_scf'].inputs['pw.structure'])
wg.to_html()

[5]:

Graph builder

One can compare it with the other method: using the workgraph as a task inside another workgraph.

[6]:

wg = WorkGraph('Relax_Bands_PDOS_graph_builder')
relax_task = wg.add_task(PwRelaxWorkChain, name='relax')
bands_job = wg.add_task(bands_workgraph, name='bands_group')
pdos_job = wg.add_task(pdos_workgraph, name='pdos_group')
wg.add_link(relax_task.outputs['output_structure'], bands_job.inputs['structure'])
wg.add_link(relax_task.outputs['output_structure'], pdos_job.inputs['structure'])
wg.to_html()

[6]:

Pros

Show the top level workflow.

Cons

Only know the exact task graph at runtime.