io.py 10.8 KB
Newer Older
1
"""### Input/Output tools
Qusai Al Shidi's avatar
Qusai Al Shidi committed
2
3

Here are tools to read and write files relating to SWMF.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
4
5
6
"""

import datetime as dt
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from .tools import _make_line


def write_imf_input(imf_data, filename='IMF.dat', **kwargs):
    """Creates the IMF.dat input file for the SWMF BATS-R-US geospace model.

    `imf_data` must be a dictionary of array_like objects with same length
    in data. In swmfpy Pythonic versions are always preferred so the 'times'
    must be `datetime.datetime` array.
    imf_data = dict(times, bx, by, bz, vx, vy, vz, density, temperature)

    Args:
        imf_data (dict): This dictionary contains the solar wind data.
        filename (str): (default: 'IMF.dat') Filename to write to.
        **kwargs:
            commands ([str]): (default: None) List of commands to write to imf
                              input file (indexed by line then by tabs on same
                              line). *Note*: Must be a list if have one command
                              str.

    Raises:
        TypeError: If commands is not a list or tuple. It must be at least a
                   one element list of strings.
30
        AssertionError: If inputs aren't prepared properly (key names)
31
32

    Examples:
33
34
35
36
37
38
        ```python
        from swmfpy.io import write_imf_input

        # Prepare imf dictionary: imf_data
        write_imf_input(imf_data, filename='run/IMF.dat')
        ```
39
40
41
42
43
    """

    columns_dat = ['year', 'month', 'day', 'hour', 'min', 'sec', 'msec',
                   'bx', 'by', 'bz', 'vx', 'vy', 'vz',
                   'density', 'temperature']
44
45
46
47
48
    column_keys = ['times',
                   'bx', 'by', 'bz', 'vx', 'vy', 'vz',
                   'density', 'temperature']
    if kwargs.get('column_keys', None):
        column_keys = kwargs.get('column_keys', column_keys)
49
50
51
52
53
54
55
56

    def _time_repr(time_raw):
        'Represent time in a suitable format'
        raw_str = dt.datetime.strftime(time_raw, '%Y %m %d %H %M %S %f')[:-3]
        return raw_str.split()

    with open(filename, 'w') as file_imf:
        # header
57
58
        file_imf.write('Made with swmfpy ')
        file_imf.write('(https://gitlab.umich.edu/swmf_software/swmfpy)\n\n')
59
60
61
62
63
64
65
66
67
68
69
70

        # write commands
        commands = kwargs.get('commands', None)
        if commands:
            try:
                isinstance(commands, (list, tuple))
            except TypeError:
                raise TypeError(f'{__name__}: commands must be list or tuple')
            for command in commands:
                file_imf.write(_make_line(command)+'\n')

        # write dat file
71
72
73
74
75
        def justified(str_list, width=7):
            'Returns justified'
            return [s.rjust(width) for s in str_list]
        file_imf.write('\n'+' '.join(justified(columns_dat)))
        file_imf.write('\n#START\n')
76
        lines = []
77
78
79
80
81
82
        for index, _time in enumerate(imf_data[column_keys[0]]):
            line = _time_repr(_time)
            for key in column_keys[1:]:
                line += [str(round(imf_data[key][index], 2))]
            line = ' '.join(justified(line))
            lines += line + '\n'
83
        file_imf.writelines(lines)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
84
85
86


def read_wdc_ae(wdc_filename):
87
88
    """Read an auroral electrojet (AE) indeces from Kyoto's World Data Center
       text file into a dictionary of lists.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
89
90

    Args:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
91

Qusai Al Shidi's avatar
Qusai Al Shidi committed
92
93
94
        wdc_filename (str): Filename of wdc data from
                            http://wdc.kugi.kyoto-u.ac.jp/
    Returns:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
95
96
97
        dict: Auroral indeces 'AL', 'AE', 'AO', 'AU'
            datetime.datetime: 'times'
            int: 'values'
Qusai Al Shidi's avatar
Qusai Al Shidi committed
98
    """
99
100
101
102
103
104
105
106

    # Initialize return data
    return_data = {'AL': {'times': [], 'values': []},
                   'AE': {'times': [], 'values': []},
                   'AO': {'times': [], 'values': []},
                   'AU': {'times': [], 'values': []}}

    # Open and make sure it is correct file
107
    with open(wdc_filename, 'rt') as wdc_file:
108
        header = wdc_file.readline()
109
110
        assert header[:8] == 'AEALAOAU', \
            'Does not seem to be a WDC AE file. First 8 chars: ' + header[:8]
111
112

        # Parse
Qusai Al Shidi's avatar
Qusai Al Shidi committed
113
        for line in wdc_file:
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
            data = line.split()
            year_suffix = int(data[1][:2])
            if year_suffix < 50:
                year = 2000 + year_suffix
            else:
                year = 1990 + year_suffix
            month = int(data[1][2:4])
            day = int(data[1][4:6])
            hour = int(data[1][7:9])
            index = data[1][-2:]
            values_60 = [int(val) for val in data[3:60+3]]

            # Fill
            for minute, value in enumerate(values_60):
                return_data[index]['values'].append(value)
                return_data[index]['times'].append(
                    dt.datetime(year, month, day, hour, minute))

    return return_data
Qusai Al Shidi's avatar
Qusai Al Shidi committed
133
134


135
def read_wdc_asy_sym(wdc_filename):
136
137
138
139
140
141
142
    """Reads a WDC file for ASY/SYM data.

    Reads an ASY/SYM file downloaded from
    http://wdc.kugi.kyoto-u.ac.jp/aeasy/index.html
    and puts it into a dictionary.

    Args:
143
        wdc_filename (str): Relative filename (or file handle no.) to read.
144
145

    Returns:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
146
        dict: of values. {'[ASY-SYM]-[D-H]': 'times': [], 'values': []}
147
148

    Examples:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
149
        ```python
150
151
152
153
154
155
156
157
158
159
        indeces = swmfpy.io.read_wdc_asy_sym('wdc.dat')
        # Plot data
        plt.plot(indeces['SYM-H']['times'],
                 indeces['SYM-H']['values'],
                 label='SYM-H [nT]'
                 )
        plt.xlabel('Time [UT]')
        ```

    Important to note if there is bad data it will be filled as None.
160
161
    """

162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    return_data = {
        'ASY-D': {
            'times': [],
            'values': [],
            },
        'ASY-H': {
            'times': [],
            'values': [],
            },
        'SYM-D': {
            'times': [],
            'values': [],
            },
        'SYM-H': {
            'times': [],
            'values': []
            }
        }
180
181

    with open(wdc_filename) as wdc_file:
182
183
184
185
186
187
188
189
        # Check for correct file
        header = wdc_file.readline()
        assert header[:12] == 'ASYSYM N6E01', ('File does not seem to be'
                                               + 'an ASY/SYM file from wdc.'
                                               + 'First 12 characters: '
                                               + header)
        return_data['edition'] = header[24:34]

190
        for line in wdc_file:
191
192
            # Parse
            year = int(line[12:14])
193
194
195
196
            if year < 70:  # Starts from 1970 but only gives 2 digits
                year += 2000
            else:
                year += 1900
197
198
199
200
201
202
203
            month = int(line[14:16])
            day = int(line[16:18])
            hour = int(line[19:21])
            comp = line[18]
            index = line[21:24]

            # Fill 60 min data
204
            data = line.split()
205
206
207
208
209
210
211
212
213
214
215
            values_60 = [int(val) for val in data[2:62]]
            for minute, value in enumerate(values_60):
                return_data[index+'-'+comp]['times'].append(
                    dt.datetime(year, month, day, hour, minute))
                # Check if data is bad
                if value != 99999:
                    return_data[index+'-'+comp]['values'].append(
                        value)
                else:
                    return_data[index+'-'+comp]['values'].append(
                        None)
216

217
    return return_data
218
219


220
def _coarse_filtering(data, coarseness=3):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
221
222
223
224
225
226
227
    """Applies coarse filtering to a pandas.DataFrame"""
    for column in data.columns:
        mean = data[column].abs().mean()
        sigma = data[column].std()
        data[column] = data[data[column].abs() < mean+coarseness*sigma][column]


228
def read_gm_log(filename, colnames=None, dtypes=None, index_time=True):
229
    """Make a dictionary out of the indeces outputted
Qusai Al Shidi's avatar
Qusai Al Shidi committed
230
231
232
    from the GM model log.

    Args:
233
        filename (str): The relative filename as a string. (or file handle no.)
234
235
236
        colnames ([str]): (default: None) Supply the name of the columns.
                                          If None it will use second line
                                          of log file.
237
238
        dtypes ([types]): (default: None) Provide types for the columns, if
                                          None then all will be float.
239
240
241
        index_time (bool): (default: True) Make a column of dt.datetime objects
                                           in dictionary key 'Time [UT]'.

Qusai Al Shidi's avatar
Qusai Al Shidi committed
242
    Returns:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
243
        dict: Dictionary of the log file
Qusai Al Shidi's avatar
Qusai Al Shidi committed
244
245

    Examples:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
246
        To plot AL and Dst get the log files
Qusai Al Shidi's avatar
Qusai Al Shidi committed
247
        ```python
248
249
        geo = swmfpy.io.read_gm_log('run/GM/IO2/geoindex_e20140215-100500.log')
        dst = swmfpy.io.read_gm_log('run/GM/IO2/log_e20140215-100500.log')
Qusai Al Shidi's avatar
Qusai Al Shidi committed
250

251
        # Plot AL indeces
252
        plt.plot(geo['times', geo['AL'])
Qusai Al Shidi's avatar
Qusai Al Shidi committed
253
        ```
Qusai Al Shidi's avatar
Qusai Al Shidi committed
254

255
    """
Qusai Al Shidi's avatar
Qusai Al Shidi committed
256

Qusai Al Shidi's avatar
Qusai Al Shidi committed
257
    # If column names were not specified
258
259
260
261
262
263
264
265
    return_data = {}
    with open(filename, 'r') as logfile:

        # Find column names and initialize
        description = logfile.readline()
        return_data['description'] = description
        # Usually the names are in the second line
        if not colnames:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
266
            colnames = logfile.readline().split()
267
268
269
270
271
272
        try:
            colnames = _fix_str_duplicates(colnames)
        except RuntimeWarning:
            print(f'{__name__}: Warning: '
                  + 'Found duplicates in column names. '
                  + 'Changes made to column names.')
273
274
275
276
277
278
        for col in colnames:
            return_data[col] = []

        # Fill data dictionary
        for line_num, line in enumerate(logfile):
            if line_num > 2:  # First two lines are usually metadata
279
                for col, data in enumerate(line.strip().split()):
280
281
282
283
                    if dtypes:
                        data = dtypes[col](data)
                    else:
                        data = float(data)
284
285
286
287
                    return_data[colnames[col]].append(data)

        # datetime index
        if index_time:
288
            return_data['times'] = []
289
            for row, year in enumerate(return_data[colnames[1]]):
290
                return_data['times'].append(
291
292
293
294
295
296
297
298
299
                    dt.datetime(int(year),
                                int(return_data[colnames[2]][row]),  # month
                                int(return_data[colnames[3]][row]),  # day
                                int(return_data[colnames[4]][row]),  # hour
                                int(return_data[colnames[5]][row]),  # min
                                int(return_data[colnames[6]][row]),  # sec
                                int(return_data[colnames[7]][row])))  # ms

    return return_data
300
301
302
303
304
305
306
307
308
309
310
311
312
313


def _fix_str_duplicates(str_list):
    """Returns a list and bool if a fix was made.
       The fix is adding an _[index] to avoid duplicates.
    """
    duplicate_found = False
    for index, string in enumerate(str_list):
        if str_list.count(string) > 1:
            duplicate_found = True
            str_list[index] = string + f'_{index}'
    if duplicate_found:
        raise RuntimeWarning
    return str_list