paramin.py 5.7 KB
Newer Older
Qusai Al Shidi's avatar
Qusai Al Shidi committed
1
2
3
"""Tools to manipulate or create param.in files

PARAM.in Tools
Qusai Al Shidi's avatar
Qusai Al Shidi committed
4
--------------
Qusai Al Shidi's avatar
Qusai Al Shidi committed
5
6

These tools are to help script the editing of PARAM.in files.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
7
"""
Qusai Al Shidi's avatar
Qusai Al Shidi committed
8
9
10
11
__all__ = [
    'read_command',
    'replace_command'
    ]
12
13
14
__author__ = 'Qusai Al Shidi'
__email__ = 'qusai@umich.edu'

Qusai Al Shidi's avatar
Qusai Al Shidi committed
15
16
import logging

Qusai Al Shidi's avatar
Qusai Al Shidi committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def get_command(line):
    """Returns the '#COMMAND' if on line.

    Args:
        line (str): The line in the PARAM.in file.

    Returns:
        (str): '#COMMAND' if found and None if not.
    """
    to_check = line.split()
    if isinstance(to_check, (str, tuple, list)):
        if isinstance(to_check, str) and to_check.startswitch('#'):
            return to_check
        return get_command(to_check[0])
    return None

Qusai Al Shidi's avatar
Qusai Al Shidi committed
33

34
def replace_command(parameters, input_file, output_file='PARAM.in'):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
35
36
37
38
39
    """Replace values for the parameters in a PARAM.in file.

    Note, if you have repeat commands this will replace all the repeats.

    Args:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
40
        parameters (dict): Dictionary of strs with format
41
                        replace = {'#COMMAND': ['value', 'comments', ...]}
Qusai Al Shidi's avatar
Qusai Al Shidi committed
42
                        This is case sensitive.
43
44
        input_file (str): String of PARAM.in file name.
        output_file (str): (default 'PARAM.in') The output file to write to.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
45
46
47
48
49
50
                           A value of None will not output a file.
    Returns:
        A list of lines of the PARAM.in file that would be outputted.

    Examples:
        ```
51
52
        change['#SOLARWINDFILE'] = [['T', 'UseSolarWindFile'],
                                    ['new_imf.dat', 'NameSolarWindFile']]
Qusai Al Shidi's avatar
Qusai Al Shidi committed
53
        # This will overwrite PARAM.in
54
        swmfpy.paramin.replace('PARAM.in.template', change)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
55
56
        ```
    """
57
58
59
    # Author: Qusai Al Shidi
    # Email: qusai@umich.edu

Qusai Al Shidi's avatar
Qusai Al Shidi committed
60
61
    # TODO This will replace all for repeat commands.
    logger = logging.getLogger()  # For debugging
62

Qusai Al Shidi's avatar
Qusai Al Shidi committed
63
64
    # Read and replace paramin file by making a temp list
    with open(input_file, 'rt') as paramin:
65
        command = None  # Current command
Qusai Al Shidi's avatar
Qusai Al Shidi committed
66
67
68
69
        # Compile lines in a list before editing/writing it
        lines = list(paramin)
        for line_num, line in enumerate(lines):
            words = line.split()
70
71

            # If the current command is what we want
Qusai Al Shidi's avatar
Qusai Al Shidi committed
72
            if words and words[0] in parameters.keys():
73
74
75
                command = words[0]

                # Replace code
Qusai Al Shidi's avatar
Qusai Al Shidi committed
76
                for param, value in enumerate(parameters[command]):
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
                    newline = ''
                    # Allow additions of comments if list(str)
                    if isinstance(value, list):
                        for text in value:
                            newline += text + '\t\t\t'
                        logger.info('Replacing: %s\n with: %s\n',
                                    line, newline)
                        # Lines will be replaced in order
                        lines[line_num+param+1] = newline + '\n'
                    # Else just make a line
                    elif isinstance(value, str):
                        logger.info('Replacing: %s\n with: %s\n', line, value)
                        # Lines will be replaced in order
                        lines[line_num+param+1] = value + '\n'

Qusai Al Shidi's avatar
Qusai Al Shidi committed
92
93
94
95
96
97
    # Write the PARAM.in file
    if output_file is None:
        return lines  # Break if None output_file (not default behaviour)
    with open(output_file, 'w') as outfile:
        for line in lines:
            outfile.write(line)
98

Qusai Al Shidi's avatar
Qusai Al Shidi committed
99
    return lines
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
160
161
162
163
164
165


def read_command(command, paramin_file='PARAM.in', **kwargs):
    """Get parameters of a certain command in PARAM.in file.

    This will find the #COMMAND and return a list of values for the parameters.

    Args:
        command (str): This is the #COMMAND you're looking for.
        paramin_file (str): (default: 'PARAM.in') The file in which you're
                            looking for the command values.
        **kwargs:
            num_of_values (int): (default: None) Number of values to take from
                                 command.

    Returns:
        list: Values found for the #COMMAND in file. Index 0 is #COMMAND and
              the values follow (1 for first argument...)

    Raises:
        ValueError: When the #COMMAND is not found.

    Examples:
        start_time = swmfpy.paramin.read_command('#STARTTIME')
        end_time = swmfpy.paramin.read_command('#ENDTIME')
        print('Starting month is ', start_time[1])
        print('Ending month is ', end_time[1])

    This will treat all following lines as values for the command. To suppress
    this, try using the `num_of_values` keyword. This is helpful if your
    PARAM.in is comment heavy.
    """
    # Author: Qusai Al Shidi
    # Email: qusai@umich.edu

    logger = logging.getLogger()  # For debugging

    with open(paramin_file) as paramin:
        return_values = []
        found_command = False  # to know if worked
        at_command = False  # when after command needed
        for line in paramin:
            words = line.split()
            if words and words[0] == command:
                logger.info('Found command: %s', command)
                found_command = True
                at_command = True
                return_values.append(command)
            elif words and words[0][0] == '#':
                at_command = False
            elif words and at_command:
                value = words[0]
                logger.info('Value added: %s', value)
                return_values.append(value)

        # Error handling
        # Unable to find #COMMAND
        if not found_command:
            raise ValueError(command + ' not found.')

        # To ignore additional lines
        value_limit = kwargs.get('num_of_values', None)
        if value_limit:
            return_values = return_values[:value_limit+1]

        return return_values  # empty list might mean command not found