Use protocol to set input parameters

Introduction

One can manually prepare all the inputs for a WorkGraph. One can also set the input parameters using the protocol, as long as the WorkChain support get input parameters based on the input protocol. For example, in the aiida-quantumespresso package, the PwRelaxWorkChain has a get_builder_from_protocol method. In this tutorial, we will show how to use the protocol to set the input parameters inside the WorkGraph.

Note: For the moment, it only suppor the ProtocolMixin from aiida-quantumespresso.

Load the AiiDA profile.

[1]:
%load_ext aiida
from aiida import load_profile
load_profile()
[1]:
Profile<uuid='4a9501fb20364cecb99920cadf27b238' name='xing'>

Use protocol to set input parameters

[18]:
from aiida_workgraph import WorkGraph
from aiida_quantumespresso.workflows.pw.relax import PwRelaxWorkChain
from aiida import orm
from ase.build import bulk
from pprint import pprint

code = orm.load_code("pw-7.2@localhost")
wg = WorkGraph("test_pw_relax")
structure_si = orm.StructureData(ase=bulk("Si"))
pw_relax1 = wg.add_task(PwRelaxWorkChain, name="pw_relax1")
# set the inputs from the protocol
# this will call the `PwRelaxWorkChain.get_builder_from_protocol` method
# to set the inputs of the workchain
pw_relax1.set_from_protocol(
    code,
    structure_si,
    protocol="fast",
    pseudo_family="SSSP/1.2/PBEsol/efficiency"
)
# we can now inspect the inputs of the workchain
print("The inputs for the PwBaseWorkchain are:")
print("-"*80)
pprint(pw_relax1.inputs["base"].value)
print("\nThe input parameters for pw are:")
print("-"*80)
pprint(pw_relax1.inputs["base"].value['pw']['parameters'].get_dict())


The inputs for the PwBaseWorkchain are:
--------------------------------------------------------------------------------
{'kpoints_distance': <Float: uuid: ae0202b9-594c-42fc-8975-c8ed380b946f (unstored) value: 0.5>,
 'kpoints_force_parity': <Bool: uuid: eca308be-0b51-4490-b01f-8c261316977f (unstored) value: False>,
 'metadata': {},
 'pw': {'code': <InstalledCode: Remote code 'pw-7.2' on localhost pk: 158, uuid: cd7f5a2a-3d4e-4335-b6a4-d114dba06ab9>,
        'metadata': {'options': {'max_wallclock_seconds': 43200,
                                 'resources': {'num_machines': 1},
                                 'stash': {},
                                 'withmpi': True}},
        'monitors': {},
        'parameters': <Dict: uuid: bc13abf6-fc1f-4d69-a809-5975e712dc8a (unstored)>,
        'pseudos': {'Si': <UpfData: uuid: 86dda7cc-d764-4291-a827-5cd9d2b7a07d (pk: 59)>}}}

The input parameters for pw are:
--------------------------------------------------------------------------------
{'CELL': {'cell_dofree': 'all', 'press_conv_thr': 0.5},
 'CONTROL': {'calculation': 'vc-relax',
             'etot_conv_thr': 0.0002,
             'forc_conv_thr': 0.001,
             'tprnfor': True,
             'tstress': True},
 'ELECTRONS': {'conv_thr': 8e-10, 'electron_maxstep': 80, 'mixing_beta': 0.4},
 'SYSTEM': {'degauss': 0.01,
            'ecutrho': 240.0,
            'ecutwfc': 30.0,
            'nosym': False,
            'occupations': 'smearing',
            'smearing': 'cold'}}

Then, we submit the workgraph

[7]:
wg.submit(wait=True, timeout=200)

Warning: RabbitMQ v3.9.13 is not supported and will cause unexpected problems!
Warning: It can cause long-running workflows to crash and jobs to be submitted multiple times.
Warning: See https://github.com/aiidateam/aiida-core/wiki/RabbitMQ-version-to-use for details.
Manager for incoming INPUT_WORK links for node pk=15022

Adjust the input parameters from the protocol

Sometimes, we want to adjust the input parameters from the protocol. For example, we want to remove the base_final_scf from the inputs, so that the PwRelaxWorkChain will not run the base_final_scf step.

[ ]:
# skip the `base_final_scf` step
pw_relax1.inputs["base_final_scf"].value = None