Added dpwrspc. tdpwrspc

Previously, those functions were in pyspedas

meta.yaml

@@ -1,5 +1,5 @@
 {% set name = "pytplot" %}
-{% set version = "2.2.1" %}
+{% set version = "2.2.2" %}
 
 package:
   name: "{{ name|lower }}"

pytplot/__init__.py

@@ -165,7 +165,7 @@ def newlayout(self, layout):
 #from pytplot.tplot_math import *
 from .tplot_math import add_across, add, avg_res_data, clip, crop, deflag, degap, derive, divide, flatten
 from .tplot_math import interp_nan, tinterp, join_vec, multiply, resample, spec_mult, split_vec, subtract, pwr_spec
-from .tplot_math import clean_spikes, tsmooth, smooth, tdotp, tcrossp, pwrspc, tpwrspc
+from .tplot_math import clean_spikes, tsmooth, smooth, tdotp, tcrossp, pwrspc, tpwrspc, dpwrspc, tdpwrspc
 #from .tplot_math import tdeflag
 from .tplot_math import tnormalize, subtract_median, subtract_average, time_clip, tkm2re, makegap
 

pytplot/tplot_math/__init__.py

@@ -30,3 +30,5 @@
 from .tdeflag import tdeflag
 from .pwrspc import pwrspc
 from .tpwrspc import tpwrspc
+from .dpwrspc import dpwrspc
+from .tdpwrspc import tdpwrspc

pytplot/tplot_math/dpwrspc.py

@@ -0,0 +1,230 @@
+"""
+Compute power spectra for data.
+
+Notes
+-----
+Similar to dpwrspc.pro in IDL SPEDAS.
+
+"""
+import logging
+import numpy as np
+
+
+def dpwrspc(
+    time,
+    quantity,
+    nboxpoints=256,
+    nshiftpoints=128,
+    bin=3,
+    tbegin=-1.0,
+    tend=-1.0,
+    noline=False,
+    nohanning=False,
+    notperhz=False,
+    notmvariance=False,
+    tm_sensitivity=None,
+):
+    """
+    Compute power spectra.
+
+    Parameters
+    ----------
+    time: list of float
+        Time array.
+    quantity: list of float
+        Data array.
+    nboxpoints: int, optional
+        The number of points to use for the hanning window.
+        The default is 256.
+    nshiftpoints: int, optional
+        The number of points to shift for each spectrum.
+        The default is 128.
+    bin: int, optional
+        Size for binning of the data along the frequency domain.
+        The default is 3.
+    tbegin: float, optional
+        Start time for the calculation.
+        If -1.0, the start time is the first time in the time array.
+    tend: float, optional
+        End time for the calculation.
+        If -1.0, the end time is the last time in the time array.
+    nohanning: bool, optional
+        If True, no hanning window is applied to the input.
+        The default is False.
+    noline: bool, optional
+        If True, no straight line is subtracted.
+        The default is False.
+    notperhz: bool, optional
+        If True, the output units are the square of the input units.
+        The default is False.
+    notmvariance: bool, optional
+        If True, replace output spectrum for any windows that have variable
+        cadence with NaNs.
+        The default is False.
+    tm_sensitivity: float, optional
+        If noTmVariance is set, this number controls how much of a dt anomaly
+        is accepted.
+        The default is None.
+
+    Returns
+    -------
+    tdps: array of float
+        The time array for the dynamic power spectrum, the center time of the
+        interval used for the spectrum.
+    fdps: array of float
+        The frequency array (units =1/time units).
+    dps: array of float
+        The power spectrum, (units of quantity)^2/frequency_units.
+
+    """
+    tdps, fdps, dps = np.array(-1.0), np.array(-1.0), np.array(-1.0)
+    tdps0, fdps0, dps0 = np.array(-1.0), np.array(-1.0), np.array(-1.0)
+
+    if nohanning is False:
+        window = np.array(np.hanning(nboxpoints), dtype=np.float64)
+    else:
+        window = np.ones(nboxpoints)
+
+    time = np.array(time, dtype=np.float64)
+    quantity = np.array(quantity, dtype=np.float64)
+
+    if tbegin == -1.0:
+        tbegin = time[0]
+    if tend == -1.0:
+        tend = time[-1]
+
+    # Find the time array that correspond to the start and end times
+    tbegin_idx = np.where(time >= tbegin)[0][0]
+    tend_idx = np.where(time <= tend)[0][-1]
+    if tend_idx - tbegin_idx < nboxpoints:
+        logging.error("Not enough points for a calculation")
+        return tdps0, fdps0, dps0
+    time = time[tbegin_idx : tend_idx + 1]
+    quantity = quantity[tbegin_idx : tend_idx + 1]
+
+    # remove NaNs from the data
+    where_finite = np.where(np.isnan(quantity) == False)
+    quantity2process = quantity[where_finite[0]]
+    times2process = time[where_finite[0]]
+    nboxpnts = int(nboxpoints)
+    nshiftpnts = nshiftpoints
+
+    totalpoints = len(times2process)
+    nspectra = int((totalpoints - nboxpnts / 2.0) / nshiftpnts)
+
+    # test nspectra, if the value of nshiftpnts is much smaller than
+    # nboxpnts/2 strange things happen
+
+    nbegin = np.array([nshiftpnts * i for i in range(nspectra)], dtype=np.int64)
+    nend = nbegin + nboxpnts
+
+    okspec = np.where(nend <= totalpoints - 1)
+
+    if len(okspec[0]) <= 0:
+        logging.error("Not enough points for a calculation")
+        return tdps0, fdps0, dps0
+
+    tdps = np.zeros(nspectra)
+    nfreqs = int(int(nboxpnts / 2) / bin)
+
+    if nfreqs <= 1:
+        logging.error("Not enough frequencies for a calculation")
+        return tdps0, fdps0, dps0
+
+    dps = np.zeros([nspectra, nfreqs])
+    fdps = np.zeros([nspectra, nfreqs])
+
+    # Main calculation loop
+    for nthspectrum in range(0, nspectra):
+        nbegin = int(nthspectrum * nshiftpnts)
+        nend = nbegin + nboxpnts
+
+        if nend <= totalpoints:
+            t = times2process[nbegin:nend]
+            t0 = t[0]
+            t = t - t0
+            x = quantity2process[nbegin:nend]
+
+            # Use center time
+            tdps[nthspectrum] = (times2process[nbegin] + times2process[nend - 1]) / 2.0
+
+            if noline is False:
+                coef = np.polyfit(t, x, 1)
+                poly1d_fn = np.poly1d(coef)
+                line = poly1d_fn(t)
+                x = x - line
+
+            if nohanning is False:
+                x = x * window
+
+            bign = nboxpnts
+
+            if bign % 2 != 0:
+                logging.warning(
+                    "dpwrspc: needs an even number of data points, dropping last point..."
+                )
+                t = t[0 : bign - 1]
+                x = x[0 : bign - 1]
+                bign = bign - 1
+
+            n_tm = len(t)
+
+            # time variance can break power spectrum
+            # this keyword skips over those gaps
+            if notmvariance and n_tm > 1:
+                if tm_sensitivity is not None:
+                    tmsn = tm_sensitivity
+                else:
+                    tmsn = 100.0
+
+                tdiff = t[1:n_tm] - t[0 : n_tm - 1]
+                med_diff = np.median(tdiff)
+
+                idx = np.where(np.abs(tdiff / med_diff - 1) > 1.0 / tmsn)
+
+                if len(idx[0]) > 0:
+                    dps[nthspectrum, :] = float("nan")
+                    fdps[nthspectrum, :] = float("nan")
+
+                    continue
+
+            # following Numerical recipes in Fortran, p. 421, sort of...
+            # (actually following the IDL implementation)
+            k = np.array(range(int(bign / 2) + 1))
+            tres = np.median(t[1 : len(t)] - t[0 : len(t) - 1])
+            fk = k / (bign * tres)
+
+            xs2 = np.abs(np.fft.fft(x)) ** 2
+
+            pwr = np.zeros(int(bign / 2 + 1))
+            pwr[0] = xs2[0] / bign**2
+            ptmp = 1 + np.array(range(int(bign / 2 - 1)))
+            pwr[1 : int(bign / 2)] = (
+                xs2[1 : int(bign / 2)] + xs2[bign - ptmp]
+            ) / bign**2
+            pwr[int(bign / 2)] = xs2[int(bign / 2)] / bign**2
+
+            if nohanning is False:
+                wss = float(bign) * float(np.sum(window**2))
+                pwr = bign**2 * pwr / wss
+
+            dfreq = bin * (fk[1] - fk[0])
+
+            npwr = len(pwr) - 1
+            nfinal = int(npwr / bin)
+            iarray = np.array(range(nfinal))
+            power = np.zeros(nfinal)
+
+            # Note: zeroth point includes zero freq. power.
+            freqcenter = (fk[iarray * bin + 1] + fk[iarray * bin + bin]) / 2.0
+
+            for i in range(bin):
+                power = power + pwr[iarray * bin + i + 1]
+
+            if notperhz is False:
+                power = power / dfreq
+
+            dps[nthspectrum, :] = power
+            fdps[nthspectrum, :] = freqcenter
+
+    return tdps, fdps, dps