"""
Splits up 2D data into many 1D tplot variables.
Notes
-----
Similar to split_vec.pro in IDL SPEDAS.
"""
import logging
import pytplot
import numpy as np
import logging
[docs]
def split_vec(
tvar,
polar=False,
newname=None,
new_name=None,
columns='all',
suffix=None
):
"""
Splits up 2D data into many 1D tplot variables. Takes a stored tplot vector like Vp
and stores tplot variables Vp_x, Vp_y, Vp_z
Parameters
----------
tvar : str
Name of tplot variable to split up
polar : bool, optional
If True, the input data is in polar coordinates.
Suffix will be set to ['_mag', '_th', '_phi'].
Default: False
new_name : int/list, optional (Deprecated)
The names of the new tplot variables.
This must be the same length as the number of variables created.
newname : int/list, optional
The names of the new tplot variables.
This must be the same length as the number of variables created.
Default: None
columns : list of ints, optional
The specific column numbers to grab from the data.
Default: 'all' (splits all columns)
suffix: str
Suffix str to be added to end of tplot variable name
Default: None
Returns
-------
list[str]
List of variables created
Examples
--------
>>> pyspedas.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]]})
>>> pyspedas.split_vec('b')
"""
# new_name is deprecated in favor of newname
if new_name is not None:
logging.info("split_vec: The new_name parameter is deprecated. Please use newname instead.")
newname = new_name
# Make sure the tvar is found
if tvar not in pytplot.data_quants:
logging.error(f"Error: {tvar} not found in memory.")
return None
# Give a default to the new name
if newname is None:
newname = tvar
# Gather data from the tvar
alldata = pytplot.get_data(tvar)
time = alldata[0]
data = alldata[1]
dim = data.shape
metadata = pytplot.get_data(tvar, metadata=True)
# If already size one, simply return
if len(dim) == 1:
return [tvar]
vec_length = dim[1]
# Determine what the suffix list will be
if suffix is not None:
if vec_length > len(suffix):
logging.error(f"split_vec error: number of columns ({vec_length}) is greater than the number of suffix entered")
else:
if vec_length == 3:
if polar:
suffix = ['_mag', '_th', '_phi']
else:
suffix = ["_x", "_y", "_z"]
else:
suffix = []
for i in range(vec_length):
suffix.append("_"+str(i))
created_variables = []
# grab column data
if columns == 'all':
columns = range(vec_length)
for i in columns:
# if not a list
if isinstance(i, list):
range_start = i[0]
range_end = i[1]
else:
range_start = i
range_end = i
split_col = list(range(range_start, range_end+1))
split_name = newname + suffix[i]
created_variables = created_variables + [split_name]
data_for_tplot = {'x': time, 'y': data[:, split_col].squeeze()}
if not pytplot.store_data(split_name, data=data_for_tplot, attr_dict=metadata):
raise Exception(f"Failed to store {split_name} in pytplot.")
return created_variables