Commit 3c757873 authored by Qusai Al Shidi's avatar Qusai Al Shidi 💬
Browse files

Working and tested. Made documentation. Added new feature swmfpy.write_imf_from_omni

parent e4f879c8
# Table of Contents
* [swmfpy](#.swmfpy)
* [swmfpy.tools](#.swmfpy.tools)
* [swmfpy.io](#.swmfpy.io)
* [read\_wdc\_ae](#.swmfpy.io.read_wdc_ae)
* [read\_wdc\_asy\_sym](#.swmfpy.io.read_wdc_asy_sym)
* [read\_gm\_log](#.swmfpy.io.read_gm_log)
* [write\_imf\_from\_omni](#.swmfpy.write_imf_from_omni)
* [swmfpy.web](#.swmfpy.web)
* [get\_omni\_data](#.swmfpy.web.get_omni_data)
* [download\_magnetogram\_hmi](#.swmfpy.web.download_magnetogram_hmi)
* [download\_magnetogram\_adapt](#.swmfpy.web.download_magnetogram_adapt)
* [swmfpy.io](#.swmfpy.io)
* [write\_imf\_input](#.swmfpy.io.write_imf_input)
* [read\_wdc\_ae](#.swmfpy.io.read_wdc_ae)
* [read\_wdc\_asy\_sym](#.swmfpy.io.read_wdc_asy_sym)
* [read\_gm\_log](#.swmfpy.io.read_gm_log)
* [swmfpy.paramin](#.swmfpy.paramin)
* [replace\_command](#.swmfpy.paramin.replace_command)
* [read\_command](#.swmfpy.paramin.read_command)
* [swmfpy.tools](#.swmfpy.tools)
<a name=".swmfpy"></a>
## swmfpy
......@@ -34,114 +36,36 @@ These are not automatically imported. Might have extra dependancies.
*None yet.*
<a name=".swmfpy.tools"></a>
## swmfpy.tools
Tools to be used in swmfpy functions and classes. Some of the functions are
*hidden functions*.
<a name=".swmfpy.io"></a>
## swmfpy.io
### Input/Output tools
Here are tools to read and write files relating to SWMF.
<a name=".swmfpy.io.read_wdc_ae"></a>
#### read\_wdc\_ae
<a name=".swmfpy.write_imf_from_omni"></a>
#### write\_imf\_from\_omni
```python
read_wdc_ae(wdc_filename)
write_imf_from_omni(time_from, time_to, filename='IMF.dat', **kwargs)
```
Read an auroral electrojet (AE) indeces from Kyoto's World Data Center
text file into a dictionary of lists.
Writes an IMF.dat file for the geospace model runs for a specific time
period.
**Arguments**:
- `wdc_filename` _str_ - Filename of wdc data from
http://wdc.kugi.kyoto-u.ac.jp/
**Returns**:
- `dict` - Auroral indeces 'AL', 'AE', 'AO', 'AU'
- `datetime.datetime` - 'times'
- `int` - 'values'
<a name=".swmfpy.io.read_wdc_asy_sym"></a>
#### read\_wdc\_asy\_sym
```python
read_wdc_asy_sym(wdc_filename)
```
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.
**Arguments**:
- `wdc_filename` _str_ - Relative filename (or file handle no.) to read.
**Returns**:
- `dict` - of values. {'[ASY-SYM]-[D-H]': 'times': [], 'values': []}
**Examples**:
```python
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.
<a name=".swmfpy.io.read_gm_log"></a>
#### read\_gm\_log
```python
read_gm_log(filename, colnames=None, dtypes=None, index_time=True)
```
Make a dictionary out of the indeces outputted
from the GM model log.
**Arguments**:
- `filename` _str_ - The relative filename as a string. (or file handle no.)
- `colnames` _[str]_ - (default: None) Supply the name of the columns.
If None it will use second line
of log file.
- `dtypes` _[types]_ - (default: None) Provide types for the columns, if
None then all will be float.
- `index_time` _bool_ - (default: True) Make a column of dt.datetime objects
in dictionary key 'Time [UT]'.
**Returns**:
- `dict` - Dictionary of the log file
- `time_from` _datetime.datetime_ - Time to begin omni data retrieval
- `time_to` _datetime.datetime_ - Time to end omni data retrieval
- `filename` _str_ - The filename for the dat file, defaults to 'IMF.dat'.
**kwargs:
see `swmfpy.io.write_imf_input()` and `swmfpy.web.get_omni_data()`
**Examples**:
To plot AL and Dst get the log files
Using this function is simple:
```python
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')
# Plot AL indeces
plt.plot(geo['times', geo['AL'])
import swmfpy
import datetime as dt
times = (dt.datetime(2014, 2, 2), dt.datetime(2014, 2, 4))
# Usually the kwargs are unecessary
swmfpy.write_imf_input(*times)
# Sometimes this
swmfpy.write_imf_input(*times, filename='run/IMF.dat')
```
<a name=".swmfpy.web"></a>
......@@ -173,6 +97,14 @@ use swmfpy.io.read_omni_data().
data that you want to receive.
- `time_to` _datetime.datetime_ - The end time of the solar wind data
you want to receive.
**kwargs:
- `original_colnames` _bool_ - Use the original column names from the
spdf specification. The alternative is
nicer and shorter names. Defaults to
False.
- `resolution` _str_ - (default: 'high') Here you can choose 'high' or
'low' resolution omni data. Some columns appear
in one but not the other.
**Returns**:
......@@ -191,6 +123,10 @@ use swmfpy.io.read_omni_data().
storm_end = datetime.datetime(year=2000, month=2, day=15)
data = swmfpy.web.get_omni_data(time_from=storm_start,
time_to=storm_end)
# or for low res
data = swmfpy.web.get_omni_data(time_from=storm_start,
time_to=storm_end,
resolution='low')
```
<a name=".swmfpy.web.download_magnetogram_hmi"></a>
......@@ -310,6 +246,151 @@ pattern: adapt4[0,1]3*yyyymmddhh
download_dir='./mymaps/')
```
<a name=".swmfpy.io"></a>
## swmfpy.io
### Input/Output tools
Here are tools to read and write files relating to SWMF.
<a name=".swmfpy.io.write_imf_input"></a>
#### write\_imf\_input
```python
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)
**Arguments**:
- `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.
- `AssertionError` - If inputs aren't prepared properly (key names)
**Examples**:
```python
from swmfpy.io import write_imf_input
# Prepare imf dictionary: imf_data
write_imf_input(imf_data, filename='run/IMF.dat')
```
<a name=".swmfpy.io.read_wdc_ae"></a>
#### read\_wdc\_ae
```python
read_wdc_ae(wdc_filename)
```
Read an auroral electrojet (AE) indeces from Kyoto's World Data Center
text file into a dictionary of lists.
**Arguments**:
- `wdc_filename` _str_ - Filename of wdc data from
http://wdc.kugi.kyoto-u.ac.jp/
**Returns**:
- `dict` - Auroral indeces 'AL', 'AE', 'AO', 'AU'
- `datetime.datetime` - 'times'
- `int` - 'values'
<a name=".swmfpy.io.read_wdc_asy_sym"></a>
#### read\_wdc\_asy\_sym
```python
read_wdc_asy_sym(wdc_filename)
```
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.
**Arguments**:
- `wdc_filename` _str_ - Relative filename (or file handle no.) to read.
**Returns**:
- `dict` - of values. {'[ASY-SYM]-[D-H]': 'times': [], 'values': []}
**Examples**:
```python
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.
<a name=".swmfpy.io.read_gm_log"></a>
#### read\_gm\_log
```python
read_gm_log(filename, colnames=None, dtypes=None, index_time=True)
```
Make a dictionary out of the indeces outputted
from the GM model log.
**Arguments**:
- `filename` _str_ - The relative filename as a string. (or file handle no.)
- `colnames` _[str]_ - (default: None) Supply the name of the columns.
If None it will use second line
of log file.
- `dtypes` _[types]_ - (default: None) Provide types for the columns, if
None then all will be float.
- `index_time` _bool_ - (default: True) Make a column of dt.datetime objects
in dictionary key 'Time [UT]'.
**Returns**:
- `dict` - Dictionary of the log file
**Examples**:
To plot AL and Dst get the log files
```python
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')
# Plot AL indeces
plt.plot(geo['times', geo['AL'])
```
<a name=".swmfpy.paramin"></a>
## swmfpy.paramin
......@@ -331,7 +412,7 @@ Note, if you have repeat commands this will replace all the repeats.
**Arguments**:
- `parameters` _dict_ - Dictionary of strs with format
replace = '#COMMAND': ['value', 'comments', ...]
replace = '\#COMMAND': ['value', 'comments', ...]
This is case sensitive.
- `input_file` _str_ - String of PARAM.in file name.
- `output_file` _str_ - (default 'PARAM.in') The output file to write to.
......@@ -350,7 +431,7 @@ Note, if you have repeat commands this will replace all the repeats.
**Examples**:
```python
change['#SOLARWINDFILE'] = [['T', 'UseSolarWindFile'],
change['\#SOLARWINDFILE'] = [['T', 'UseSolarWindFile'],
['new_imf.dat', 'NameSolarWindFile']]
# This will overwrite PARAM.in
swmfpy.paramin.replace('PARAM.in.template', change)
......@@ -365,12 +446,12 @@ read_command(command, paramin_file='PARAM.in', **kwargs)
Get parameters of a certain command in PARAM.in file.
This will find the #COMMAND and return a list of
This will find the \#COMMAND and return a list of
values for the parameters.
**Arguments**:
- `command` _str_ - This is the #COMMAND you're looking for.
- `command` _str_ - This is the \#COMMAND you're looking for.
- `paramin_file` _str_ - (default: 'PARAM.in') The file in which you're
looking for the command values.
**kwargs:
......@@ -380,20 +461,20 @@ values for the parameters.
**Returns**:
- `list` - Values found for the #COMMAND in file. Index 0 is
#COMMAND and the values follow (1 for first argument...)
- `list` - Values found for the \#COMMAND in file. Index 0 is
\#COMMAND and the values follow (1 for first argument...)
**Raises**:
- `ValueError` - When the #COMMAND is not found.
- `ValueError` - When the \#COMMAND is not found.
**Examples**:
```python
start_time = swmfpy.paramin.read_command('#STARTTIME')
end_time = swmfpy.paramin.read_command('#ENDTIME')
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])
```
......@@ -402,3 +483,9 @@ values for the parameters.
this, try using the `num_of_values` keyword. This is helpful if your
PARAM.in is comment heavy.
<a name=".swmfpy.tools"></a>
## swmfpy.tools
Tools to be used in swmfpy functions and classes. Some of the functions are
*hidden functions*.
......@@ -4,5 +4,5 @@
# run `pip install pydoc-markdown --user`
# then run this script in the project root directory
pydoc-markdown -v -I. --render-toc -m swmfpy -m swmfpy.web -m swmfpy.io -m swmfpy.paramin -m swmfpy.tools | sed 's/\\043/#/g' - > DOCUMENTATION.markdown
pydoc-markdown --py3 -v -I. --render-toc -m swmfpy -m swmfpy.web -m swmfpy.io -m swmfpy.paramin -m swmfpy.tools | sed 's/\\043/#/g' - > DOCUMENTATION.markdown
......@@ -27,3 +27,38 @@ from . import paramin
from . import io
from . import web
assert sys.version_info >= (3, 6), "swmfpy requires Python >=3.6. Sorry :(."
def write_imf_from_omni(time_from, time_to, filename='IMF.dat', **kwargs):
"""Writes an IMF.dat file for the geospace model runs for a specific time
period.
Args:
time_from (datetime.datetime): Time to begin omni data retrieval
time_to (datetime.datetime): Time to end omni data retrieval
filename (str): The filename for the dat file, defaults to 'IMF.dat'.
**kwargs:
see #swmfpy.io.write_imf_input() and #swmfpy.web.get_omni_data()
Examples:
Using this function is simple:
```python
import swmfpy
import datetime as dt
times = (dt.datetime(2014, 2, 2), dt.datetime(2014, 2, 4))
# Usually the kwargs are unecessary
swmfpy.write_imf_input(*times)
# Sometimes this
swmfpy.write_imf_input(*times, filename='run/IMF.dat')
```
"""
omni_data = web.get_omni_data(time_from, time_to, **kwargs)
commands = ['#COOR', 'GSE']
if kwargs.get('commands', None):
kwargs['commands'] += commands
else:
kwargs['commands'] = commands
column_keys = ['times',
'bx', 'by_gse', 'bz_gse', 'vx_gse', 'vy_gse', 'vz_gse',
'density', 'temperature']
io.write_imf_input(omni_data, filename, column_keys=column_keys, **kwargs)
......@@ -27,16 +27,25 @@ def write_imf_input(imf_data, filename='IMF.dat', **kwargs):
Raises:
TypeError: If commands is not a list or tuple. It must be at least a
one element list of strings.
AssertionError: If inputs aren't prepared properly (key names)
Examples:
```python
from swmfpy.io import write_imf_input
# Prepare imf dictionary: imf_data
write_imf_input(imf_data, filename='run/IMF.dat')
```
"""
columns_dat = ['year', 'month', 'day', 'hour', 'min', 'sec', 'msec',
'bx', 'by', 'bz', 'vx', 'vy', 'vz',
'density', 'temperature']
columns_dict = ['times',
'bx', 'by', 'bz', 'vx', 'vy', 'vz',
'density', 'temperature']
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)
def _time_repr(time_raw):
'Represent time in a suitable format'
......@@ -59,15 +68,18 @@ def write_imf_input(imf_data, filename='IMF.dat', **kwargs):
file_imf.write(_make_line(command)+'\n')
# write dat file
file_imf.write('\n'+'\t'.join(columns_dat))
file_imf.write('#START\n')
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')
lines = []
for index, _time in enumerate(imf_data[columns_dict[0]]):
line = [_time_repr(_time)]
for key in columns_dict[1:]:
line += [str(imf_data[key][index])]
line = '\t'.join(line)
lines += line
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'
file_imf.writelines(lines)
......
......@@ -5,118 +5,10 @@
__author__ = 'Qusai Al Shidi'
__email__ = 'qusai@umich.edu'
import numpy as np
from scipy.optimize import newton, check_grad
def get_v_gsm_from_omni(dict_omni):
"""Adds gsm values to the omni dictionary
This is helpful for the omni dict created with #swmfpy.web.get_omni_data()
Args:
dict_omni(dict): Dictionary with np.arrays in 'bz', 'by', 'bz_gse',
'by_gse', 'vx_gse', 'vy_gse'
Returns:
dict_omni with added 'vz' and 'vy' keys for gsm coordinates.
Raises:
KeyError: If above values don't exist.
Examples:
```python
import datetime as dt
from swmfpy.web import get_omni_data
from swmfpy.toos import get_v_gsm_from_omni
time_period = (dt.datetime(2012, 2, 1), dt.datetime(2012, 4, 5))
data = get_v_gsm_from_omni(get_omni_data(*time_period))
# You should now have the omni dict with gsm values
plot(data['times'], data['vz'])
```
"""
# Author: Qusai Al Shidi
# Email: qusai@umich.edu
b_gsm = [dict_omni['bx'], dict_omni['by'], dict_omni['bz']]
b_gse = [dict_omni['bx'], dict_omni['by_gse'], dict_omni['bz_gse']]
v_gse = [dict_omni['vx_gse'], dict_omni['vy_gse'], dict_omni['vz_gse']]
def check_prime(func, fprime, *args):
'Make sure my derivatives are correct'
assert check_grad(func, fprime, [0.1], *args) < 1.e-3, 'bad grad'
def rot_y(theta, vec):
'rotate y in yz plane'
return vec[1]*np.cos(theta)-vec[2]*np.sin(theta)
def rot_y_prime(theta, vec):
'rotate y in yz plane'
return -vec[1]*np.sin(theta)-vec[2]*np.cos(theta)
check_prime(rot_y, rot_y_prime, (1, 1, 1))
def rot_y_prime2(theta, vec):
'rotate y in yz plane'
return -vec[1]*np.cos(theta)+vec[2]*np.sin(theta)
check_prime(rot_y_prime, rot_y_prime2, (1, 1, 1))
def rot_z(theta, vec):
'rotate z in yz plane'
return vec[1]*np.sin(theta)+vec[2]*np.cos(theta)
def rot_z_prime(theta, vec):
'rotate z in yz plane'
return vec[1]*np.cos(theta)-vec[2]*np.sin(theta)
check_prime(rot_z, rot_z_prime, (1, 1, 1))
def rot_z_prime2(theta, vec):
'rotate z in yz plane'
return -vec[1]*np.sin(theta)-vec[2]*np.cos(theta)
check_prime(rot_z_prime, rot_z_prime2, (1, 1, 1))
def opt_func(theta, vec_p, vec):
'The two functions to optimize'
opt_y = vec_p[1] - rot_y(theta, vec)
opt_z = vec_p[2] - rot_z(theta, vec)
return opt_y + opt_z
def opt_func_prime(theta, _, vec):
'The two functions to optimize'
opt_y = -rot_y_prime(theta, vec)
opt_z = -rot_z_prime(theta, vec)
return opt_y + opt_z
check_prime(opt_func, opt_func_prime, (1, 1, 1), (1, 1, 1))
def opt_func_prime2(theta, _, vec):
'The two functions to optimize'
opt_y = -rot_y_prime2(theta, vec)
opt_z = -rot_z_prime2(theta, vec)
return opt_y + opt_z
check_prime(opt_func_prime, opt_func_prime2, (1, 1, 1), (1, 1, 1))
angle = newton(opt_func, np.zeros_like(b_gsm[0]), args=(b_gsm, b_gse),
fprime=opt_func_prime, fprime2=opt_func_prime2, tol=1.e-2)
dict_omni.update({
'vx': dict_omni['vx_gse'],
'vy': rot_y(angle, v_gse),
'vz': rot_z(angle, v_gse),
})
return dict_omni
def _import_error_string(library):
return (
'Error importing drms. '
f'Error importing {library}. '
f'Maybe try `pip install -U {library} --user ` .'
)
......
......@@ -9,8 +9,7 @@ __author__ = 'Qusai Al Shidi'
__email__ = 'qusai@umich.edu'
import datetime as dt
import ftplib