aiidateam / aiida-core

The official repository for the AiiDA code
https://aiida-core.readthedocs.io
Other
431 stars 187 forks source link

Allow `calcfunction` to return a `list`? #6354

Open mbercx opened 5 months ago

mbercx commented 5 months ago

Is your feature request related to a problem? Please describe

Currently, a calcfunction can return multiple outputs by returning a dict where the keys will be used as the labels of the CREATE links. Take this very basic example (which I've actually used in a real-life work chain):

from aiida import orm, engine, load_profile

load_profile()

@engine.calcfunction
def split_list(list_node: orm.List):
    return {
        f'element_{i}': orm.Float(item)
        for i, item in enumerate(list_node)
    }

However, if you try to return a list:

@engine.calcfunction
def split_list(list_node: orm.List):
    return [orm.Float(item) for item in list_node]

AiiDA will complain:

TypeError: Function process returned an output with unsupported type '<class 'list'>'
Must be a Data type or a mapping of {string: Data}

Describe the solution you'd like

I'm wondering if we can't simply use a set of sensible defaults for the CREATE link (e.g. output_i) in this case, to make the life of the user easier?

Additional context

This came up when trying to implement the cleanest version possible of a calcfunction that can rescale structures for e.g. an EOS:

@engine.calcfunction
def rescale_list(structure: orm.StructureData, factor_list: orm.List):

    scaled_structure_dict = {}

    for index, scaling_factor in enumerate(factor_list.get_list()):

        ase_structure = structure.get_ase()

        new_cell = ase_structure.get_cell() * scaling_factor
        ase_structure.set_cell(new_cell, scale_atoms=True)

        scaled_structure_dict[f'structure_{index}'] = orm.StructureData(ase=ase_structure)

    return scaled_structure_dict
mbercx commented 5 months ago

One note here is that the calcfunction in this case should also return a list of nodes, else the behavior would be confusing.