Hi @Xing ,
Here I am again with questions about using aiida-workgraph. I am trying to build the following workflow,
- Run a DFT calculation that optimizes the cell (keeping atoms fixed)
- Run a DFT calculation that optimizes atomic positions (keeping cell fixed), use as input structure the output structure of (1)
- Check output forces and stress of (2), if not below threshold start over at (1), passing the output structure of (2) as its input structure
I hope I am explaining this clearly enough. I have been trying to follow the example here (Use while loop - AiiDA WorkGraph documentation) for using the while task and the example here (Use Context to pass data between tasks - AiiDA WorkGraph documentation) for passing data between tasks. So far I have come up with the following,
wg = WorkGraph("cell_atoms_relax")
# Initialize context
wg.context = {
'struct': node1.inputs.structure,
'folder': node1.outputs.remote_folder,
'forces_stress': node1.outputs.forces_and_stress,
}
# Create while task
struct_converged = wg.add_task(structure_converged_siesta, name="converged", forces_and_stress='{{forces_stress}}', forces_threshold=orm.Float(0.04), stress_threshold=orm.Float(1 / cs.RyAng3Tokbar))
while1 = wg.add_task("While", max_iterations=100, conditions=struct_converged.outputs["result"])
# Create tasks within while loop
struct_ctx1 = wg.add_task("workgraph.from_context", name="struct_ctx1", key="struct")
folder_ctx1 = wg.add_task("workgraph.from_context", name="folder_ctx1", key="folder")
relax_cell = wg.add_task(SiestaBaseWorkChain, name="relax_cell")
relax_cell.set({
'pseudos': node1.inputs.pseudos,
'basis': node1.inputs.basis,
'clean_workdir': node1.inputs.clean_workdir,
'code': node1.inputs.code,
'kpoints': node1.inputs.kpoints,
'max_iterations': orm.Int(5),
'options': node2.inputs.options,
'parameters': orm.Dict(dict=pdict_cell.get_dict()),
'structure': struct_ctx1.outputs["result"],
'parent_calc_folder': folder_ctx1.outputs["result"],
})
relax_cell.set_context({
'output_structure': 'struct',
'remote_folder': 'folder',
})
struct_ctx2 = wg.add_task("workgraph.from_context", name="struct_ctx2", key="struct")
folder_ctx2 = wg.add_task("workgraph.from_context", name="folder_ctx2", key="folder")
struct_ctx2.waiting_on.add("relax_cell")
folder_ctx2.waiting_on.add("relax_cell")
relax_atoms = wg.add_task(SiestaBaseWorkChain, name="relax_atoms")
relax_atoms.set({
'lua.md_run': node3.inputs.lua.md_run,
'lua.script': node3.inputs.lua.script,
'pseudos': node1.inputs.pseudos,
'basis': node1.inputs.basis,
'clean_workdir': node1.inputs.clean_workdir,
'code': node1.inputs.code,
'kpoints': node1.inputs.kpoints,
'max_iterations': orm.Int(5),
'options': node2.inputs.options,
'parameters': orm.Dict(dict=pdict_atoms.get_dict()),
'structure': struct_ctx2.outputs["result"],
'parent_calc_folder': folder_ctx2.outputs["result"],
})
relax_atoms.set_context({
'output_structure': 'struct',
'remote_folder': 'folder',
'forces_and_stress': 'forces_stress',
})
while1.children.add([
"struct_ctx1",
"folder_ctx1",
"relax_cell",
"struct_ctx2",
"folder_ctx2",
"relax_atoms",
])
Here node1, node2, node3 are SiestaBaseWorkChain nodes from earlier calculations, and pdict_cell and pdict_relax are dictionaries with input parameters I prepared beforehand. Also the comparison function is defined like this,
from aiida_workgraph import task
import numpy as np
@task.calcfunction()
def structure_converged_siesta(forces_and_stress, forces_threshold, stress_threshold):
# Check if forces are below threshold
forces_conv = np.all(np.abs(forces_and_stress.get_array('forces')) < forces_threshold.value)
# Check if stress is below threshold
stress_conv = np.all(np.abs(forces_and_stress.get_array('stress')) < stress_threshold.value)
return forces_conv and stress_conv
However, if I try to submit this workgraph I get a very obscure error message:
aiida_workgraph/utils/analysis.py:91, in WorkGraphSaver.build_task_link(self)
89 output["links"] = []
90 for link in self.wgdata["links"]:
---> 91 to_socket = [
92 socket
93 for name, socket in self.wgdata["tasks"][link["to_node"]][
94 "inputs"
95 ].items()
96 if name == link["to_socket"]
97 ][0]
98 from_socket = [
99 socket
100 for name, socket in self.wgdata["tasks"][link["from_node"]][
(...)
103 if name == link["from_socket"]
104 ][0]
105 to_socket["links"].append(link)
IndexError: list index out of range
This sounds like an error somewhere in the internals of aiida-workgraph and doesn’t really help me to figure out what is going wrong. Do you have any ideas on how to debug this? Or can you already spot an error in how I am setting up the workgraph? Or do you have other suggestions for creating a workflow like this with aiida-workgraph that is easier and/or less error prone? Any help/insight is appreciated.
Let’s hope the rest goes smoothly.
. Maybe to mimic the builder, e.g.,