Skip to content

Commit

Permalink
Merge pull request #689 from nickssl/master
Browse files Browse the repository at this point in the history 
Minor modifications to better match the IDL functions
  • Loading branch information
@nickssl
nickssl authored Jan 5, 2024
2 parents babf612 + 6f3eca0 commit e3f90ea
Show file tree
Hide file tree
Showing 2 changed files with 177 additions and 78 deletions.
127 changes: 85 additions & 42 deletions pyspedas/analysis/dpwrspc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,20 @@
import numpy as np


def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
nohanning=False, noline=False, notperhz=False, notmvariance=False,
tm_sensitivity=None):
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.
Expand All @@ -28,9 +39,15 @@ def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
nshiftpoints: int, optional
The number of points to shift for each spectrum.
The default is 128.
binsize: int, optional
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.
Expand Down Expand Up @@ -60,8 +77,30 @@ def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
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.hanning(nboxpoints)
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)
Expand All @@ -71,32 +110,33 @@ def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
nshiftpnts = nshiftpoints

totalpoints = len(times2process)
nspectra = int((totalpoints-nboxpnts/2.)/nshiftpnts)
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)
nbegin = np.array([nshiftpnts * i for i in range(nspectra)], dtype=np.int64)
nend = nbegin + nboxpnts

okspec = np.where(nend <= totalpoints-1)
okspec = np.where(nend <= totalpoints - 1)

if len(okspec[0]) <= 0:
logging.error('Not enough points for a calculation')
return
logging.error("Not enough points for a calculation")
return tdps0, fdps0, dps0

tdps = np.zeros(nspectra)
nfreqs = int(int(nboxpnts/2)/binsize)
nfreqs = int(int(nboxpnts / 2) / bin)

if nfreqs <= 1:
logging.error('Not enough frequencies for a calculation')
return
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)
nbegin = int(nthspectrum * nshiftpnts)
nend = nbegin + nboxpnts

if nend <= totalpoints:
Expand All @@ -106,7 +146,7 @@ def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
x = quantity2process[nbegin:nend]

# Use center time
tdps[nthspectrum] = (times2process[nbegin]+times2process[nend-1])/2.0
tdps[nthspectrum] = (times2process[nbegin] + times2process[nend - 1]) / 2.0

if noline is False:
coef = np.polyfit(t, x, 1)
Expand All @@ -115,15 +155,16 @@ def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
x = x - line

if nohanning is False:
x = x*window
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]
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)
Expand All @@ -136,50 +177,52 @@ def dpwrspc(time, quantity, nboxpoints=256, nshiftpoints=128, binsize=3,
else:
tmsn = 100.0

tdiff = t[1:n_tm]-t[0:n_tm-1]
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)
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')
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)
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
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
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 = bign*np.sum(window**2)
pwr = bign**2*pwr/wss
wss = float(bign) * float(np.sum(window**2))
pwr = bign**2 * pwr / wss

dfreq = binsize*(fk[1]-fk[0])
dfreq = bin * (fk[1] - fk[0])

npwr = len(pwr)-1
nfinal = int(npwr/binsize)
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*binsize+1]+fk[iarray*binsize+binsize])/2.
freqcenter = (fk[iarray * bin + 1] + fk[iarray * bin + bin]) / 2.0

for i in range(binsize):
power = power+pwr[iarray*binsize+i+1]
for i in range(bin):
power = power + pwr[iarray * bin + i + 1]

if notperhz is False:
power = power/dfreq
power = power / dfreq

dps[nthspectrum, :] = power
fdps[nthspectrum, :] = freqcenter
Expand Down
Loading

0 comments on commit e3f90ea

Please sign in to comment.