cp2kinputengine.py 6.19 KB
Newer Older
1
2
#! /usr/bin/env python
# -*- coding: utf-8 -*-
Lauri Himanen's avatar
Lauri Himanen committed
3
from collections import defaultdict
4
from cp2kparser.generics.nomadlogging import *
Lauri Himanen's avatar
Lauri Himanen committed
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50


#===============================================================================
class CP2KInputEngine(object):
    """Used to parse out a CP2K input file.

    When given a file handle to a CP2K input file, this class attemts to parse
    out it's structure into an accessible object tree. Because the input file
    has such a clearly defined structure (unlike the output file of CP2K), it
    is better to use a dedicated parser instead of regular expressions.
    """
    def __init__(self, cp2k_parser):
        """
        Args:
            cp2k_parser: Instance of a CP2KParser or it's subclass. Allows
            access to e.g. unified file reading methods.
        """
        self.cp2k_parser = cp2k_parser
        self.root_section = None

    def parse_input(self):
        """Parses the given CP2K input string"""

        # The input file should be quite small, so just get the entire contents
        inp = self.cp2k_parser.get_file_contents("input")

        root_section = InputSection('CP2K_INPUT')
        section_stack = [root_section]

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

            if line.upper().startswith('&END'):
                s = section_stack.pop()
            elif line[0] == '&':
                parts = line.split(' ', 1)
                name = parts[0][1:]
                if len(parts) > 1:
                    s = InputSection(name=name, params=parts[1].strip())
                else:
                    s = InputSection(name=name)
                section_stack[-1].subsections[name.upper()].append(s)
                section_stack.append(s)
            else:
51
52
53
54
                split = line.split(' ', 1)
                keyword_name = split[0]
                keyword_value = split[1]
                section_stack[-1].keywords[keyword_name].append(keyword_value)
Lauri Himanen's avatar
Lauri Himanen committed
55
56
57
58
59
60

        self.root_section = root_section

    def get_subsection(self, path, index=0):
        return self.root_section.get_subsection(path, index)

61
62
63
64
65
66
67
68
69
70
71
72
73
    def get_keyword(self, path, index=0):
        split = path.rsplit('/', 1)
        section_path = split[0]
        keyword = split[1]
        section = self.root_section.get_subsection(section_path, index)
        if section is not None:
            return section.get_keyword(keyword)

    def get_parameter(self, path, index=0):
        section = self.root_section.get_subsection(path, index)
        if section is not None:
            return section.get_parameter()

Lauri Himanen's avatar
Lauri Himanen committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

#===============================================================================
class InputSection(object):
    """Represents a section in a CP2K input file"""
    def __init__(self, name, params=None):
        self.name = name.upper()
        self.params = params
        self.keywords = defaultdict(list)
        self.subsections = defaultdict(list)

    def write(self):
        """Outputs input section as string"""
        output = []
        for name, k_list in self.keywords.iteritems():
            for value in k_list:
                output.append(value)
        for name, s_list in self.subsections.iteritems():
            for s in s_list:
                if s.params:
                    output.append('&%s %s' % (s.name, s.params))
                else:
                    output.append('&%s' % s.name)
                for l in s.write():
                    output.append('   %s' % l)
                output.append('&END %s' % s.name)
        return output

    def get_subsection(self, path, index=0):
        """Finds a subsection specified by a string where subsections are
        separated by a slash. If multiple subsections are found with the same
        path, the one specified by the given index (default 0) is returned.

        Example: get_subsection("FORCE_EVAL/PRINT/FORCES")

        Args:
            path: String indicating the path to the subsection
            index: In case of repeating subsections, return the one specified
            by this index.

        Returns:
            The InputSection object if found.

        """
        parts = path.upper().split('/', 1)
        candidates = self.subsections.get(parts[0])    # [s for s in self.subsections if s.name == parts[0]]
        if not candidates:
            print_debug("Subsection '{}' not found.".format(parts[0]))
            return None
        elif len(candidates) > 1:
            print_warning("Multiple subsections with the same name found with name '{}' If no index is given, the first occurence in the input file is returned.".format(parts[0]))
        try:
            subsection = candidates[index]
        except IndexError:
            print_error("Invalid subsection index given.")

        if len(parts) == 1:
            return subsection
        return subsection.get_subsection(parts[1])

    def get_keyword(self, keyword, index=0):
        """Finds a keyword specified by a string. If multiple keywords are found with the same
        name, the one specified by the given index (default 0) is returned.

        Args:
            keyword: String indicating the name of the keyword. The name is the
            first word in the line.
            index: In case of repeating keywords, return the one specified
            by this index.

        Returns:
            The keyword value (everything else than the first word on the line).

        """
        candidates = self.keywords.get(keyword)    # [s for s in self.subsections if s.name == parts[0]]
        if not candidates:
            print_debug("No keywords with name '{}' found in subsection '{}'".format(keyword, self.name))
            return None
        elif len(candidates) > 1:
            print_warning("Multiple keywords with the same name found with name '{}' If no index is given, the first occurence in the input file is returned.".format(parts[0]))
        try:
            result = candidates[index]
        except IndexError:
            print_error("Invalid keyword index given.")
        return result

    def get_parameter(self):
160
161
        if self.params is None:
            print_debug("The section '{}' has no parameters set".format(self.name))
Lauri Himanen's avatar
Lauri Himanen committed
162
        return self.params