swmfpy.py 11.1 KB
Newer Older
Qusai Al Shidi's avatar
Qusai Al Shidi committed
1
#!/usr/bin/env python3
Qusai Al Shidi's avatar
Qusai Al Shidi committed
2
"""swmfpy contains tools for the Space Weather Modeling Framework
Qusai Al Shidi's avatar
Qusai Al Shidi committed
3
4
5
6
7
8
9
"""
__author__ = "Qusai Al Shidi"
__license__ = "MIT"
__version__ = "0.0.1"
__maintainer__ = "Qusai Al Shidi"
__email__ = "qusai@umich.edu"

Qusai Al Shidi's avatar
Qusai Al Shidi committed
10
import logging
Qusai Al Shidi's avatar
Qusai Al Shidi committed
11
import datetime as dt
Qusai Al Shidi's avatar
Qusai Al Shidi committed
12
13
14
15
import numpy as np
import pandas as pd


Qusai Al Shidi's avatar
Qusai Al Shidi committed
16
17
18
def read_wdc_ae(wdc_filename):
    """Read an AE WDC text file into a dictionary of arrays.

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


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

52
    Args:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
53
54
55
56
57
58
59
60
61
62
63
64
        fnames: dict with filenames from omni .lst files. The keys must be:
            density, temperature, magnetic_field, velocity
        filtering: default=False Remove points where the value
                          is >sigma (default: sigma=3) from mean.

    Returns: pandas.DataFrame object with solar wind data

    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.

65
    Other Args:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
66
67
68
        coarseness: default=3, Number of standard deviations above which to
                    remove if filtering=True.
        clean: default=True, Clean the omni data of bad data points
Qusai Al Shidi's avatar
Qusai Al Shidi committed
69
70
71

    """
    # Read the csv files and set the index to dates
Qusai Al Shidi's avatar
Qusai Al Shidi committed
72
73
74
    colnames = ['Time', 'Bx [nT]', 'By [nT]', 'Bz [nT]',
                'Vx [km/s]', 'Vy [km/s]', 'Vz [km/s]',
                'Rho [n/cc]', 'T [K]']
Qusai Al Shidi's avatar
Qusai Al Shidi committed
75
    with open(filename, 'r') as datafile:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
76
        data = pd.read_csv(datafile, names=colnames, skiprows=1)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    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:
        coarse_filtering(data, kwargs.get('coarseness', 3))
    return data.interpolate().bfill().ffill()

Qusai Al Shidi's avatar
Qusai Al Shidi committed
96

Qusai Al Shidi's avatar
Qusai Al Shidi committed
97
98
99
100
101
102
103
104
def coarse_filtering(data, coarseness=3):
    """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]


Qusai Al Shidi's avatar
Qusai Al Shidi committed
105
106
107
108
def write_sw_input(data, outfilename="IMF.dat", enable_rb=True, **kwargs):
    """Writes the pandas.DataFrame into an input file
    that SWMF can read as input IMF (IMF.dat).

109
    Args:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
110
111
112
113
114
115
116
        data: pandas.DataFrame object with solar wind data
        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)

    Other paramaters:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
117
        gse: (default=False)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
118
119
120
121
            Use GSE coordinate system for the file instead of GSM default.
    """
    # Generate BATS-R-US solar wind input file
    with open(outfilename, 'w') as outfile:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
122
123
124
        outfile.write("CSV files downloaded from ")
        outfile.write("https://cdaweb.gsfc.nasa.gov/\n")
        if kwargs.get('gse', False):
Qusai Al Shidi's avatar
Qusai Al Shidi committed
125
126
127
128
129
130
            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] + ' ')
Qusai Al Shidi's avatar
Qusai Al Shidi committed
131
132
133
134
135
136
137
138
            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] + ' ')
Qusai Al Shidi's avatar
Qusai Al Shidi committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
            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
Qusai Al Shidi's avatar
Qusai Al Shidi committed
164
165
                                           + rows['Vy [km/s]']**2
                                           + rows['Vz [km/s]']**2))[:8]
Qusai Al Shidi's avatar
Qusai Al Shidi committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
                             + '\n')


def convert(infile, outfile="IMF.dat"):
    """Start the process of conversion of OMNI file to
    SWMF IMF input file.
    """
    # Write out the header
    outfile.write("OMNI file downloaded from \
                   https://omniweb.gsfc.nasa.gov/\n")
    outfile.write("yr mn dy hr min sec msec bx by bz vx vy vz dens temp\n")
    outfile.write("#START\n")
    # Begin conversion line by line
    for line in infile:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
180
        date = dt.datetime.strptime(line[:14], "%Y %j %H %M")
Qusai Al Shidi's avatar
Qusai Al Shidi committed
181
182
183
184
185
186
187
188
        correctline = date.strftime("%Y %m %d %H %M %S")\
            + ' 000' + line[14:]
        outfile.write(correctline)
    # Close files
    outfile.close()
    infile.close()


Qusai Al Shidi's avatar
Qusai Al Shidi committed
189
190
191
192
def read_gm_log(filename, colnames=None, index_by_time=True):
    """Make a pandas.DataFrame out of the Dst indeces outputted
    from the GM model log.

193
194
195
196
    Args:
        filename (str): The filename as a string.
        colnames (list(str)): Supply the name of the columns
        index_by_time (bool): Change the index to a time index
Qusai Al Shidi's avatar
Qusai Al Shidi committed
197
    Returns:
198
        pandas.DataFrame of the log file
Qusai Al Shidi's avatar
Qusai Al Shidi committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222

    Examples:
        # To plot AL and Dst get the log files
        geo = swmfpy.read_gm_log("run/GM/IO2/geoindex_e20140215-100500.log")
        dst = swmfpy.read_gm_log("run/GM/IO2/log_e20140215-100500.log")
        # Then plot using pandas features
        dst["dst_sm"].plot.line()
        geo["AL"].plot.line()
    """
    # If column names were not specified
    if not colnames:
        with open(filename, 'r') as logfile:
            # Usually the names are in the second line
            logfile.readline()
            colnames = logfile.readline().split()
    data = pd.read_fwf(filename, header=None, skiprows=2, names=colnames)
    if index_by_time:
        data.set_index(pd.to_datetime({'year': data['year'],
                                       'month': data['mo'],
                                       'day': data['dy'],
                                       'h': data['hr'],
                                       'm': data['mn'],
                                       's': data['sc']}),
                       inplace=True)
Qusai Al Shidi's avatar
Qusai Al Shidi committed
223
        data.index.names = ['Time [UT]']
Qusai Al Shidi's avatar
Qusai Al Shidi committed
224
225
226
    return data


Qusai Al Shidi's avatar
Qusai Al Shidi committed
227
def paramin_replace(input_file, replace, output_file="PARAM.in"):
228
    """Replace values for the parameters in a PARAM.in file.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
229

230
231
232
    Note, if you have repeat commands this will replace all the repeats.

    Args:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
233
234
235
        input_file (str): String of PARAM.in file name.
        replace (dict): Dictionary of strs with format
                        replace = {"#COMMAND": ["value", "comments", ...]}
236
237
238
                        This is case sensitive.
        output_file (str): (default "PARAM.in") The output file to write to.
                           A value of None will not output a file.
239
240
241
    Returns:
        A list of lines of the PARAM.in file that would be outputted.

Qusai Al Shidi's avatar
Qusai Al Shidi committed
242
243
244
245
    Examples:
        ```
        change["#SOLARWINDFILE"] = [["T", "UseSolarWindFile"],
                                    ["new_imf.dat", "NameSolarWindFile"]]
246
        # This will overwrite PARAM.in
Qusai Al Shidi's avatar
Qusai Al Shidi committed
247
248
        swmfpy.paramin_replace("PARAM.in.template", change)
        ```
Qusai Al Shidi's avatar
Qusai Al Shidi committed
249
    """
250
    # TODO This will replace all for repeat commands.
Qusai Al Shidi's avatar
Qusai Al Shidi committed
251
252
253
    logger = logging.getLogger()  # For debugging
    # 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
254
255
256
257
258
        command = None  # Top level #COMMAND
        # Compile lines in a list before editing/writing it
        lines = list(paramin)
        for line_num, line in enumerate(lines):
            words = line.split()
259
260
            if words and words[0] in replace.keys():
                command = words[0]  # Current command
Qusai Al Shidi's avatar
Qusai Al Shidi committed
261
262
263
264
265
266
267
268
                for param, value in enumerate(replace[command]):
                    newline = ""
                    for text in value:
                        newline += text + "\t\t\t"
                    logger.info("Replacing:\n" + line
                                + "with:\n" + newline)
                    # Lines will be replaced in order
                    lines[line_num+param+1] = newline + '\n'
Qusai Al Shidi's avatar
Qusai Al Shidi committed
269
    # Write the PARAM.in file
270
271
    if output_file is None:
        return lines  # Break if None output_file (not default behaviour)
272
    with open(output_file, 'w') as outfile:
Qusai Al Shidi's avatar
Qusai Al Shidi committed
273
274
275
        for line in lines:
            outfile.write(line)
    return lines