def test_avg_data(self):
"""Test avg_data."""
avg_data('aaabbbccc', width=2) # Test non-existent name
avg_data('test', width=2)
d = get_data('test-avg')
self.assertTrue((d[1] == [4.0, 11.5, 10.5]).all())
avg_data('test', width=2, overwrite=True) # Test overwrite
avg_data('test', dt=4.0, noremainder=False) # Test dt option
store_data('test',
data={
'x': [1., 2., 3., 4., 5., 6.],
'y': [3., 5., 8., -4., 20., 1.]
})
avg_data('test', width=2, new_names='aabb') # Test new_names
d = get_data('aabb')
# Test multiple names
avg_data(['test', 'aabb'], new_names='aaabbb', width=2)
dn = [[3., 5., 8.], [15., 20., 1.], [3., 5., 8.], [15., 20., 1.],
[23., 15., 28.], [15., 20., 1.]]
store_data('test1', data={'x': [1., 12., 13., 14., 15., 16.], 'y': dn})
avg_data('test1', width=2) # Test 3-d data
avg_data('test1', new_names='test2', dt=2.) # Test a reasonable dt
avg_data('test1', dt=-1.) # Test dt error
avg_data('test1', dt=1.e8) # Test dt error
d2 = get_data('test2')
self.assertTrue(len(d) > 0)
self.assertTrue(d2[1][-1][0] == 19.0)
def setUp(self):
"""Create a tplot variable to be used in tests."""
store_data('test',
data={
'x': [1., 2., 3., 4., 5., 6.],
'y': [3., 5., 8., 15., 20., 1.]
})
def test_timeclip(self):
"""Test time_clip."""
time_clip('aaabbbccc', 1577308800, 1577598800) # Test non-existent
tn = [
1577112800, 1577308800, 1577598800, 1577608800, 1577998800,
1587998800
]
dn = [3., 5., 8., 15., 20., 1.]
store_data('test1', data={'x': tn, 'y': dn})
time_clip('aaabbb', 1577308800, 1577598800)
time_clip('test1', 1577112800, 1577608800)
d = get_data('test1-tclip')
dd = d[1]
time_clip('test', 1577308800, 1577598800, new_names='name-clip')
time_clip(['test', 'name-clip'],
1577308800,
1577598800,
new_names='name1-ci')
time_clip('test', 1577308800, 1577598800, overwrite=1)
time_clip('test', 1577308800, 1577598800, new_names="testtest")
time_clip(['test', 'test1'],
1577308800,
1577598800,
new_names="testtest2")
time_clip('test1', 1677112800, 1577608800)
self.assertTrue((dd == [3., 5., 8., 15.]).all())
def ex_create():
# Delete any existing pytplot variables
pytplot.del_data()
# Create a sin wave plot
a = list(range(0, 101))
b = [2.0 / 100.0 * numpy.pi * s for s in a]
c = pyspedas.time_float('2017-01-01')
x = list()
y = list()
for i in range(len(b)):
x.append(c + 60.0 / (2 * numpy.pi) * 60.0 * b[i])
y.append(1000.0 * numpy.sin(b[i]))
# Store data
pytplot.store_data('sinx', data={'x': x, 'y': y})
# Apply tclip
pyspedas.tclip('sinx', -800.0, 800.0)
# Remove NaN values
pyspedas.tdeflag('sinx-clip')
# Interpolate
pyspedas.tinterpol(['sinx-clip-deflag'], ['sinx'], 'quadratic')
# Plot
pytplot.ylim('sinx', -1100.0, 1100.0)
pytplot.ylim('sinx-clip', -1100.0, 1100.0)
pytplot.ylim('sinx-clip-deflag', -1100.0, 1100.0)
pytplot.ylim('sinx-clip-deflag-itrp', -1100.0, 1100.0)
pytplot.tplot_options('title', 'Interpolation example')
pytplot.tplot(
['sinx', 'sinx-clip', 'sinx-clip-deflag', 'sinx-clip-deflag-itrp'])
def mms_pgs_mphigeo(mag_temp, pos_temp):
"""
Generates the 'mphigeo' transformation matrix
"""
pos_data = get_data(pos_temp)
if pos_data is None:
logging.error('Error with position data')
return
# the following is heisted from the IDL version
# transformation to generate other_dim dim for mphigeo from thm_fac_matrix_make
# All the conversions to polar and trig simplifies to this.
# But the reason the conversion is why this is the conversion that is done, is lost on me.
# The conversion swaps the x & y components of position, reflects over x=0,z=0 then projects into the xy plane
pos_conv = np.stack(
(-pos_data.y[:, 1], pos_data.y[:, 0], np.zeros(len(pos_data.times))))
pos_conv = np.transpose(pos_conv, [1, 0])
store_data(pos_temp, data={'x': pos_data.times, 'y': pos_conv})
# transform into GSE because the particles are in GSE
cotrans(name_in=pos_temp,
name_out=pos_temp,
coord_in='gei',
coord_out='gse')
# create orthonormal basis set
z_basis = tnormalize(mag_temp, return_data=True)
x_basis = tcrossp(z_basis, pos_temp, return_data=True)
x_basis = tnormalize(x_basis, return_data=True)
y_basis = tcrossp(z_basis, x_basis, return_data=True)
return (x_basis, y_basis, z_basis)
def ex_wavelet(plot=True):
"""Demonstrates how to use wavelets with pyspedas."""
# Delete any existing pytplot variables
pytplot.del_data()
# Create a pytplot variable.
t = np.arange(4000.)
y = np.sin(2 * np.pi * t / 32.)
y2 = np.sin(2 * np.pi * t / 64.)
y[1000:3000] = y2[1000:3000]
var = 'sin_wav'
time = time_float('2010-01-01') + 10 * t
pytplot.store_data(var, data={'x': time, 'y': y})
# Complex wavelet transformation.
powervar = wavelet(var, wavename='cmorl0.5-1.0')
# Also try the following and compare:
# powervar = wavelet(var, wavename='gaus1')
pvar = powervar[0]
# Plot.
pytplot.options(pvar, 'colormap', 'jet')
pytplot.options(pvar, 'ylog', True)
pytplot.options(pvar, 'ytitle', pvar)
pytplot.ylim(pvar, 0.001, 1.0)
if plot:
pytplot.tplot([var, pvar])
return 1
def tdeflag(names, method=None, flag=None, new_names=None, suffix=None):
old_names = pyspedas.tnames(names)
if len(old_names) < 1:
print('tdeflag error: No pytplot names were provided.')
return
if suffix is None:
suffix = '-deflag'
if flag is None:
flag = float('nan')
if new_names is None:
n_names = [s + suffix for s in old_names]
elif new_names == '':
n_names = old_names
else:
n_names = new_names
if len(n_names) != len(old_names):
n_names = [s + suffix for s in old_names]
for i in range(len(old_names)):
time, data = pytplot.get_data(old_names[i])
new_time = []
new_data = []
for j in range(len(time)):
if not numpy.isnan(data[j]):
new_time.append(time[j])
new_data.append(data[j])
pytplot.store_data(n_names[i], data={'x': new_time, 'y': new_data})
print('tdeflag was applied to: ' + n_names[i])
def tclip(names, ymin, ymax, flag=None, new_names=None, suffix=None):
old_names = pyspedas.tnames(names)
if len(old_names) < 1:
print('tclip error: No pytplot names were provided.')
return
if suffix is None:
suffix = '-clip'
if flag is None:
flag = float('nan')
if new_names is None:
n_names = [s + suffix for s in old_names]
elif new_names == '':
n_names = old_names
else:
n_names = new_names
if len(n_names) != len(old_names):
n_names = [s + suffix for s in old_names]
for i in range(len(old_names)):
alldata = pytplot.get_data(old_names[i])
time = alldata[0]
data = alldata[1]
new_data = numpy.array(data)
new_data[new_data <= ymin] = flag
new_data[new_data >= ymax] = flag
pytplot.store_data(n_names[i], data={'x': time, 'y': new_data})
print('tclip was applied to: ' + n_names[i])
def test_altitude_plot():
pytplot.netcdf_to_tplot(current_directory +
"/testfiles/g15_xrs_2s_20170619_20170619.nc",
time='time_tag')
pytplot.store_data(
'altitude',
data={
'x':
pytplot.data_quants['A_COUNT'].coords['time'].values,
'y':
np.arange(
0,
len(pytplot.data_quants['A_COUNT'].coords['time'].values),
step=1)
})
pytplot.link('A_COUNT', 'altitude')
pytplot.xlim('2017-06-19 02:00:00', '2017-06-19 04:00:00')
pytplot.ylim("A_COUNT", 17000, 18000)
pytplot.timebar('2017-06-19 03:00:00',
"A_COUNT",
color=(100, 255, 0),
thick=3)
pytplot.timebar('2017-06-19 03:30:00', "A_COUNT", color='g')
pytplot.options("A_COUNT", 'alt', 1)
pytplot.tplot(2, testing=True)
pytplot.tplot(2, testing=True, bokeh=True)
def subtract_median(names, new_names=None, suffix=None):
old_names = pyspedas.tnames(names)
if len(old_names) < 1:
print('Subtract Median error: No pytplot names were provided.')
return
if suffix is None:
suffix = '-m'
if new_names is None:
n_names = [s + suffix for s in old_names]
elif new_names == '':
n_names = old_names
else:
n_names = new_names
if len(n_names) != len(old_names):
n_names = [s + suffix for s in old_names]
for i in range(len(old_names)):
alldata = pytplot.get_data(old_names[i])
time = alldata[0]
data = alldata[1]
new_data = data-numpy.median(data, axis=0)
pytplot.store_data(n_names[i], data={'x': time, 'y': new_data})
print('Subtract Median was applied to: ' + n_names[i])
def mms_feeps_correct_energies(probes, data_rate, level='l2', suffix=''):
types = ['top', 'bottom']
sensors = range(1, 13)
units_types = ['intensity', 'count_rate', 'counts']
for probe in probes:
for sensor_type in types:
for sensor in sensors:
if sensor >= 6 and sensor <= 8:
species = 'ion'
else:
species = 'electron'
for units in units_types:
var_name = 'mms' + probe + '_epd_feeps_' + data_rate + '_' + level + '_' + species + '_' + sensor_type + '_' + units + '_sensorid_' + str(
sensor)
times, data = get_data(var_name + suffix)
energies = pytplot.data_quants[var_name +
suffix].spec_bins.values
energy_map = mms_feeps_energy_table(
probe, sensor_type[0:3], sensor)
store_data(var_name + suffix,
data={
'x': times,
'y': data,
'v': energy_map
})
def test_tnormalize(self):
""" tests for normalizing tplot variables"""
store_data('test_tnormalize',
data={
'x': [1, 2, 3, 4, 5],
'y': [[3, 2, 1], [1, 2, 3], [10, 5, 1], [8, 10, 14],
[70, 20, 10]]
})
norm = tnormalize('test_tnormalize')
normalized_data = get_data(norm)
self.assertTrue(
np.round(normalized_data.y[0, :], 4).tolist() ==
[0.8018, 0.5345, 0.2673])
self.assertTrue(
np.round(normalized_data.y[1, :], 4).tolist() ==
[0.2673, 0.5345, 0.8018])
self.assertTrue(
np.round(normalized_data.y[2, :], 4).tolist() ==
[0.8909, 0.4454, 0.0891])
self.assertTrue(
np.round(normalized_data.y[3, :], 4).tolist() ==
[0.4216, 0.527, 0.7379])
self.assertTrue(
np.round(normalized_data.y[4, :], 4).tolist() ==
[0.9526, 0.2722, 0.1361])
def tdeflag(names, method=None, new_names=None, suffix=None, overwrite=None):
"""
Remove NaNs from tplot variables.
Parameters
----------
names: str/list of str
List of pytplot names.
method: str, optional
Method to apply. Default is 'remove_nan.
Other options 'repeat' (repeat last good value).
new_names: str/list of str, optional
List of new_names for pytplot variables.
If '', then pytplot variables are replaced.
If not given, then a suffix is applied.
suffix: str, optional
A suffix to apply. Default is '-clip'.
overwrite: bool, optional
Replace the existing tplot name.
Returns
-------
None.
"""
old_names = pyspedas.tnames(names)
if len(old_names) < 1:
print('tdeflag error: No pytplot names were provided.')
return
if suffix is None:
suffix = '-deflag'
if overwrite is not None:
n_names = old_names
elif new_names is None:
n_names = [s + suffix for s in old_names]
else:
n_names = new_names
if isinstance(n_names, str):
n_names = [n_names]
if len(n_names) != len(old_names):
n_names = [s + suffix for s in old_names]
for i in range(len(old_names)):
alldata = pytplot.get_data(old_names[i])
time = alldata[0]
data = alldata[1]
new_time = []
new_data = []
for j in range(len(time)):
if not numpy.isnan(data[j]):
new_time.append(time[j])
new_data.append(data[j])
pytplot.store_data(n_names[i], data={'x': new_time, 'y': new_data})
print('tdeflag was applied to: ' + n_names[i])
def tinterpol(names, interp_names=None, method=None, new_names=None,
suffix=None):
old_names = pyspedas.tnames(names)
if len(old_names) < 1:
print('tinterpol error: No pytplot names were provided.')
return
if suffix is None:
suffix = '-itrp'
if method is None:
method = 'linear'
if new_names is None:
n_names = [s + suffix for s in old_names]
elif new_names == '':
n_names = old_names
else:
n_names = new_names
if len(n_names) != len(old_names):
n_names = [s + suffix for s in old_names]
for i in range(len(old_names)):
time, data = pytplot.get_data(old_names[i])
new_time, data1 = pytplot.get_data(interp_names[i])
data = numpy.asarray(data).squeeze()
f2 = interp1d(time, data, kind=method)
new_data = f2(new_time)
pytplot.store_data(n_names[i], data={'x': new_time, 'y': new_data})
print('tinterpol (' + method + ') was applied to: ' + n_names[i])
def avg_res_data(tvar1,res,new_tvar=None):
"""
Averages the variable over a specified period of time.
Parameters:
tvar1 : str
Name of tplot variable.
res : int/float
The new data resolution
new_tvar : str
Name of new tvar for averaged data. If not set, then the data in tvar1 is replaced.
Returns:
None
Examples:
>>> #Average the data over every two seconds
>>> pytplot.store_data('d', data={'x':[2,5,8,11,14,17,21], 'y':[[1,1,50],[2,2,3],[100,4,47],[4,90,5],[5,5,99],[6,6,25],[7,7,-5]]})
>>> pytplot.avg_res_data('d',2,'d2res')
>>> print(pytplot.data_quants['d'].values)
"""
tvar = pytplot.data_quants[tvar1].coarsen(time=res, boundary='trim').mean()
tvar.name = pytplot.data_quants[tvar1].name
if new_tvar is None:
pytplot.data_quants[tvar1] = tvar
else:
if 'spec_bins' in pytplot.data_quants[tvar1].coords:
pytplot.store_data(new_tvar, data={'x': tvar.coords['time'].values, 'y': tvar.values,
'v': tvar.coords['spec_bins'].values})
else:
pytplot.store_data(new_tvar, data={'x': tvar.coords['time'].values, 'y': tvar.values})
def spec_mult(tvar,new_tvar=None):
"""
Multiplies the data by the stored spectrogram bins and created a new tplot variable
.. note::
This analysis routine assumes the data is no more than 2 dimensions. If there are more, they may become flattened!
Parameters:
tvar : str
Name of tplot variable
times : int/list
Desired times for interpolation.
new_tvar : str
Name of new tvar in which to store interpolated data. If none is specified, a name will be created.
Returns:
None
Examples:
>>> pytplot.store_data('h', data={'x':[0,4,8,12,16,19,21], 'y':[[8,1,1],[100,2,3],[4,2,47],[4,39,5],[5,5,99],[6,6,25],[7,-2,-5]],'v':[[1,1,50],[2,2,3],[100,4,47],[4,90,5],[5,5,99],[6,6,25],[7,7,-5]]})
>>> pytplot.spec_mult('h','h_specmult')
>>> print(pytplot.data_quants['h_specmult'].data)
"""
if new_tvar is None:
new_tvar = tvar+'_specmult'
if 'spec_bins' not in pytplot.data_quants[tvar].coords:
print("Specified variable must have spec bins stored. Returning...")
return
d, s = pytplot.tplot_utilities.convert_tplotxarray_to_pandas_dataframe(tvar)
dataframe = d.values
specframe = s.values
new_df = pd.DataFrame(dataframe*specframe, columns=d.columns, index=d.index)
pytplot.store_data(new_tvar,data={'x': new_df.index,'y': new_df.values})
return
def derive(tvar, new_tvar=None):
"""
Takes the derivative of the tplot variable.
Parameters:
tvar : str
Name of tplot variable.
new_tvar : str
Name of new tplot variable. If not set, then the data in tvar is replaced.
Returns:
None
Examples:
>>> pytplot.store_data('b', data={'x':[2,5,8,11,14,17,20], 'y':[[1,1,1,1,1,1],[2,2,5,4,1,1],[100,100,3,50,1,1],[4,4,8,58,1,1],[5,5,9,21,1,1],[6,6,2,2,1,1],[7,7,1,6,1,1]]})
>>> pytplot.derive('b','dbdt')
>>> print(pytplot.data_quants['dbdt'].values)
"""
a = pytplot.data_quants[tvar].differentiate('time')
if new_tvar is None:
a.name = tvar
a.attrs = copy.deepcopy(pytplot.data_quants[tvar].attrs)
pytplot.data_quants[tvar] = a
else:
data = {'x': a.coords['time'], 'y': a.values}
for coord in a.coords:
if coord != 'time' and coord != 'spec_bins':
data[coord] = a.coords[coord].values
pytplot.store_data(new_tvar, data=data)
pytplot.data_quants[new_tvar].attrs = copy.deepcopy(
pytplot.data_quants[tvar].attrs)
return
def spd_pgs_make_tplot(name,
x=None,
y=None,
z=None,
units='',
ylog=False,
zlog=True,
colorbar='jet',
ytitle=None):
"""
Create tplot variable with standard spectrogram settings
Input:
name: str
Name of the new tplot variable to create
Parameters:
x: numpy.ndarray
X-axis values (time)
y: numpy.ndarray
Y-axis values
z: numpy.ndarray
Z-axis values (data)
units: str
Units string to store in the metadata
ylog: bool
Set the y-axis to log scale (default: False)
zlog: bool
Set the z-axis to log scale (default: True)
colorbar: str
PyTplot 'Colormap' option (default: 'jet')
Returns:
String containing new variable name
"""
if not isinstance(x, np.ndarray) or not isinstance(
y, np.ndarray) or not isinstance(z, np.ndarray):
print('Error, must specify x, y and z parameters')
return
if ytitle is None:
ytitle = name
store_data(name, data={'x': x, 'y': z, 'v': y})
options(name, 'ylog', ylog)
options(name, 'zlog', zlog)
options(name, 'Spec', True)
options(name, 'ytitle', ytitle)
options(name, 'ztitle', spd_units_string(units, units_only=True))
options(name, 'Colormap', colorbar)
return name
def full_flatten(tvar1, new_tvar):
df = pytplot.data_quants[tvar1].data
df_index = list(df.columns)
#divide by column average
for i in df_index:
df[i] = df[i] / df[i].mean()
pytplot.store_data(new_tvar, data={'x': df.index, 'y': df})
return new_tvar
def add_data_across(tvar1, new_tvar):
#separate and add data
time = pytplot.data_quants[tvar1].data.index
data1 = pytplot.data_quants[tvar1].data
data = data1.sum(axis=1)
#store added data
pytplot.store_data(new_tvar, data={'x': time, 'y': data})
return new_tvar
def spec_mult(tvar1, new_tvar):
dataframe = pytplot.data_quants[tvar1].data
specframe = pytplot.data_quants[tvar1].spec_bins
new_df = pd.DataFrame(dataframe.values * specframe.values,
columns=dataframe.columns,
index=dataframe.index)
pytplot.store_data(new_tvar, data={'x': new_df.index, 'y': new_df.values})
return new_tvar
def join_vec(tvars, new_tvar=None):
"""
Joins 1D tplot variables into one tplot variable.
.. note::
This analysis routine assumes the data is no more than 2 dimensions. If there are more, they may become flattened!
Parameters:
tvars : list of str
Name of tplot variables to join together
new_tvar : str, optional
The name of the new tplot variable. If not specified, a name will be assigned.
Returns:
None
Examples:
>>> pytplot.store_data('d', data={'x':[2,5,8,11,14,17,21], 'y':[[1,1,50],[2,2,3],[100,4,47],[4,90,5],[5,5,99],[6,6,25],[7,7,-5]]})
>>> pytplot.store_data('e', data={'x':[2,5,8,11,14,17,21], 'y':[[np.nan,1,1],[np.nan,2,3],[4,np.nan,47],[4,np.nan,5],[5,5,99],[6,6,25],[7,np.nan,-5]]})
>>> pytplot.store_data('g', data={'x':[0,4,8,12,16,19,21], 'y':[[8,1,1],[100,2,3],[4,2,47],[4,39,5],[5,5,99],[6,6,25],[7,-2,-5]]})
>>> pytplot.join_vec(['d','e','g'],'deg')
>>> print(pytplot.data_quants['deg'].values)
"""
if not isinstance(tvars, list):
tvars = [tvars]
if new_tvar is None:
new_tvar = '-'.join(tvars) + '_joined'
for i, val in enumerate(tvars):
if i == 0:
if 'spec_bins' in pytplot.data_quants[tvars[i]].coords:
df, s = pytplot.tplot_utilities.convert_tplotxarray_to_pandas_dataframe(
tvars[i])
else:
df = pytplot.tplot_utilities.convert_tplotxarray_to_pandas_dataframe(
tvars[i])
s = None
else:
if 'spec_bins' in pytplot.data_quants[tvars[i]].coords:
d, _ = pytplot.tplot_utilities.convert_tplotxarray_to_pandas_dataframe(
tvars[i])
else:
d = pytplot.tplot_utilities.convert_tplotxarray_to_pandas_dataframe(
tvars[i])
df = pd.concat([df, d], axis=1)
if s is None:
pytplot.store_data(new_tvar, data={'x': df.index, 'y': df.values})
else:
pytplot.store_data(new_tvar,
data={
'x': df.index,
'y': df.values,
'v': s.values
})
return
def mms_fgm_remove_flags(probe, data_rate, level, instrument, suffix=''):
"""
This function removes data flagged by the FGM 'flag' variable (flags > 0),
in order to only show science quality data by default.
Parameters:
probe : str or list of str
probe or list of probes, valid values for MMS probes are ['1','2','3','4'].
data_rate : str or list of str
instrument data rates for FGM include 'brst' 'fast' 'slow' 'srvy'. The
default is 'srvy'.
level : str
indicates level of data processing. the default if no level is specified is 'l2'
instrument : str
instrument; probably 'fgm'
suffix: str
The tplot variable names will be given this suffix. By default,
no suffix is added.
"""
if not isinstance(probe, list): probe = [probe]
if not isinstance(data_rate, list): data_rate = [data_rate]
if not isinstance(level, list): level = [level]
tplot_vars = set(tnames())
for this_probe in probe:
for this_dr in data_rate:
for this_lvl in level:
if level == 'ql':
flag_var = 'mms' + str(
this_probe
) + '_' + instrument + '_' + this_dr + '_' + this_lvl + '_flag' + suffix
else:
flag_var = 'mms' + str(
this_probe
) + '_' + instrument + '_flag_' + this_dr + '_' + this_lvl + suffix
flagged = get_data(flag_var)
if flagged == None:
continue
times, flags = flagged
flagged_data = np.where(flags != 0.0)[0]
for var_specifier in [
'_b_gse_', '_b_gsm_', '_b_dmpa_', '_b_bcs_'
]:
var_name = 'mms' + str(
this_probe
) + '_' + instrument + var_specifier + this_dr + '_' + this_lvl + suffix
if var_name in tplot_vars:
times, var_data = get_data(var_name)
var_data[flagged_data] = np.nan
store_data(var_name, data={'x': times, 'y': var_data})
def join_vec(tvars, newtvar):
df = pytplot.data_quants[tvars[0]].data
for i, val in enumerate(tvars):
if i == 0:
pass
else:
df = pd.concat([df, pytplot.data_quants[val].data], axis=1)
pytplot.store_data(newtvar, data={'x': df.index, 'y': df})
return newtvar
def test_tinterpol(self):
"""Test tinterpol."""
tinterpol('aaabbbccc', 'test') # Test non-existent name
tn = [1., 1.5, 4.6, 5.8, 6.]
dn = [10., 15., 46., 58., 60.]
store_data('test1', data={'x': tn, 'y': dn})
tinterpol('test1', 'test')
d = get_data('test1-itrp')
self.assertTrue(d[1][1] == 20.)
def subtract(tvar1, tvar2, new_tvar=None):
"""
Subtracts two tplot variables. Will interpolate if the two are not on the same time cadence.
Parameters:
tvar1 : str
Name of first tplot variable.
tvar2 : int/float
Name of second tplot variable
new_tvar : str
Name of new tvar for added data. If not set, then the data in tvar1 is replaced.
Returns:
None
Examples:
>>> pytplot.store_data('a', data={'x':[0,4,8,12,16], 'y':[1,2,3,4,5]})
>>> pytplot.store_data('c', data={'x':[0,4,8,12,16,19,21], 'y':[1,4,1,7,1,9,1]})
>>> pytplot.subtract('a','c','a-c')
"""
#interpolate tvars
tv2 = pytplot.tplot_math.tinterp(tvar1, tvar2)
#separate and subtract data
data1 = pytplot.data_quants[tvar1].values
data2 = pytplot.data_quants[tv2].values
data = data1 - data2
#store subtracted data
if new_tvar is None:
pytplot.data_quants[tvar1].values = data
return tvar1
if 'spec_bins' in pytplot.data_quants[tvar1].coords:
pytplot.store_data(
new_tvar,
data={
'x': pytplot.data_quants[tvar1].coords['time'].values,
'y': data,
'v': pytplot.data_quants[tvar1].coords['spec_bins'].values
})
pytplot.data_quants[new_tvar].attrs = copy.deepcopy(
pytplot.data_quants[tvar1].attrs)
else:
pytplot.store_data(
new_tvar,
data={
'x': pytplot.data_quants[tvar1].coords['time'].values,
'y': data
})
pytplot.data_quants[new_tvar].attrs = copy.deepcopy(
pytplot.data_quants[tvar1].attrs)
return new_tvar
def tinterpol(names, interp_to, method=None, newname=None, suffix=None):
if not isinstance(names, list):
names = [names]
if not isinstance(newname, list):
newname = [newname]
old_names = tnames(names)
if len(old_names) < 1:
print('tinterpol error: No pytplot names were provided.')
return
if suffix == None:
suffix = '-itrp'
if method == None:
method = 'linear'
if (newname == None) or (len(newname) == 1 and newname[0] == None):
n_names = [s + suffix for s in old_names]
elif newname == '':
n_names = old_names
else:
n_names = newname
interp_to_data = get_data(interp_to)
if interp_to_data == None:
print('Error, tplot variable: ' + interp_to + ' not found.')
return
interp_to_times = interp_to_data[0]
for name_idx, name in enumerate(old_names):
xdata = get_data(name, xarray=True)
xdata_interpolated = xdata.interp({'time': interp_to_times},
method=method)
if 'spec_bins' in xdata.coords:
store_data(n_names[name_idx],
data={
'x': interp_to_times,
'y': xdata_interpolated.values,
'v': xdata_interpolated.coords['spec_bins'].values
})
else:
store_data(n_names[name_idx],
data={
'x': interp_to_times,
'y': xdata_interpolated.values
})
print('tinterpol (' + method + ') was applied to: ' +
n_names[name_idx])
def degap(tvar,dt,margin,func='nan',new_tvar = None):
'''
Fills gaps in the data either with NaNs or the last number.
Parameters:
tvar : str
Name of tplot variable to modify
dt : int/float
Step size of the data in seconds
margin : int/float, optional
The maximum deviation from the step size allowed before degapping occurs. In other words, if you'd like to fill in data every 4 seconds
but occasionally the data is 4.1 seconds apart, set the margin to .1 so that a data point is not inserted there.
func : str, optional
Either 'nan' or 'ffill', which overrides normal interpolation with NaN
substitution or forward-filled values.
new_tvar : str, optional
The new tplot variable name to store the data into. If None, then the data is overwritten.
Returns:
None
Examples:
>>> # TODO
'''
gap_size = np.diff(pytplot.data_quants[tvar].coords['time'])
gap_index_locations = np.where(gap_size > dt+margin)
new_tvar_index = pytplot.data_quants[tvar].coords['time']
values_to_add = np.array([])
for i in gap_index_locations[0]:
values_to_add = np.append(values_to_add, np.arange(new_tvar_index[i], new_tvar_index[i+1], dt))
new_index = np.sort(np.unique(np.concatenate((values_to_add, new_tvar_index))))
if func == 'nan':
method = None
if func == 'ffill':
method = 'ffill'
a = pytplot.data_quants[tvar].reindex({'time': new_index}, method=method)
if new_tvar is None:
a.name = tvar
a.attrs = copy.deepcopy(pytplot.data_quants[tvar].attrs)
pytplot.data_quants[tvar] = copy.deepcopy(a)
else:
if 'spec_bins' in a.coords:
pytplot.store_data(new_tvar, data={'x': a.coords['time'], 'y': a.values, 'v': a.coords['spec_bins']})
pytplot.data_quants[new_tvar].attrs = copy.deepcopy(pytplot.data_quants[tvar].attrs)
else:
pytplot.store_data(new_tvar, data={'x': a.coords['time'], 'y': a.values})
pytplot.data_quants[new_tvar].attrs = copy.deepcopy(pytplot.data_quants[tvar].attrs)
return
def tnormalize(variable, newname=None, return_data=False):
"""
Normalize all the vectors stored in a tplot variable
Input
----------
variable: str or np.ndarray
tplot variable (or numpy array) containing the vectors to be normalized
Parameters
----------
newname: str
name of the output variable; default: variable_normalized
return_data: bool
return the normalized vectors instead of creating a tplot variable
Returns
----------
name of the tplot variable created or normalized vectors if return_data
is set
"""
metadata_in = {}
if isinstance(variable, str):
data_in = get_data(variable)
metadata_in = get_data(variable, metadata=True)
data = data_in[1]
times = data_in[0]
else:
data = np.atleast_2d(variable)
times = np.zeros(data.shape[0])
n = np.sqrt(np.nansum(data**2, axis=1))
# to do element-wise division, the magnitude needs to be repeated for each component
norm_reshaped = np.reshape(n, [len(times), 1])
norm_mag = np.repeat(norm_reshaped, len(data[0, :]), axis=1)
data_norm = data / norm_mag
if return_data:
return data_norm
else:
if newname is None:
newname = variable + '_normalized'
store_data(newname,
data={
'x': times,
'y': data_norm
},
attr_dict=metadata_in)
return newname
def add_data(tvar1, tvar2, new_tvar, interp='linear'):
#interpolate tvars
tv1, tv2 = fn_interp(tvar1, tvar2, interp=interp)
#separate and add data
time = pytplot.data_quants[tv1].data.index
data1 = pytplot.data_quants[tv1].data
data2 = pytplot.data_quants[tv2].data
data = data1 + data2
#store added data
pytplot.store_data(new_tvar, data={'x': time, 'y': data})
return new_tvar
