diff --git a/pyspedas/analysis/tvectot.py b/pyspedas/analysis/tvectot.py new file mode 100644 index 00000000..548fdc43 --- /dev/null +++ b/pyspedas/analysis/tvectot.py @@ -0,0 +1,46 @@ +from xarray_einstats import linalg +from pytplot import split_vec, join_vec, get_data, store_data, options + +def _tvectot(tvar: str, new_name: str, join_component: bool): + data = get_data(tvar, xarray=True) + new_data = linalg.norm(data, dims="v_dim") + store_data(new_name, new_data, xarray=True) + + if join_component: + join_vec(split_vec(tvar)+[new_name], new_name) + options(new_name, 'legend_names', ['x', 'y', 'z', 'Magnitude']) + else: + options(new_name, 'legend_names', 'Magnitude') + return new_name + +def tvectot(tvars: str | list[str], newnames: str | list[str] = None, suffix="_mag", join_component=False) -> str | list[str]: + """ + Computes the magnitude of a vector time series. + + Parameters + ---------- + - tvars : Names of the tplot variables. + - new_names: Names for the resultant magnitude tplot variables. If not provided, it appends the suffix to `tvars`. + - suffix: The suffix to append to tensor_names to form new_names if new_names is not provided. + - join_component: If True, the magnitude tplot variable is joined with the component tplot variables. + + Returns + ------- + Names of the magnitude tplot variables. + """ + tvars_type = type(tvars) + if tvars_type == str: + tvars = [tvars] + if join_vec: + suffix = "_tot" + + if newnames is None: + newnames = [tvar + suffix for tvar in tvars] + + for tvar, newname in zip(tvars, newnames): + _tvectot(tvar, newname, join_component) + + if tvars_type == str: + return newnames[0] + else: + return newnames diff --git a/requirements.txt b/requirements.txt index 51aaa78e..cba726bd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ astropy geopack>=1.0.10 hapiclient>=0.2.2 viresclient +xarray-einstats # Stats, linear algebra and einops for xarray