from aiida_workgraph import task, WorkGraph
from ase import Atoms
from ase.build import bulk

@task.pythonjob(outputs=[{"name": "structure_lst", "identifier": "Namespace"}])
def generate_structures(structure: Atoms, factor_lst: list) -> list[Atoms]:
    structure_lst = []
    for factor in factor_lst:
        structure_strain = structure.copy()
        structure_strain.set_cell(
            structure_strain.cell * factor,
            scale_atoms=True
        )
        structure_lst.append(structure_strain)
    return structure_lst

# user can also decorate a existing function
#generate_structures = task.pythonjob(outputs=[{"name": "structure_lst", "identifier": "Namespace"}])(generate_structures)


atoms = bulk("Si")
wg = WorkGraph("EOS-WG")
generate_structures_task = wg.add_task(
        generate_structures,
        name="generate_structures",
        factor_lst=[0.95, 1.0, 1.05],
        structure=atoms,
    )
wg.submit()