Note: The default ITS GitLab runner is a shared resource and is subject to slowdowns during heavy usage.
You can run your own GitLab runner that is dedicated just to your group if you need to avoid processing delays.

paramin.py 5.47 KB
Newer Older
Qusai Al Shidi's avatar
Qusai Al Shidi committed
1
""" These tools are to help script the editing of PARAM.in files.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
2
"""
Qusai Al Shidi's avatar
Qusai Al Shidi committed
3
4
__all__ = [
    'read_command',
Qusai Al Shidi's avatar
Qusai Al Shidi committed
5
    'replace_command'
Qusai Al Shidi's avatar
Qusai Al Shidi committed
6
    ]
7
8
9
__author__ = 'Qusai Al Shidi'
__email__ = 'qusai@umich.edu'

Qusai Al Shidi's avatar
Qusai Al Shidi committed
10
11
import logging

Qusai Al Shidi's avatar
Qusai Al Shidi committed
12

13
def replace_command(parameters, input_file, output_file='PARAM.in'):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
14
15
16
17
18
    """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
19
        parameters (dict): Dictionary of strs with format
20
                        replace = {'#COMMAND': ['value', 'comments', ...]}
Qusai Al Shidi's avatar
Qusai Al Shidi committed
21
                        This is case sensitive.
22
23
        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
24
25
26
27
                           A value of None will not output a file.
    Returns:
        A list of lines of the PARAM.in file that would be outputted.

Qusai Al Shidi's avatar
Qusai Al Shidi committed
28
29
30
    Raises:
        TypeError: If a value given couldn't be converted to string.

Qusai Al Shidi's avatar
Qusai Al Shidi committed
31
    Examples:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
32
        ```python
33
34
        change['#SOLARWINDFILE'] = [['T', 'UseSolarWindFile'],
                                    ['new_imf.dat', 'NameSolarWindFile']]
Qusai Al Shidi's avatar
Qusai Al Shidi committed
35
        # This will overwrite PARAM.in
36
        swmfpy.paramin.replace('PARAM.in.template', change)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
37
38
        ```
    """
39
40
41
    # Author: Qusai Al Shidi
    # Email: qusai@umich.edu

Qusai Al Shidi's avatar
Qusai Al Shidi committed
42
    logger = logging.getLogger()  # For debugging
43

Qusai Al Shidi's avatar
Qusai Al Shidi committed
44
45
    # Read and replace paramin file by making a temp list
    with open(input_file, 'rt') as paramin:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
46

Qusai Al Shidi's avatar
Qusai Al Shidi committed
47
48
49
        # Compile lines in a list before editing/writing it
        lines = list(paramin)
        for line_num, line in enumerate(lines):
50
51

            # If the current command is what we want
Qusai Al Shidi's avatar
Qusai Al Shidi committed
52
            command = _get_command(line)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
53
            if command in parameters.keys():
54

Qusai Al Shidi's avatar
Qusai Al Shidi committed
55
                for param, value in enumerate(parameters[command]):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
56
57
58
59
60
                    newline = _make_line(value)
                    logger.info('Replacing: %s\n with: %s\n',
                                line, newline)
                    # Lines will be replaced in order
                    lines[line_num+param+1] = newline + '\n'
61

Qusai Al Shidi's avatar
Qusai Al Shidi committed
62
63
64
65
66
67
    # 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)
68

Qusai Al Shidi's avatar
Qusai Al Shidi committed
69
    return lines
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92


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:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
93
        ```python
94
95
96
97
        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])
Qusai Al Shidi's avatar
Qusai Al Shidi committed
98
99
        ```

100
101
102
103
104
105
106
107
108
109
110
    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 = []
Qusai Al Shidi's avatar
Qusai Al Shidi committed
111
112
        command_found = False  # to know if worked
        in_command = False  # when after command needed
113
        for line in paramin:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
114
            if _get_command(line) == command:
115
                logger.info('Found command: %s', command)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
116
117
                command_found = True
                in_command = True
118
                return_values.append(command)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
119
            elif in_command and _get_command(line):  # Make sure not out of cmd
Qusai Al Shidi's avatar
Qusai Al Shidi committed
120
121
122
                in_command = False
            elif line.split() and in_command:
                value = line.split()[0]
123
124
125
126
127
                logger.info('Value added: %s', value)
                return_values.append(value)

        # Error handling
        # Unable to find #COMMAND
Qusai Al Shidi's avatar
Qusai Al Shidi committed
128
        if not command_found:
129
130
131
132
133
134
135
136
            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
Qusai Al Shidi's avatar
Qusai Al Shidi committed
137
138


Qusai Al Shidi's avatar
Qusai Al Shidi committed
139
# HIDDEN FUNCTIONS
Qusai Al Shidi's avatar
Qusai Al Shidi committed
140
141
142
143
def _make_line(value):
    """Makes the paramin line based on value type recursively"""
    if isinstance(value, str):
        return value
Qusai Al Shidi's avatar
Qusai Al Shidi committed
144
    if isinstance(value, list):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
145
146
        return '\t\t\t'.join([_make_line(v) for v in value])
    return str(value)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165


def _get_command(line):
    """Returns the '#COMMAND' if on line.

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

    Returns:
        (str): '#COMMAND' if found and None if not.
    """
    if isinstance(line, (str, list, tuple)):  # Raises type error otherwise
        if isinstance(line, str):
            to_check = line.split()
            if to_check and to_check[0].startswith('#'):
                return to_check[0]
            return None
        return _get_command(line[0])
    return None