Source code for aiida_quantumespresso.workflows.functions.get_marked_structures

# -*- coding: utf-8 -*-
"""CalcFunction to create structures with a marked atom for each site in a list."""
from aiida import orm
from aiida.common import ValidationError
from aiida.engine import calcfunction
from aiida.orm.nodes.data.structure import Kind, Site, StructureData


@calcfunction
def get_marked_structures(structure, atoms_list, marker='X'):
    """Read a StructureData object and return structures for XPS calculations.

    :param atoms_list:  the atoms_list of atoms to be marked.
    :param marker: a Str node defining the name of the marked atom Kind. Default is 'X'.
    :returns: StructureData objects for the generated structure.
    """
    marker = marker.value
    elements_present = [kind.symbol for kind in structure.kinds]
    if marker in elements_present:
        raise ValidationError(
            f'The marker ("{marker}") should not match an existing Kind in '
            f'the input structure ({elements_present}.'
        )

    output_params = {}
    result = {}

    for index in atoms_list.get_list():
        marked_structure = StructureData()
        kinds = {kind.name: kind for kind in structure.kinds}
        marked_structure.set_cell(structure.cell)

        for i, site in enumerate(structure.sites):
            if i == index:
                marked_kind = Kind(name=marker, symbols=site.kind_name)
                marked_site = Site(kind_name=marked_kind.name, position=site.position)
                marked_structure.append_kind(marked_kind)
                marked_structure.append_site(marked_site)
                output_params[f'site_{index}'] = {'symbol': site.kind_name, 'multiplicity': 1}
            else:
                if site.kind_name not in [kind.name for kind in marked_structure.kinds]:
                    marked_structure.append_kind(kinds[site.kind_name])
                new_site = Site(kind_name=site.kind_name, position=site.position)
                marked_structure.append_site(new_site)
        result[f'site_{index}'] = marked_structure

    result['output_parameters'] = orm.Dict(dict=output_params)

    return result