Source code for aiida_raspa.utils.base_input_generator

# -*- coding: utf-8 -*-
"""Basic raspa input generator."""
from copy import deepcopy

ORDERED_ITEMS_COMPONENT_SECTION = [
    "NumberOfIdentityChanges",
    "IdentityChangesList",
]


class RaspaInput:  # pylint: disable=too-few-public-methods
    """Convert input dictionary into input file"""

    def __init__(self, params):
        self.params = deepcopy(params)  # make sure the original object is not modified
        try:
            self.system_order = sorted(params["System"].keys())  # we sort the keys to keep the order in the input
        except KeyError:
            raise KeyError("The input dictionary should contain the System subdictionary.")
        except AttributeError:
            raise AttributeError("The System subdictionary should be of type dict.")
        # file persistent this is essential for the restarts.

        if not self.system_order:
            raise ValueError("The System subdictionary should not be empty.")

    # --------------------------------------------------------------------------
    def render(self):
        """Perform conversion"""
        output = ["!!! Generated by AiiDA !!!"]
        params = deepcopy(self.params)  # make sure it is idempotent

        # General settings can be be rendered without any problems
        section = params.pop("GeneralSettings")
        self._render_section(output, section)

        # System section has to make sure to keep the order persistent, this is why we use
        # self.system_order to iterate through
        section = params.pop("System")
        for my_id, name in enumerate(self.system_order):
            system = section[name]
            s_type = system.pop("type")
            output.append("{} {}".format(s_type, my_id))
            if s_type == "Framework":
                output.append("{}FrameworkName {}".format(' ' * 3, name))
            self._render_section(output, system, indent=3)

        # Section Component may contain parameters thar are for several systems
        # like:
        # CreateNumberOfMolecules n1, n2  <--- n1 is for system 0, n2 is for system 1
        section = params.pop("Component")
        for my_id, (name, molecule) in enumerate(section.items()):
            output.append('Component {} MoleculeName {}'.format(my_id, name))
            if "BlockPocketsFileName" in molecule:
                if isinstance(molecule["BlockPocketsFileName"], dict):
                    try:
                        bps = self._dict_to_ordered_list(molecule["BlockPocketsFileName"])
                    except KeyError as err:
                        raise KeyError("You did not provide BlockPocketsFileName parameter for the system '{}'".format(
                            str(err)))

                    molecule["BlockPocketsFileName"] = bps
                    molecule["BlockPockets"] = ["yes" if bp else "no" for bp in bps]
                elif isinstance(molecule["BlockPocketsFileName"], str):
                    molecule["BlockPockets"] = "yes"

            if "CreateNumberOfMolecules" in molecule:
                if isinstance(molecule["CreateNumberOfMolecules"], dict):
                    try:
                        molecule["CreateNumberOfMolecules"] = self._dict_to_ordered_list(
                            molecule["CreateNumberOfMolecules"])
                    except KeyError as err:
                        raise KeyError("You did not provide CreateNumberOfMolecules parameter for the system {}".format(
                            str(err)))

            for item in ORDERED_ITEMS_COMPONENT_SECTION:
                if item in molecule:
                    self._render_item(output, item, molecule.pop(item), indent=3)

            self._render_section(output, molecule, indent=3)

        return "\n".join(output) + "\n"

    def _dict_to_ordered_list(self, input_dict):
        """Convert dict to ordered list.

        Convert dictionary {"framework_name1": something1, "framework_name2": something2}
        to [something2, something1], where the order is defined by self.system_order list

        :param input_dict: dictionary that looks as folows {"framework_name1": something1,
            "framework_name2": something2}
        :type input_dict: dict


        :returns: a list of items ordered according to the order of frameworks in the input file
        :rtype: list
        """
        ordered_list = [None] * len(self.system_order)
        for i, name in enumerate(self.system_order):
            ordered_list[i] = input_dict[name]
        return ordered_list

    # --------------------------------------------------------------------------
    @staticmethod
    def _render_section(output, params, indent=0):
        """
        It takes a dictionary and recurses through.

        For key-value pair it checks whether the value is a dictionary
        and prepends the key with &
        It passes the valued to the same function, increasing the indentation
        If the value is a list, I assume that this is something the user
        wants to store repetitively
        """
        #        output.append("enter")
        #        output.append("This what comes:" + str(params))
        for key, val in sorted(params.items()):
            RaspaInput._render_item(output, key, val, indent=indent)

    @staticmethod
    def _render_item(output, key, val, indent=0):
        """
        It takes one key-value item and adds to the output file
        """
        if isinstance(val, list):
            output.append('{}{} {}'.format(' ' * indent, key, ' '.join(str(p) for p in val)))
        elif isinstance(val, bool):
            val_str = 'yes' if val else 'no'
            output.append('{}{} {}'.format(' ' * indent, key, val_str))
        else:
            output.append('{}{} {}'.format(' ' * indent, key, val))