How to return an arbitrary number of items from a calcfunction/workfunction

General Usage
1

Hello,

Let’s say I want to use AiiDA to wrap some python functions I have that generate many molecule configurations so that I can compare them and pipe them into other workflows. As my functions are currently implemented, the number of molecules returned is variable.

def build_configuration(atoms) -> Atoms:
    atoms.append(random_atom) 
    return atoms

@calcfunction
def generate_configurations(atoms, num_configs):
    configurations = List()
    for config in num_configs:
        config_atoms = build_configuration(atoms)
        config_struct = StructureData(ase=config_atoms)
        configurations.append(config_struct)
    return configurations

I want all of these configurations stored in provenance, but I don’t care how that’s done. One method I’ve tried is to turn all of these into StructureData objects, and pack them all into an aiida List that the calcfunction returns (shown above in pseudocode). I suspect part of why this doesn’t work is because I’d need to return the StructureData nodes individually from calc functions and then pack them into a List.

Do I have to pack all of this into a WorkChain? It isn’t clear to me how to do that given that I still need to define (I think) a set number of outputs in the spec.

I’ve looked around the documentation to find an example where multiple things are being returned, but couldn’t find anything.

Thanks AiiDA team!

2

Hi @cote3804

you can use a dict to return multiple results from a calcfunction or workfunction. Taking your example above, this could look like the following:

@calcfunction
def generate_configurations(atoms, num_configs):
    configurations = []
    for config in num_configs:
        config_atoms = build_configuration(atoms)
        config_struct = StructureData(ase=config_atoms)
        configurations.append(config_struct)
    return {f'configuration{i}': config for i, config in enumerate(configurations)}

If you don’t care about a specific label, you can of course simply iterate the dict.values() afterwards.

2 Likes
3

Please check this related issue by @mbercx .

1 Like