From 384cfbd9a81569e941f3360fcebc067cf70ca047 Mon Sep 17 00:00:00 2001
From: Nick H <34072991+nickssl@users.noreply.github.com>
Date: Thu, 14 Dec 2023 15:37:11 -0800
Subject: [PATCH] Added pwrspc

---
 pytplot/__init__.py            |  2 +-
 pytplot/tplot_math/__init__.py |  6 +--
 pytplot/tplot_math/pwrspc.py   | 97 ++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 6 deletions(-)
 create mode 100644 pytplot/tplot_math/pwrspc.py

diff --git a/pytplot/__init__.py b/pytplot/__init__.py
index a1321c8a..93f1960e 100644
--- a/pytplot/__init__.py
+++ b/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
+from .tplot_math import clean_spikes, tsmooth, smooth, tdotp, tcrossp, pwrspc
 #from .tplot_math import tdeflag
 from .tplot_math import tnormalize, subtract_median, subtract_average, time_clip, tkm2re, makegap
 
diff --git a/pytplot/tplot_math/__init__.py b/pytplot/tplot_math/__init__.py
index a71b609a..cd096b5c 100644
--- a/pytplot/tplot_math/__init__.py
+++ b/pytplot/tplot_math/__init__.py
@@ -1,8 +1,3 @@
-# Copyright 2018 Regents of the University of Colorado. All Rights Reserved.
-# Released under the MIT license.
-# This software was developed at the University of Colorado's Laboratory for Atmospheric and Space Physics.
-# Verify current version before use at: https://github.com/MAVENSDC/Pytplot
-
 from .add_across import add_across
 from .add import add
 from .avg_res_data import avg_res_data
@@ -33,3 +28,4 @@
 from .tkm2re import tkm2re
 from .makegap import makegap
 from .tdeflag import tdeflag
+from .pwrspc import pwrspc
diff --git a/pytplot/tplot_math/pwrspc.py b/pytplot/tplot_math/pwrspc.py
new file mode 100644
index 00000000..85f647ad
--- /dev/null
+++ b/pytplot/tplot_math/pwrspc.py
@@ -0,0 +1,97 @@
+import logging
+import numpy as np
+from scipy.stats import linregress
+
+
+def pwrspc(time, quantity, noline=False, nohanning=False, bin=3, notperhz=False):
+    """
+    Compute the power spectrum of a given time series.
+
+    Parameters:
+        time (array):
+            The time array.
+        quantity (array):
+            The data array for which the power spectrum is to be computed.
+        noline (bool):
+            If True, straight line is not subtracted from the data.
+        nohanning (bool):
+            If True, no Hanning window is applied to the data.
+        bin (int):
+            Bin size for binning the data. Default is 3.
+        notperhz (bool):
+            If True, the output units are the square of the input units.
+
+    Returns:
+        tuple: Tuple containing:
+            - freq (array):
+                The frequency array.
+            - power (array):
+                The computed power spectrum.
+
+    Notes:
+        This is similar to IDL pwrspc.pro routine.
+
+    Example:
+        >>> # Compute the power spectrum of a given time series
+        >>> from pytplot import pwrspc
+        >>> time = [1, 2, 3, 4, 5]
+        >>> quantity = [1, 2, 3, 4, 5]
+        >>> freq, power = pwrspc(time, quantity)
+    """
+
+    t = np.array(time, dtype=np.float64) - time[0]
+    x = np.array(quantity, dtype=np.float64)
+
+    if not noline:
+        slope, intercept, _, _, _ = linregress(t, x)
+        x -= (slope * t + intercept)
+
+    binsize = bin
+
+    window = 0.0
+    if not nohanning:
+        window = np.hanning(len(x))
+        x *= window
+
+    nt = len(t)
+    if nt % 2 != 0:
+        logging.info('needs an even number of data points, dropping last point...')
+        t = t[:-1]
+        x = x[:-1]
+        nt -= 1
+
+    xs2 = np.abs(np.fft.fft(x)) ** 2
+    dbign = float(nt)
+    logging.info('bign=', dbign)
+
+    k = np.arange(0, dbign // 2 + 1)
+    tres = np.median(np.diff(t))
+    fk = k / (dbign * tres)
+
+    pwr = np.zeros(nt // 2 + 1)
+    pwr[0] = xs2[0] / dbign**2
+    pwr[1:nt // 2] = (xs2[1:nt // 2] + xs2[nt:nt // 2:-1]) / dbign**2
+    pwr[-1] = xs2[-1] / dbign**2
+
+    if not nohanning:
+        wss = dbign * np.sum(window ** 2)
+        pwr = pwr*dbign**2/wss
+
+    dfreq = binsize * (fk[1] - fk[0])
+    npwr = len(pwr) - 1
+    nfinal = int(npwr / binsize)
+    iarray = np.arange(nfinal)
+    power = np.zeros(nfinal)
+
+    idx = (iarray + 0.5) * binsize + 1
+    freq = [fk[int(i)] for i in idx]
+
+    for i in range(binsize):
+        power += pwr[iarray * binsize + i + 1]
+
+    if not notperhz:
+        power /= dfreq
+
+    logging.info('dfreq=', dfreq)
+
+    return freq, power