import numpy as np
from copy import deepcopy
from pyspedas.tplot_tools import get_data, store_data, options, tplot_wildcard_expand
from typing import Union, List
import logging
def _tvectot(tvar: str, new_name: str, join_component: bool):
data = get_data(tvar)
md = get_data(tvar,metadata=True)
new_data = np.linalg.norm(data.y, axis=1)
if join_component:
# Caution: this can clobber existing variable names! Maybe better to manipulate the
# data directly instead of calling split_vec/join_vec
#join_vec(split_vec(tvar)+[new_name], newname=new_name)
joined_data = np.append(data.y, new_data[:,np.newaxis],1)
store_data(new_name, data={'x':data.times,'y':joined_data}, attr_dict=deepcopy(md))
options(new_name, 'legend_names', ['x', 'y', 'z', 'Magnitude'])
else:
store_data(new_name, data={'x': data.times, 'y': new_data}, attr_dict=md)
options(new_name, 'legend_names', 'Magnitude')
return new_name
[docs]
def tvectot(tvars: Union[str, List[str]],
newname=None,
suffix=None,
join_component=True) -> Union[str , List[str]]:
"""
Computes the magnitude of a vector time series.
Parameters
----------
tvars : str or list[str]
Names of the tplot variables.
newname: str or list[str]
Names for the resultant magnitude tplot variables. If not provided, it appends the suffix to `tvars`.
suffix: str
The suffix to append to tvars to form newname if newname is not provided.
join_component: bool
If True, the magnitude tplot variable is joined with the component tplot variables.
Returns
-------
str or list[str]
Names of the magnitude tplot variables.
"""
tvars = tplot_wildcard_expand(tvars)
if len(tvars) == 0:
return
tvars_type = type(tvars)
if tvars_type == str:
tvars = [tvars]
if suffix is None:
if join_component:
suffix = "_tot"
else:
suffix='_mag'
if newname is None:
newname = [tvar + suffix for tvar in tvars]
# If newname is a scalar string, zip will iterate over characters
if type(newname) is not list:
newname = [newname]
for tvar, newname_single in zip(tvars, newname):
_tvectot(tvar, newname_single, join_component)
if tvars_type == str:
return newname[0]
else:
return newname