Source code for aim2dat.aiida_workflows.cp2k.partial_charges_work_chain

"""
Aiida work chains for cp2k to find parameters that converge the Kohn-Sham equations.
"""

# 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_set_parameter, dict_create_tree

Critic2Calculation = CalculationFactory("aim2dat.critic2")
ChargemolCalculation = CalculationFactory("aim2dat.chargemol")


[docs]class PartialChargesWorkChain(_BaseCoreWorkChain):
    """AiiDA work chain to calculate the partial charges of the atoms."""

    _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 electron density cube file (only supported for chargemol "
            "and critic2).",
        )
        spec.expose_inputs(
            Critic2Calculation,
            namespace="critic2",
            exclude=("charge_density_folder", "kind_info"),
            namespace_options={"required": False, "populate_defaults": False, "help": "..."},
        )
        spec.expose_inputs(
            ChargemolCalculation,
            namespace="chargemol",
            exclude=("charge_density_folder", "charge_density_filename", "kind_info"),
            namespace_options={"required": False, "populate_defaults": False, "help": "..."},
        )
        spec.expose_outputs(
            Critic2Calculation,
            namespace="critic2",
            namespace_options={"required": False, "help": "Output parameters of critic2."},
        )
        spec.expose_outputs(
            ChargemolCalculation,
            namespace="chargemol",
            namespace_options={"required": False, "help": "Output parameters of chargemol."},
        )
        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_external_partial_charge_analysis,
            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"
        parameters = self.ctx.inputs.parameters.get_dict()
        dict_create_tree(parameters, ["GLOBAL"])
        dict_set_parameter(parameters, ["GLOBAL", "PRINT_LEVEL"], "MEDIUM")
        if "critic2" in self.inputs or "chargemol" in self.inputs:
            extra_sections = {"E_DENSITY_CUBE": {"STRIDE": 1}}
            dict_create_tree(parameters, ["FORCE_EVAL", "DFT", "PRINT"])
            parameters["FORCE_EVAL"]["DFT"]["PRINT"].update(extra_sections)
        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)
        self.ctx.inputs.parameters = aiida_orm.Dict(dict=parameters)

    def setup_external_partial_charge_analysis(self):
        """Set input parameters for external post-processing codes."""
        if "critic2" in self.inputs:
            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
            running = self.submit(Critic2Calculation, **inputs)
            self.report(f"Launching  <{running.pk}>.")
            self.to_context(critic2=running)
        if "chargemol" in self.inputs:
            inputs = AttributeDict(self.exposed_inputs(ChargemolCalculation, "chargemol"))
            inputs.charge_density_folder = self.ctx.children[-1].outputs.remote_folder
            inputs.charge_density_filename = aiida_orm.Str("aiida-ELECTRON_DENSITY-1_0.cube")
            inputs.kind_info = self.ctx.children[-1].outputs.output_kind_info
            running = self.submit(ChargemolCalculation, **inputs)
            self.report(f"Launching  <{running.pk}>.")
            self.to_context(chargemol=running)

    def wc_specific_post_processing(self):
        """Expose outputs of the external codes."""
        if "critic2" in self.inputs:
            self.out_many(
                self.exposed_outputs(self.ctx.critic2, Critic2Calculation, namespace="critic2")
            )
        if "chargemol" in self.inputs:
            self.out_many(
                self.exposed_outputs(
                    self.ctx.chargemol, ChargemolCalculation, namespace="chargemol"
                )
            )

    @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)