Source code for pyspedas.tplot_tools.tplot_math.tdpwrspc

"""
Compute power spectra for a tplot variable.

Calls dpwrspc for the actual computation.

Notes
-----
Similar to tdpwrspc.pro in IDL SPEDAS.

"""
import logging
import numpy as np
from pyspedas.tplot_tools import dpwrspc
from pyspedas.tplot_tools import split_vec
from pyspedas.tplot_tools import get_data, store_data, options, time_double, set_units


[docs] def tdpwrspc( varname, newname=None, trange=["0.0", "0.0"], nboxpoints=None, nshiftpoints=None, polar=False, bin=3, nohanning=False, noline=False, notperhz=False, notmvariance=False, ): """ Compute power spectra for a tplot variable. Parameters ---------- varname: str Name of tplot variable. newname: str, optional Name of new tplot variable to save data to. Default: None. If newname is not set '_dpwrspc' will be appended to the varname trange: list of float, optional Start and end times for the calculation. nboxpoints: int, optional The number of points to use for the hanning window. Default: 256 nshiftpoints: int, optional The number of points to shift for each spectrum. Default: 128 polar: bool, optional If True, the input data is in polar coordinates. Default: False bin: int, optional Size for binning of the data along the frequency domain. Default: 3 nohanning: bool, optional If True, no hanning window is applied to the input. Default: False noline: bool, optional If True, no straight line is subtracted. Default: False notperhz: bool, optional If True, the output units are the square of the input units. Default: False notmvariance: bool, optional If True, replace output spectrum for any windows that have variable. cadence with NaNs. Default: False Returns ------- str Name of new tplot variable. Example ------- >>> # Compute the power spectrum of a given time series >>> import pyspedas >>> import numpy as np >>> pyspedas.store_data('a', data={'x': range(3000), 'y': np.random.random(3000)}) >>> pyspedas.tdpwrspc('a') """ if newname is None: newname = varname + "_dpwrspc" data_tuple = get_data(varname) metadata = get_data(varname, metadata=True) # Find units inputunits = "#" has_data_att = False has_units = False newunits = "" if "data_att" in metadata: mm = metadata["data_att"] has_data_att = True if "units" in mm: inputunits = mm["units"] has_units = True if inputunits != "#": if notperhz: newunits = "(" + inputunits + ")^2" else: newunits = "(" + inputunits + ")^2/Hz" if data_tuple is not None: if data_tuple[1][0].shape != (): split_vars = split_vec(varname, polar=polar) out_vars = [] for var in split_vars: out_vars.append( tdpwrspc( var, newname=var + "_dpwrspc", trange=trange, nboxpoints=nboxpoints, nshiftpoints=nshiftpoints, polar=polar, bin=bin, nohanning=nohanning, noline=noline, notperhz=notperhz, notmvariance=notmvariance, ) ) return out_vars else: t = data_tuple[0] y = data_tuple[1] if trange is not None and len(trange) == 2 and trange[1] > trange[0]: # Only use points inside the time range tr = time_double(trange) ok = np.argwhere((t >= tr[0]) & (t < tr[1])) if len(ok) == 0: logging.error("No data in time range") logging.error(f"{tr}") return None t = t[ok] y = y[ok] # Filter out NaNs ok = np.isfinite(y) if len(ok) == 0: logging.error("No finite data in time range") return None t = t[ok] y = y[ok] t00 = data_tuple[0][0] t = t - t00 # Only do this if there are enough data points, default nboxpoints to # 64 and nshiftpoints to 32, and use larger values when there are more # points if nboxpoints is None: nbp = np.max([2 ** (np.floor(np.log(len(ok)) / np.log(2)) - 5), 8]) else: nbp = nboxpoints if nshiftpoints is None: nsp = nbp / 2.0 else: nsp = nshiftpoints if len(ok) <= nbp: logging.error("Not enough data in time range") return None pwrspc = dpwrspc( t, y, nboxpoints=nbp, nshiftpoints=nsp, bin=bin, nohanning=nohanning, noline=noline, notperhz=notperhz, notmvariance=notmvariance, ) if pwrspc is not None: store_data( newname, data={"x": pwrspc[0] + t00, "y": pwrspc[2], "v": pwrspc[1]} ) options(newname, "ysubtitle", "[Hz]") options(newname, "spec", True) options(newname, "ylog", True) options(newname, "zlog", True) options(newname, "Colormap", "spedas") if has_units: set_units(newname, newunits) if notperhz: options(newname, "ztitle", "(" + inputunits + ")^2") else: options(newname, "ztitle", "(" + inputunits + ")^2/Hz") return newname return None