paramin.py 5.44 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', ...]
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
        change['#SOLARWINDFILE'] = [['T', 'UseSolarWindFile'],
34
                                    ['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


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

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

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

    Raises:
91
        ValueError: When the COMMAND is not found.
92
93

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

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


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


def _get_command(line):
151
    """Returns the 'COMMAND' if on line.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

    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