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.49 KB
Newer Older
1
2
3
"""### Editing PARAM.in files

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

Qusai Al Shidi's avatar
Qusai Al Shidi committed
12
13
import logging

Qusai Al Shidi's avatar
Qusai Al Shidi committed
14

15
def replace_command(parameters, input_file, output_file='PARAM.in'):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
16
17
18
19
20
    """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
21
        parameters (dict): Dictionary of strs with format
22
              replace = 'COMMAND': ['value', 'comments', ...]
23
              This is case sensitive.
24
25
        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
26
27
28
29
                           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
30
31
32
    Raises:
        TypeError: If a value given couldn't be converted to string.

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

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

Qusai Al Shidi's avatar
Qusai Al Shidi committed
46
47
    # 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
48

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

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

Qusai Al Shidi's avatar
Qusai Al Shidi committed
57
                for param, value in enumerate(parameters[command]):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
58
59
60
61
62
                    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'
63

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

Qusai Al Shidi's avatar
Qusai Al Shidi committed
71
    return lines
72
73
74
75
76


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

77
78
    This will find the COMMAND and return a list of
    values for the parameters.
79
80

    Args:
81
        command (str): This is the COMMAND you're looking for.
82
83
84
85
86
87
88
        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:
89
90
        list: Values found for the COMMAND in file. Index 0 is
        COMMAND and the values follow (1 for first argument...)
91
92

    Raises:
93
        ValueError: When the COMMAND is not found.
94
95

    Examples:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
96
        ```python
97
98
        start_time = swmfpy.paramin.read_command('\043STARTTIME')
        end_time = swmfpy.paramin.read_command('\043ENDTIME')
99
100
        print('Starting month is ', start_time[1])
        print('Ending month is ', end_time[1])
Qusai Al Shidi's avatar
Qusai Al Shidi committed
101
102
        ```

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


Qusai Al Shidi's avatar
Qusai Al Shidi committed
142
# HIDDEN FUNCTIONS
Qusai Al Shidi's avatar
Qusai Al Shidi committed
143
144
145
146
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
147
    if isinstance(value, list):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
148
149
        return '\t\t\t'.join([_make_line(v) for v in value])
    return str(value)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
150
151
152


def _get_command(line):
153
    """Returns the 'COMMAND' if on line.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
154
155
156
157
158

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

    Returns:
159
        (str): '\043COMMAND' if found and None if not.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
160
161
162
163
164
165
166
167
168
    """
    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