aiida_fleur.data.fleurinpmodifier — AiiDA-FLEUR 1.1.0 documentation

# Source code for aiida_fleur.data.fleurinpmodifier

    # -*- coding: utf-8 -*-
    # pylint: disable=inconsistent-return-statements,protected-access,missing-docstring
    ###############################################################################
    # Copyright (c), Forschungszentrum Jülich GmbH, IAS-1/PGI-1, Germany.         #
    #                All rights reserved.                                         #
    # This file is part of the AiiDA-FLEUR package.                               #
    #                                                                             #
    # The code is hosted on GitHub at https://github.com/JuDFTteam/aiida-fleur    #
    # For further information on the license, see the LICENSE.txt file            #
    # For further information please visit http://www.flapw.de or                 #
    # http://aiida-fleur.readthedocs.io/en/develop/                               #
    ###############################################################################
    """
    In this module is the FleurinpModifier class, which is used to manipulate
    FleurinpData objects in a way which keeps the provernance.
    """

    from __future__ import absolute_import
    from __future__ import print_function
    import os

    from lxml import etree

    from aiida.plugins import DataFactory
    from aiida.engine.processes.functions import calcfunction as cf
    from aiida_fleur.data.fleurinp import FleurinpData

    class FleurinpModifier(object):
        """
        A class which represents changes to the :class:`~aiida_fleur.data.fleurinp.FleurinpData` object.
        """

        def __init__(self, original):
            """
            Initiation of FleurinpModifier.
            """
            assert isinstance(original, FleurinpData), 'Wrong AiiDA data type'

            self._original = original
            self._tasks = []

        @staticmethod
        def apply_modifications(fleurinp_tree_copy, modification_tasks, schema_tree=None):
            """
            Applies given modifications to the fleurinp lxml tree.
            It also checks if a new lxml tree is validated against schema.
            Does not rise an error if inp.xml is not validated, simple prints a message about it.

            :param fleurinp_tree_copy: a fleurinp lxml tree to be modified
            :param modification_tasks: a list of modification tuples

            :returns: a modified fleurinp lxml tree
            """
            from aiida_fleur.tools.xml_util import xml_set_attribv_occ, xml_set_first_attribv
            from aiida_fleur.tools.xml_util import xml_set_all_attribv, xml_set_text
            from aiida_fleur.tools.xml_util import xml_set_text_occ, xml_set_all_text
            from aiida_fleur.tools.xml_util import create_tag, replace_tag, delete_tag
            from aiida_fleur.tools.xml_util import delete_att, set_species
            from aiida_fleur.tools.xml_util import change_atomgr_att, add_num_to_att
            from aiida_fleur.tools.xml_util import change_atomgr_att_label, set_species_label
            from aiida_fleur.tools.xml_util import set_inpchanges, set_nkpts, set_kpath, shift_value
            from aiida_fleur.tools.xml_util import shift_value_species_label
            from aiida_fleur.tools.xml_util import clear_xml

            def xml_set_attribv_occ1(fleurinp_tree_copy, xpathn, attributename, attribv, occ=None, create=False):
                if occ is None:
                    occ = [0]
                xml_set_attribv_occ(fleurinp_tree_copy, xpathn, attributename, attribv, occ=occ, create=create)
                return fleurinp_tree_copy

            def xml_set_first_attribv1(fleurinp_tree_copy, xpathn, attributename, attribv, create=False):
                xml_set_first_attribv(fleurinp_tree_copy, xpathn, attributename, attribv, create=create)
                return fleurinp_tree_copy

            def xml_set_all_attribv1(fleurinp_tree_copy, xpathn, attributename, attribv, create=False):
                xml_set_all_attribv(fleurinp_tree_copy, xpathn, attributename, attribv, create=create)
                return fleurinp_tree_copy

            def xml_set_text1(fleurinp_tree_copy, xpathn, text, create=False):
                xml_set_text(fleurinp_tree_copy, xpathn, text, create=create)
                return fleurinp_tree_copy

            def xml_set_text_occ1(fleurinp_tree_copy, xpathn, text, create=False, occ=0):
                xml_set_text_occ(fleurinp_tree_copy, xpathn, text, create=create, occ=occ)
                return fleurinp_tree_copy

            def xml_set_all_text1(fleurinp_tree_copy, xpathn, text, create=False):
                xml_set_all_text(fleurinp_tree_copy, xpathn, text, create=create)
                return fleurinp_tree_copy

            def create_tag1(fleurinp_tree_copy, xpath, newelement, create=False):
                fleurinp_tree_copy = create_tag(fleurinp_tree_copy, xpath, newelement, create=create)
                return fleurinp_tree_copy

            def delete_att1(fleurinp_tree_copy, xpath, attrib):
                fleurinp_tree_copy = delete_att(fleurinp_tree_copy, xpath, attrib)
                return fleurinp_tree_copy

            def delete_tag1(fleurinp_tree_copy, xpath):
                fleurinp_tree_copy = delete_tag(fleurinp_tree_copy, xpath)
                return fleurinp_tree_copy

            def replace_tag1(fleurinp_tree_copy, xpath, newelement):
                fleurinp_tree_copy = replace_tag(fleurinp_tree_copy, xpath, newelement)
                return fleurinp_tree_copy

            def set_species1(fleurinp_tree_copy, species_name, attributedict, create=False):
                fleurinp_tree_copy = set_species(fleurinp_tree_copy, species_name, attributedict, create=create)
                return fleurinp_tree_copy

            def set_species2(fleurinp_tree_copy, at_label, attributedict, create=False):
                fleurinp_tree_copy = set_species_label(fleurinp_tree_copy, at_label, attributedict, create=create)
                return fleurinp_tree_copy

            def change_atomgr_att1(fleurinp_tree_copy, attributedict, position=None, species=None, create=False):
                fleurinp_tree_copy = change_atomgr_att(fleurinp_tree_copy,
                                                       attributedict,
                                                       position=position,
                                                       species=species)
                return fleurinp_tree_copy

            def change_atomgr_att2(fleurinp_tree_copy, attributedict, atom_label, create=False):
                fleurinp_tree_copy = change_atomgr_att_label(fleurinp_tree_copy, attributedict, at_label=atom_label)
                return fleurinp_tree_copy

            def add_num_to_att1(fleurinp_tree_copy, xpathn, attributename, set_val, mode='abs', occ=None):
                if occ is None:
                    occ = [0]
                fleurinp_tree_copy = add_num_to_att(fleurinp_tree_copy, xpathn, attributename, set_val, mode=mode, occ=occ)
                return fleurinp_tree_copy

            def set_inpchanges1(fleurinp_tree_copy, change_dict):
                fleurinp_tree_copy = set_inpchanges(fleurinp_tree_copy, change_dict)
                return fleurinp_tree_copy

            def shift_value1(fleurinp_tree_copy, change_dict, mode):
                fleurinp_tree_copy = shift_value(fleurinp_tree_copy, change_dict, mode)
                return fleurinp_tree_copy

            def shift_value_species_label1(fleurinp_tree_copy, label, att_name, value, mode):
                fleurinp_tree_copy = shift_value_species_label(fleurinp_tree_copy, label, att_name, value, mode)
                return fleurinp_tree_copy

            def set_nkpts1(fleurinp_tree_copy, count, gamma):
                fleurinp_tree_copy = set_nkpts(fleurinp_tree_copy, count, gamma)
                return fleurinp_tree_copy

            def set_kpath1(fleurinp_tree_copy, kpath, count, gamma):
                fleurinp_tree_copy = set_kpath(fleurinp_tree_copy, kpath, count, gamma)
                return fleurinp_tree_copy

            actions = {
                'xml_set_attribv_occ': xml_set_attribv_occ1,
                'xml_set_first_attribv': xml_set_first_attribv1,
                'xml_set_all_attribv': xml_set_all_attribv1,
                'xml_set_text': xml_set_text1,
                'xml_set_text_occ': xml_set_text_occ1,
                'xml_set_all_text': xml_set_all_text1,
                'create_tag': create_tag1,
                'replace_tag': replace_tag1,
                'delete_tag': delete_tag1,
                'delete_att': delete_att1,
                'set_species': set_species1,
                'set_species_label': set_species2,
                'set_atomgr_att': change_atomgr_att1,
                'set_atomgr_att_label': change_atomgr_att2,
                'set_inpchanges': set_inpchanges1,
                'shift_value': shift_value1,
                'shift_value_species_label': shift_value_species_label1,
                'set_nkpts': set_nkpts1,
                'set_kpath': set_kpath1,
                'add_num_to_att': add_num_to_att1
            }

            workingtree = fleurinp_tree_copy
            if schema_tree:
                #xmlschema_doc = etree.parse(new_fleurinp._schema_file_path)
                xmlschema = etree.XMLSchema(schema_tree)

            for task in modification_tasks:
                try:
                    action = actions[task[0]]
                except KeyError:
                    raise ValueError('Unknown task {}'.format(task[0]))

                workingtree = action(workingtree, *task[1:])

            if schema_tree:
                try:
                    xmlschema.assertValid(clear_xml(workingtree))
                except etree.DocumentInvalid:
                    print('changes were not valid: {}'.format(modification_tasks))
                    raise

            return workingtree

        def get_avail_actions(self):
            """
            Returns the allowed functions from FleurinpModifier
            """
            outside_actions = {
                'xml_set_attribv_occ': self.xml_set_attribv_occ,
                'xml_set_first_attribv': self.xml_set_first_attribv,
                'xml_set_all_attribv': self.xml_set_all_attribv,
                'xml_set_text': self.xml_set_text,
                'xml_set_text_occ': self.xml_set_text_occ,
                'xml_set_all_text': self.xml_set_all_text,
                'create_tag': self.create_tag,
                'replace_tag': self.replace_tag,
                'delete_tag': self.delete_tag,
                'delete_att': self.delete_att,
                'set_species': self.set_species,
                'set_species_label': self.set_species_label,
                'set_atomgr_att': self.set_atomgr_att,
                'set_atomgr_att_label': self.set_atomgr_att_label,
                'set_inpchanges': self.set_inpchanges,
                'shift_value': self.shift_value,
                'shift_value_species_label': self.shift_value_species_label,
                'set_nkpts': self.set_nkpts,
                'add_num_to_att': self.add_num_to_att
            }
            return outside_actions

        def xml_set_attribv_occ(self, xpathn, attributename, attribv, occ=None, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.xml_set_attribv_occ()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute
            :param attributename: an attribute name
            :param attribv: an attribute value which will be set
            :param occ: a list of integers specifying number of occurrence to be set
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            """
            if occ is None:
                occ = [0]
            self._tasks.append(('xml_set_attribv_occ', xpathn, attributename, attribv, occ, create))

        def xml_set_first_attribv(self, xpathn, attributename, attribv, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.xml_set_first_attribv()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute
            :param attributename: an attribute name
            :param attribv: an attribute value which will be set
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            """
            self._tasks.append(('xml_set_first_attribv', xpathn, attributename, attribv, create))

        def xml_set_all_attribv(self, xpathn, attributename, attribv, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.xml_set_all_attribv()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute
            :param attributename: an attribute name
            :param attribv: an attribute value which will be set
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            """
            self._tasks.append(('xml_set_all_attribv', xpathn, attributename, attribv, create))

        def xml_set_text(self, xpathn, text, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.xml_set_text()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute
            :param text: text to be set
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            """
            self._tasks.append(('xml_set_text', xpathn, text, create))

        def xml_set_text_occ(self, xpathn, text, create=False, occ=0):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.xml_set_text_occ()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute
            :param text: text to be set
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            :param occ: an integer specifying number of occurrence to be set
            """
            self._tasks.append(('xml_set_text_occ', xpathn, text, create, occ))

        def xml_set_all_text(self, xpathn, text, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.xml_set_all_text()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute
            :param text: text to be set
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            """
            self._tasks.append(('xml_set_all_text', xpathn, text, create))

        def create_tag(self, xpath, newelement, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.create_tag()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path where to place a new tag
            :param newelement: a tag name to be created
            :param create: if True and there is no given xpath in the FleurinpData, creates it
            """
            self._tasks.append(('create_tag', xpath, newelement, create))

        def delete_att(self, xpath, attrib):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.delete_att()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the attribute to be deleted
            :param attrib: the name of an attribute
            """
            self._tasks.append(('delete_att', xpath, attrib))

        def delete_tag(self, xpath):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.delete_tag()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the tag to be deleted
            """
            self._tasks.append(('delete_tag', xpath))

        def replace_tag(self, xpath, newelement):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.replace_tag()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: a path to the tag to be replaced
            :param newelement: a new tag
            """
            self._tasks.append(('replace_tag', xpath, newelement))

        def set_species(self, species_name, attributedict, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.set_species()` to
            the list of tasks that will be done on the FleurinpData.

            :param species_name: a path to the tag to be replaced
            :param attributedict: attribute dictionary to be set into the specie
            :param create: if True and there is no given specie in the FleurinpData, creates it
            """
            self._tasks.append(('set_species', species_name, attributedict, create))

        def set_species_label(self, at_label, attributedict, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.set_species_label()` to
            the list of tasks that will be done on the FleurinpData.

            :param at_label: Atom label which specie will be set
            :param attributedict: attribute dictionary to be set into the specie
            :param create: if True and there is no given specie in the FleurinpData, creates it
            """
            self._tasks.append(('set_species_label', at_label, attributedict, create))

        def set_atomgr_att(self, attributedict, position=None, species=None, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.change_atomgr_att()` to
            the list of tasks that will be done on the FleurinpData.

            :param species_name: a path to the tag to be replaced
            :param attributedict: attribute dictionary to be set into the atom group
            :param create: if True and there is no given atom group in the FleurinpData, creates it
            """
            self._tasks.append(('set_atomgr_att', attributedict, position, species, create))

        def set_atomgr_att_label(self, attributedict, atom_label, create=False):
            """
            Appends a :func:`~aiida_fleur.tools.xml_util.change_atomgr_att_label()` to
            the list of tasks that will be done on the FleurinpData.

            :param attributedict: a new tag
            :param atom_label: Atom label which atom group will be set
            :param create: if True and there is no given atom group in the FleurinpData, creates it
            """
            self._tasks.append(('set_atomgr_att_label', attributedict, atom_label, create))

        def set_inpchanges(self, change_dict):
            """
            Appends a :py:func:`~aiida_fleur.tools.xml_util.set_inpchanges()` to
            the list of tasks that will be done on the FleurinpData.

            :param change_dict: a dictionary with changes

            An example of change_dict::

                change_dict = {'itmax' : 1,
                               'l_noco': True,
                               'ctail': False,
                               'l_ss': True}
            """
            self._tasks.append(('set_inpchanges', change_dict))

        def shift_value(self, change_dict, mode='abs'):
            """
            Appends a :py:func:`~aiida_fleur.tools.xml_util.shift_value()` to
            the list of tasks that will be done on the FleurinpData.

            :param change_dict: a dictionary with changes
            :param mode: 'abs' if change given is absolute, 'rel' if relative

            An example of change_dict::

                change_dict = {'itmax' : 1, dVac = -2}
            """
            self._tasks.append(('shift_value', change_dict, mode))

        def shift_value_species_label(self, label, att_name, value, mode='abs'):
            """
            Appends a :py:func:`~aiida_fleur.tools.xml_util.shift_value_species_label()` to
            the list of tasks that will be done on the FleurinpData.

            :param label: a label of an atom
            :param att_name: attrubute name of a specie
            :param value: value to set
            :param mode: 'abs' if change given is absolute, 'rel' if relative
            """
            self._tasks.append(('shift_value_species_label', label, att_name, value, mode))

        def set_nkpts(self, count, gamma='F'):
            """
            Appends a :py:func:`~aiida_fleur.tools.xml_util.set_nkpts()` to
            the list of tasks that will be done on the FleurinpData.
            """
            self._tasks.append(('set_nkpts', count, gamma))

        def set_kpath(self, kpath, count, gamma='F'):
            """
            Appends a :py:func:`~aiida_fleur.tools.xml_util.set_kpath()` to
            the list of tasks that will be done on the FleurinpData.
            """
            self._tasks.append(('set_kpath', kpath, count, gamma))

        def add_num_to_att(self, xpathn, attributename, set_val, mode='abs', occ=None):
            """
            Appends a :py:func:`~aiida_fleur.tools.xml_util.add_num_to_att()` to
            the list of tasks that will be done on the FleurinpData.

            :param xpathn: an xml path to the attribute to change
            :param attributename: a name of the attribute to change
            :param set_val: a value to be added/multiplied to the previous value
            :param mode: 'abs' if to add set_val, 'rel' if multiply
            :param occ: a list of integers specifying number of occurrence to be set
            """
            if occ is None:
                occ = [0]
            self._tasks.append(('add_num_to_att', xpathn, attributename, set_val, mode, occ))

        def validate(self):
            """
            Extracts the schema-file.
            Makes a test if all the changes lead to an inp.xml file that is validated against the
            schema.

            :return: a lxml tree representing inp.xml with applied changes
            """
            with self._original.open(key='inp.xml') as inpxmlfile:
                tree = etree.parse(inpxmlfile)

            try:  # could be not found or on another computer...
                xmlschema_tree = etree.parse(self._original._schema_file_path)
                with_schema = True
            except BaseException:
                with_schema = False
                print('No schema file found')
                return
            if with_schema:
                tree = self.apply_modifications(tree, self._tasks, schema_tree=xmlschema_tree)
            return tree

        def show(self, display=True, validate=False):
            """
            Applies the modifications and displays/prints the resulting ``inp.xml`` file.
            Does not generate a new
            :class:`~aiida_fleur.data.fleurinp.FleurinpData` object.

            :param display: a boolean that is True if resulting ``inp.xml`` has to be printed out
            :param validate: a boolean that is True if changes have to be validated

            :return: a lxml tree representing inp.xml with applied changes
            """

            if validate:
                tree = self.validate()
            else:
                with self._original.open(path='inp.xml') as inpxmlfile:
                    tree = etree.parse(inpxmlfile)
                tree = self.apply_modifications(tree, self._tasks)

            if display:
                xmltreestring = etree.tostring(tree, xml_declaration=True, pretty_print=True)
                print(xmltreestring)
            return tree

        def changes(self):
            """
            Prints out all changes given in a
            :class:`~aiida_fleur.data.fleurinpmodifier.FleurinpModifier` instance.
            """
            from pprint import pprint
            pprint(self._tasks)
            return self._tasks

        def freeze(self):
            """
            This method applies all the modifications to the input and
            returns a new stored fleurinpData object.

            :return: stored :class:`~aiida_fleur.data.fleurinp.FleurinpData` with applied changes
            """
            modifications = DataFactory('dict')(dict={'tasks': self._tasks})
            modifications.description = 'Fleurinpmodifier Tasks and inputs of these.'
            modifications.label = 'Fleurinpdata modifications'
            # This runs in a inline calculation to keep provenance
            out = modify_fleurinpdata(original=self._original,
                                      modifications=modifications,
                                      metadata={
                                          'label': 'fleurinp modifier',
                                          'description': 'This calcfunction modified an Fleurinpdataobject'
                                      })
            return out

        def undo(self, revert_all=False):
            """
            Cancels the last change or all of them

            :param revert_all: set True if need to cancel all the changes, False if the last one.
            """
            if revert_all:
                self._tasks = []
            else:
                if self._tasks:
                    self._tasks.pop()
                    #del self._tasks[-1]
            return self._tasks

    @cf
    def modify_fleurinpdata(original, modifications):
        """
        A CalcFunction that performs the modification of the given FleurinpData and stores
        the result in a database.

        :param original: a FleurinpData to be modified
        :param modifications: a python dictionary of modifications in the form of {'task': ...}

        :returns new_fleurinp: a modified FleurinpData that is stored in a database
        """

        # copy
        # get schema
        # read in inp.xml
        # add modifications
        # validate
        # save inp.xml
        # store new fleurinp (copy)
        from aiida_fleur.tools.xml_util import clear_xml

        new_fleurinp = original.clone()
        modification_tasks = modifications.get_dict()['tasks']

        xmlschema_doc = etree.parse(new_fleurinp._schema_file_path)
        xmlschema = etree.XMLSchema(xmlschema_doc)
        parser = etree.XMLParser(attribute_defaults=True, remove_blank_text=True)
        with new_fleurinp.open(path='inp.xml', mode='r') as inpxmlfile:
            tree = etree.parse(inpxmlfile, parser)

        try:
            xmlschema.assertValid(clear_xml(tree))
        except etree.DocumentInvalid:
            print('Input file is not validated against the schema')
            raise

        new_fleurtree = FleurinpModifier.apply_modifications(fleurinp_tree_copy=tree, modification_tasks=modification_tasks)

        # To include object store storage this prob has to be done differently

        inpxmlfile_new = inpxmlfile.name.replace('inp.xml', 'temp_inp.xml')
        inpxmlfile.close()

        new_fleurtree.write(inpxmlfile_new, pretty_print=True)

        new_fleurinp.del_file('inp.xml')
        new_fleurinp._add_path(str(inpxmlfile_new), 'inp.xml')
        os.remove(inpxmlfile_new)

        # default label and description
        new_fleurinp.label = 'mod_fleurinp'
        new_fleurinp.description = 'Fleurinpdata with modifications (see inputs of modify_fleurinpdata)'

        return new_fleurinp