Source code for aim2dat.aiida_workflows.cp2k.planar_fields_work_chain

"""
Aiida work chains for cp2k to calculate 2d fields based on the electron density.
"""

# Third party library imports
import aiida.orm as aiida_orm
from aiida.plugins import CalculationFactory
from aiida.engine import (
    while_,
    process_handler,
    ExitCode,
)
from aiida.common import AttributeDict

# Internal library imports
from aim2dat.aiida_workflows.cp2k.base_core_work_chain import _BaseCoreWorkChain
from aim2dat.aiida_workflows.cp2k.core_work_chain_handlers import _switch_to_atomic_scf_guess
from aim2dat.utils.dict_tools import dict_create_tree

Critic2Calculation = CalculationFactory("aim2dat.critic2")


def _validate_plane_vectors(value, _):
    value = value.get_list()
    if len(value) != 11:
        return "`plane_vectors` must have 11 entries."
    if not all(isinstance(val0, (float, int)) for val0 in value[:9]):
        return "First nine entries of `plane_vectors` need to be of type int or float."
    if not all(isinstance(val0, int) for val0 in value[9:]):
        return "Last two entries of `plane_vectors` need to be of type int."


def _validate_field_types(value, _):
    value = value.get_list()
    if not all(val0 in ["elf", "deformation_density", "total_density"] for val0 in value):
        return "Only 'elf', 'deformation_density' and 'total_density' are supported field types."
    if len(set(value)) != len(value):
        return "Field types cannot be calculated twice."


def _create_critic2_input_parameters(f_type, plane, structure):
    system = "molecule"
    if any(structure.pbc):
        system = "crystal"
    if f_type == "deformation_density":
        parameters = [
            system + " aiida-ELECTRON_DENSITY-1_0.cube",
            "load aiida-ELECTRON_DENSITY-1_0.cube core zpsp",
            'load as "$1-$0"',
            "plane " + " ".join(str(val) for val in plane) + " field 2 file rhodef",
        ]
    elif f_type == "elf":
        parameters = [
            system + " aiida-ELF_S1-1_0.cube",
            "load aiida-ELF_S1-1_0.cube",
            "plane " + " ".join(str(val) for val in plane) + " field 1 file elf",
        ]
    elif f_type == "total_density":
        parameters = [
            system + " aiida-TOTAL_DENSITY-1_0.cube",
            "load aiida-TOTAL_DENSITY-1_0.cube",
            "plane " + " ".join(str(val) for val in plane) + " field 1 file total_density",
        ]
    return aiida_orm.List(list=parameters)


[docs]class PlanarFieldsWorkChain(_BaseCoreWorkChain):
    """AiiDA work chain to calculate the planar fields."""

    _keep_scf_method_fixed = True
    _keep_smearing_fixed = True
    _initial_scf_guess = "RESTART"

    @classmethod
    def define(cls, spec):
        """Specify inputs and outputs."""
        super().define(spec)
        spec.input(
            "adjust_scf_parameters",
            valid_type=aiida_orm.Bool,
            default=lambda: aiida_orm.Bool(False),
            help="Restart calculation with adjusted parameters if SCF-clycles are not converged.",
        )
        spec.input(
            "always_add_unocc_states",
            valid_type=aiida_orm.Bool,
            default=lambda: aiida_orm.Bool(False),
            help="Always include some unoccupied states even if smearing is not used.",
        )
        spec.input(
            "store_cubes",
            valid_type=aiida_orm.Bool,
            default=lambda: aiida_orm.Bool(False),
            help="Whether to store the cube files that are input for the planes.",
        )
        spec.input(
            "plane_vectors",
            valid_type=aiida_orm.List,
            validator=_validate_plane_vectors,
            help="Plane vectors and number of points given in crystallographic coordinates and "
            "angstrom for (partial) periodic boundary conditions and non-periodic boundary "
            "conditions, respectively.",
        )
        spec.input(
            "field_types",
            valid_type=aiida_orm.List,
            validator=_validate_field_types,
            default=lambda: aiida_orm.List(list=["deformation_density", "elf", "total_density"]),
            help="The field type that is calculated, up to now only 'deformation_density' and "
            "'elf' are supported.",
        )
        spec.expose_inputs(
            Critic2Calculation,
            namespace="critic2",
            exclude=("charge_density_folder", "kind_info", "parameters"),
            namespace_options={
                "required": False,
                "populate_defaults": False,
                "help": "Input parameters of critic2.",
            },
        )
        spec.expose_outputs(
            Critic2Calculation,
            "critic2.total_density",
            namespace_options={
                "required": False,
                "help": "Critic2 outputs of the total density plane calculation.",
            },
        )
        spec.expose_outputs(
            Critic2Calculation,
            "critic2.deformation_density",
            namespace_options={
                "required": False,
                "help": "Critic2 outputs of the deformation density plane calculation.",
            },
        )
        spec.expose_outputs(
            Critic2Calculation,
            "critic2.elf",
            namespace_options={
                "required": False,
                "help": "Critic2 outputs of the elf plane calculation.",
            },
        )
        spec.outline(
            cls.setup_inputs,
            cls.setup_wc_specific_inputs,
            cls.initialize_scf_parameters,
            while_(cls.should_run_process)(
                cls.run_process,
                cls.inspect_process,
            ),
            cls.post_processing,
            cls.setup_critic2_calculation,
            cls.wc_specific_post_processing,
        )

    def setup_wc_specific_inputs(self):
        """Set input parameters to calculate partial charges."""
        self.ctx.inputs.metadata.options.parser_name = "aim2dat.cp2k.partial_charges"
        self.ctx.field_types = self.inputs.field_types.get_list()
        parameters = self.ctx.inputs.parameters.get_dict()
        extra_sections = {}
        if "deformation_density" in self.ctx.field_types:
            extra_sections["E_DENSITY_CUBE"] = {"STRIDE": 1}
        if "elf" in self.ctx.field_types:
            extra_sections["ELF_CUBE"] = {"STRIDE": 1}
        if "total_density" in self.ctx.field_types:
            extra_sections["TOT_DENSITY_CUBE"] = {"STRIDE": 1}
        dict_create_tree(parameters, ["FORCE_EVAL", "DFT", "PRINT"])
        parameters["FORCE_EVAL"]["DFT"]["PRINT"].update(extra_sections)
        self.ctx.inputs.parameters = aiida_orm.Dict(dict=parameters)
        calcjob_settings = {"output_check_scf_conv": True}
        if "store_cubes" in self.inputs:
            calcjob_settings["additional_retrieve_temporary_list"] = ["*.cube"]
        self.ctx.inputs.settings = aiida_orm.Dict(dict=calcjob_settings)

    def setup_critic2_calculation(self):
        """Set input parameters for external post-processing codes."""
        for f_type in self.ctx.field_types:
            inputs = AttributeDict(self.exposed_inputs(Critic2Calculation, "critic2"))
            inputs.charge_density_folder = self.ctx.children[-1].outputs.remote_folder
            inputs.kind_info = self.ctx.children[-1].outputs.output_kind_info
            inputs.parameters = _create_critic2_input_parameters(
                f_type, self.inputs.plane_vectors.get_list(), self.ctx.inputs.structure
            )
            running = self.submit(Critic2Calculation, **inputs)
            self.report(f"Launching {f_type} field:  <{running.pk}>.")
            self.to_context(**{f_type: running})

    def wc_specific_post_processing(self):
        """Expose outputs of the external codes."""
        for f_type in self.ctx.field_types:
            critic2_calc = getattr(self.ctx, f_type)
            if not critic2_calc.is_finished_ok:
                return self.exit_codes.ERROR_CALCULATION_ABORTED
            self.out_many(
                self.exposed_outputs(
                    getattr(self.ctx, f_type), Critic2Calculation, namespace="critic2." + f_type
                )
            )

    @process_handler(
        priority=402,
        exit_codes=ExitCode(0),
    )
    def switch_to_atomic_scf_guess(self, calc):
        """
        Switch to atomic guess for the case that the scf-cycles do not converge.
        """
        return self._execute_error_handler(calc, _switch_to_atomic_scf_guess)