io.py 9.91 KB
Newer Older
Qusai Al Shidi's avatar
Qusai Al Shidi committed
1
2
3
4
5
6
7
8
"""Input/Output tools

Input/Output
============

Here are tools to read and write files relating to SWMF.

TODO: Move pandas dependancy elsewhere.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
9
10
11
12
13
14
15
"""

import datetime as dt
import numpy as np


def read_wdc_ae(wdc_filename):
16
17
    """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
18
19
20
21
22

    Args:
        wdc_filename (str): Filename of wdc data from
                            http://wdc.kugi.kyoto-u.ac.jp/
    Returns:
23
24
25
        dict: {"time" (datetime.datetime): list of datetime objects
                                           corresponding to time in UT.
               "AL", "AE", "AO", "AU" (int): Auroral indeces.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
26
27
              }
    """
28
29
30
31
    data = {'AL': {'Time': [], 'Index': []},
            'AE': {'Time': [], 'Index': []},
            'AO': {'Time': [], 'Index': []},
            'AU': {'Time': [], 'Index': []}}
Qusai Al Shidi's avatar
Qusai Al Shidi committed
32
33
34
35
    with open(wdc_filename) as wdc_file:
        for line in wdc_file:
            ind_data = line.split()
            for minute in range(60):
36
                # TODO: Use .zfill()?
Qusai Al Shidi's avatar
Qusai Al Shidi committed
37
38
                str_min = str(minute)
                if minute < 10:
39
                    str_min = '0' + str_min
Qusai Al Shidi's avatar
Qusai Al Shidi committed
40
41
42
                time = dt.datetime.strptime(ind_data[1][:-5]
                                            + ind_data[1][7:-2]
                                            + str_min,
43
44
45
                                            '%y%m%d%H%M')
                data[ind_data[1][-2:]]['Time'] += [time]
                data[ind_data[1][-2:]]['Index'] += [int(ind_data[3+minute])]
Qusai Al Shidi's avatar
Qusai Al Shidi committed
46
47
48
    return data


49
50
51
52
53
54
55
56
def read_wdc_asy_sym(wdc_filename):
    """Docstring
    """

    return_data = {'ASY': {'Time': [], 'Index': []},
                   'SYM': {'Time': [], 'Index': []}}

    with open(wdc_filename) as wdc_file:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
57
58
59
60
61
        header = wdc_file.readline()[:6]
        assert header == 'ASYSYM', ('File does not seem to be'
                                    + 'an ASY/SYM file from wdc.'
                                    + 'First six characters: '
                                    + header)
62
63
64
65
66
67
        for line in wdc_file:
            data = line.split()

    return_data


Qusai Al Shidi's avatar
Qusai Al Shidi committed
68
69
70
71
72
def read_omni_csv(filename, filtering=False, **kwargs):
    """Take an OMNI csv file from cdaweb.sci.gsfc.nasa.gov
    and turn it into a pandas.DataFrame.

    Args:
73
74
75
76
        fnames (dict): dict with filenames from omni .lst files.
                       The keys must be: density, temperature,
                                         magnetic_field, velocity
        filtering (bool): default=False Remove points where the value
Qusai Al Shidi's avatar
Qusai Al Shidi committed
77
                          is >sigma (default: sigma=3) from mean.
78
79
80
81
        **kwargs:
            coarseness (int): default=3, Number of standard deviations
                              above which to remove if filtering=True.
            clean (bool): default=True, Clean the omni data of bad data points
Qusai Al Shidi's avatar
Qusai Al Shidi committed
82

83
84
    Returns:
        pandas.DataFrame: object with solar wind data
Qusai Al Shidi's avatar
Qusai Al Shidi committed
85
86
87
88
89
90
91
92

    Make sure to download the csv files with cdaweb.sci.gsfc.nasa.gov
    the header seperated into a json file for safety.

    This only tested with OMNI data specifically.


    """
Qusai Al Shidi's avatar
Qusai Al Shidi committed
93
    import pandas as pd
Qusai Al Shidi's avatar
Qusai Al Shidi committed
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    # Read the csv files and set the index to dates
    colnames = ['Time', 'Bx [nT]', 'By [nT]', 'Bz [nT]',
                'Vx [km/s]', 'Vy [km/s]', 'Vz [km/s]',
                'Rho [n/cc]', 'T [K]']
    with open(filename, 'r') as datafile:
        data = pd.read_csv(datafile, names=colnames, skiprows=1)
    data.set_index(pd.to_datetime(data[data.columns[0]]), inplace=True)
    data.drop(columns=data.columns[0], inplace=True)
    data.index.name = "Time [UT]"

    # clean bad data
    if kwargs.get('clean', True):
        data["By [nT]"] = data["By [nT]"][data["By [nT]"].abs() < 80.]
        data["Bx [nT]"] = data["Bx [nT]"][data["Bx [nT]"].abs() < 80.]
        data["Bz [nT]"] = data["Bz [nT]"][data["Bz [nT]"].abs() < 80.]
        data["Rho [n/cc]"] = data["Rho [n/cc]"][data["Rho [n/cc]"] < 500.]
        data["Vx [km/s]"] = data["Vx [km/s]"][data["Vx [km/s]"].abs() < 2000.]
        data["Vz [km/s]"] = data["Vz [km/s]"][data["Vz [km/s]"].abs() < 1000.]
        data["Vy [km/s]"] = data["Vy [km/s]"][data["Vy [km/s]"].abs() < 1000.]
        data["T [K]"] = data["T [K]"][data["T [K]"] < 1.e7]

    if filtering:
116
        _coarse_filtering(data, kwargs.get('coarseness', 3))
Qusai Al Shidi's avatar
Qusai Al Shidi committed
117
118
119
    return data.interpolate().bfill().ffill()


120
def _coarse_filtering(data, coarseness=3):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
121
122
123
124
125
126
127
    """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]


128
def write_imf_input(data, outfilename="IMF.dat", enable_rb=True, **kwargs):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
129
130
131
132
133
    """Writes the pandas.DataFrame into an input file
    that SWMF can read as input IMF (IMF.dat).

    Args:
        data: pandas.DataFrame object with solar wind data
Qusai Al Shidi's avatar
Qusai Al Shidi committed
134
135
136
137
        outfilename: The output file name for ballistic solar wind data.
                     (default: "IMF.dat")
        enable_rb: Enables solar wind input for the radiation belt model.
                   (default: True)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
138
139
140
141
142

    Other paramaters:
        gse: (default=False)
            Use GSE coordinate system for the file instead of GSM default.
    """
Qusai Al Shidi's avatar
Qusai Al Shidi committed
143

Qusai Al Shidi's avatar
Qusai Al Shidi committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    # Generate BATS-R-US solar wind input file
    with open(outfilename, 'w') as outfile:
        outfile.write("CSV files downloaded from ")
        outfile.write("https://cdaweb.gsfc.nasa.gov/\n")
        if kwargs.get('gse', False):
            outfile.write("#COOR\nGSE\n")
        outfile.write("yr mn dy hr min sec msec bx by bz vx vy vz dens temp\n")
        outfile.write("#START\n")
        for index, rows in data.iterrows():
            outfile.write(index.strftime("%Y %m %d %H %M %S") + ' ')
            outfile.write(index.strftime("%f")[:3] + ' ')
            outfile.write(str(rows['Bx [nT]'])[:7] + ' ')
            outfile.write(str(rows['By [nT]'])[:7] + ' ')
            outfile.write(str(rows['Bz [nT]'])[:7] + ' ')
            outfile.write(str(rows['Vx [km/s]'])[:7] + ' ')
            outfile.write(str(rows['Vy [km/s]'])[:7] + ' ')
            outfile.write(str(rows['Vz [km/s]'])[:7] + ' ')
            outfile.write(str(rows['Rho [n/cc]'])[:7] + ' ')
            outfile.write(str(rows['T [K]'])[:7] + ' ')
            outfile.write('\n')
    # Generate RBE model solar wind input file
    if enable_rb:
        with open("RB.SWIMF", 'w') as rbfile:
            # Choose first element as t=0 header (unsure if this is safe)
            rbfile.write(data.index[0].strftime("%Y, %j, %H ")
                         + "! iyear, iday, ihour corresponding to t=0\n")
            swlag_time = None
            if swlag_time in kwargs:
                rbfile.write(str(kwargs["swlag_time"]) + "  "
                             + "! swlag time in seconds "
                             + "for sw travel to subsolar\n")
            # Unsure what 11902 means but following example file
            rbfile.write("11902 data                   "
                         + "P+ NP NONLIN    P+ V (MOM)\n")
            # Quantity header
            rbfile.write("dd mm yyyy hh mm ss.ms           "
                         + "#/cc          km/s\n")
            for index, rows in data.iterrows():
                rbfile.write(index.strftime("%d %m %Y %H %M %S.%f")
                             + "     "
                             + str(rows['Rho [n/cc]'])[:8]
                             + "     "
                             # Speed magnitude
                             + str(np.sqrt(rows['Vx [km/s]']**2
                                           + rows['Vy [km/s]']**2
                                           + rows['Vz [km/s]']**2))[:8]
                             + '\n')


193
194
def read_gm_log(filename, colnames=None, index_time=True):
    """Make a dictionary out of the indeces outputted
Qusai Al Shidi's avatar
Qusai Al Shidi committed
195
196
197
198
    from the GM model log.

    Args:
        filename (str): The filename as a string.
199
200
201
202
203
204
        colnames ([str]): (default: None) Supply the name of the columns.
                                          If None it will use second line
                                          of log file.
        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
205
    Returns:
206
        Dictionary of the log file
Qusai Al Shidi's avatar
Qusai Al Shidi committed
207
208

    Examples:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
209
210
        To plot AL and Dst get the log files
        ```
211
212
        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
213

214
215
        # Plot AL indeces
        plt.plot(geo['Time [UT]', geo['AL'])
Qusai Al Shidi's avatar
Qusai Al Shidi committed
216
        ```
Qusai Al Shidi's avatar
Qusai Al Shidi committed
217

218
    """
Qusai Al Shidi's avatar
Qusai Al Shidi committed
219

Qusai Al Shidi's avatar
Qusai Al Shidi committed
220
    # If column names were not specified
221
222
223
224
225
226
227
228
    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
229
            colnames = logfile.readline().split()
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
        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
                for col, data in enumerate(line.split()):
                    return_data[colnames[col]].append(data)

        # datetime index
        if index_time:
            return_data['Time [UT]'] = []
            for row, year in enumerate(return_data[colnames[1]]):
                return_data['Time [UT]'].append(
                    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