cp2kinputparser.py 3.12 KB
Newer Older
1
import os
Lauri Himanen's avatar
Lauri Himanen committed
2
import logging
3
import cPickle as pickle
Lauri Himanen's avatar
Lauri Himanen committed
4
logger = logging.getLogger(__name__)
Lauri Himanen's avatar
Lauri Himanen committed
5
6
7


#===============================================================================
8
class CP2KInputParser(object):
Lauri Himanen's avatar
Lauri Himanen committed
9
10
11
    """Used to parse out a CP2K input file.

    When given a file handle to a CP2K input file, this class attemts to parse
12
    out it's structure into an accessible object tree.
Lauri Himanen's avatar
Lauri Himanen committed
13
    """
14
    def __init__(self):
Lauri Himanen's avatar
Lauri Himanen committed
15
        self.root_section = None
16
        self.input_tree = None
Lauri Himanen's avatar
Lauri Himanen committed
17

18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    def parse(self, inp):
        """Parses a CP2K input file into an object tree.

        Return an object tree represenation of the input augmented with the
        default values and lone keyword values from the cp2k_input.xml file
        which is version specific. Keyword aliases are also mapped to the same data.

        To query the returned tree use the following functions:
            get_keyword("GLOBAL/PROJECT_NAME")
            get_parameter("GLOBAL/PRINT")
            get_default_keyword("FORCE_EVAL/SUBSYS/COORD")

        Args:
            inp: A string containing the contents of a CP2K input file. The
            input file can be stored as string as it isn't that big.

        Returns:
            The input as an object tree.
36
        """
37
38
39
40
        # See if version is setup
        if self.input_tree is None:
            logger.error("Please setup the CP2K version before parsing")
            return
Lauri Himanen's avatar
Lauri Himanen committed
41

42
        section_stack = []
Lauri Himanen's avatar
Lauri Himanen committed
43
44
45

        for line in inp.split('\n'):
            line = line.split('!', 1)[0].strip()
46
47

            # Skip empty lines
Lauri Himanen's avatar
Lauri Himanen committed
48
49
50
            if len(line) == 0:
                continue

51
            # Section ends
Lauri Himanen's avatar
Lauri Himanen committed
52
            if line.upper().startswith('&END'):
53
                section_stack.pop()
54
            # Section starts
Lauri Himanen's avatar
Lauri Himanen committed
55
56
57
            elif line[0] == '&':
                parts = line.split(' ', 1)
                name = parts[0][1:]
58
59
60
61
62
63
64
65
66
                section_stack.append(name)

                # Form the path
                path = ""
                for index, item in enumerate(section_stack):
                    if index != 0:
                        path += '/'
                    path += item

67
68
69
                # Mark the section as accessed.
                self.input_tree.set_section_accessed(path)

70
                # Save the section parameters
Lauri Himanen's avatar
Lauri Himanen committed
71
                if len(parts) > 1:
72
                    self.input_tree.set_parameter(path, parts[1].strip())
73
            # Contents (keywords, default keywords)
Lauri Himanen's avatar
Lauri Himanen committed
74
            else:
75
76
77
                split = line.split(' ', 1)
                keyword_name = split[0]
                keyword_value = split[1]
78
                self.input_tree.set_keyword(path + "/" + keyword_name, keyword_value)
Lauri Himanen's avatar
Lauri Himanen committed
79

80
        return self.input_tree
81

Lauri Himanen's avatar
Lauri Himanen committed
82
    def setup_version(self, version_number):
83
84
85
86
        """ The pickle file which contains preparsed data from the
        cp2k_input.xml is version specific. By calling this function before
        parsing the correct file can be found.
        """
87
88
89
        pickle_path = os.path.dirname(__file__) + "/cp2kinputenginedata/cp2k_{}/cp2k_input_tree.pickle".format(version_number)
        input_tree_pickle_file = open(pickle_path, 'rb')
        self.input_tree = pickle.load(input_tree_pickle_file)