import numpy as np
try: # use scipy if available: it's faster from scipy.fftpack import fft, ifft, fftshift, ifftshift except: from numpy.fft import fft, ifft, fftshift, ifftshift
def FT_continuous(t, h, axis=-1, method=1): """Approximate a continuous 1D Fourier Transform with sampled data.
This function uses the Fast Fourier Transform to approximate
the continuous fourier transform of a sampled function, using
the convention
.. math::
H(f) = \int h(t) exp(-2 \pi i f t) dt
It returns f and H, which approximate H(f).
Parameters
----------
t : array_like
regularly sampled array of times
t is assumed to be regularly spaced, i.e.
t = t0 + Dt * np.arange(N)
h : array_like
real or complex signal at each time
axis : int
axis along which to perform fourier transform.
This axis must be the same length as t.
Returns
-------
f : ndarray
frequencies of result. Units are the same as 1/t
H : ndarray
Fourier coefficients at each frequency.
"""
assert t.ndim == 1
assert h.shape[axis] == t.shape[0]
N = len(t)
if N % 2 != 0:
raise ValueError("number of samples must be even")
Dt = t[1] - t[0]
Df = 1. / (N * Dt)
t0 = t[N / 2]
f = Df * (np.arange(N) - N / 2)
shape = np.ones(h.ndim, dtype=int)
shape[axis] = N
phase = np.ones(N)
phase[1::2] = -1
phase = phase.reshape(shape)
if method == 1:
H = Dt * fft(h * phase, axis=axis)
else:
H = Dt * fftshift(fft(h, axis=axis), axes=axis)
H *= phase
H *= np.exp(-2j * np.pi * t0 * f.reshape(shape))
H *= np.exp(-1j * np.pi * N / 2)
return f, H
def IFT_continuous(f, H, axis=-1, method=1): """Approximate a continuous 1D Inverse Fourier Transform with sampled data.
This function uses the Fast Fourier Transform to approximate
the continuous fourier transform of a sampled function, using
the convention
.. math::
H(f) = integral[ h(t) exp(-2 pi i f t) dt]
h(t) = integral[ H(f) exp(2 pi i f t) dt]
It returns t and h, which approximate h(t).
Parameters
----------
f : array_like
regularly sampled array of times
t is assumed to be regularly spaced, i.e.
f = f0 + Df * np.arange(N)
H : array_like
real or complex signal at each time
axis : int
axis along which to perform fourier transform.
This axis must be the same length as t.
Returns
-------
f : ndarray
frequencies of result. Units are the same as 1/t
H : ndarray
Fourier coefficients at each frequency.
"""
assert f.ndim == 1
assert H.shape[axis] == f.shape[0]
N = len(f)
if N % 2 != 0:
raise ValueError("number of samples must be even")
f0 = f[0]
Df = f[1] - f[0]
t0 = -0.5 / Df
Dt = 1. / (N * Df)
t = t0 + Dt * np.arange(N)
shape = np.ones(H.ndim, dtype=int)
shape[axis] = N
t_calc = t.reshape(shape)
f_calc = f.reshape(shape)
H_prime = H * np.exp(2j * np.pi * t0 * f_calc)
h_prime = ifft(H_prime, axis=axis)
h = N * Df * np.exp(2j * np.pi * f0 * (t_calc - t0)) * h_prime
return t, h
def PSD_continuous(t, h, axis=-1, method=1): """Approximate a continuous 1D Power Spectral Density of sampled data.
This function uses the Fast Fourier Transform to approximate
the continuous fourier transform of a sampled function, using
the convention
.. math::
H(f) = \int h(t) \exp(-2 \pi i f t) dt
It returns f and PSD, which approximate PSD(f) where
.. math::
PSD(f) = |H(f)|^2 + |H(-f)|^2
Parameters
----------
t : array_like
regularly sampled array of times
t is assumed to be regularly spaced, i.e.
t = t0 + Dt * np.arange(N)
h : array_like
real or complex signal at each time
axis : int
axis along which to perform fourier transform.
This axis must be the same length as t.
Returns
-------
f : ndarray
frequencies of result. Units are the same as 1/t
PSD : ndarray
Fourier coefficients at each frequency.
"""
assert t.ndim == 1
assert h.shape[axis] == t.shape[0]
N = len(t)
if N % 2 != 0:
raise ValueError("number of samples must be even")
ax = axis % h.ndim
if method == 1:
# use FT_continuous
f, Hf = FT_continuous(t, h, axis)
Hf = np.rollaxis(Hf, ax)
f = -f[N / 2::-1]
PSD = abs(Hf[N / 2::-1]) ** 2
PSD[:-1] += abs(Hf[N / 2:]) ** 2
PSD = np.rollaxis(PSD, 0, ax + 1)
else:
# A faster way to do it is with fftshift
# take advantage of the fact that phases go away
Dt = t[1] - t[0]
Df = 1. / (N * Dt)
f = Df * np.arange(N / 2 + 1)
Hf = fft(h, axis=axis)
Hf = np.rollaxis(Hf, ax)
PSD = abs(Hf[:N / 2 + 1]) ** 2
PSD[-1] = 0
PSD[1:] += abs(Hf[N / 2:][::-1]) ** 2
PSD[0] *= 2
PSD = Dt ** 2 * np.rollaxis(PSD, 0, ax + 1)
return f, PSD
def sinegauss(t, t0, f0, Q): """Sine-gaussian wavelet""" a = (f0 * 1. / Q) ** 2 return (np.exp(-a * (t - t0) ** 2) * np.exp(2j * np.pi * f0 * (t - t0)))
def sinegauss_FT(f, t0, f0, Q): """Fourier transform of the sine-gaussian wavelet.
This uses the convention
.. math::
H(f) = integral[ h(t) exp(-2pi i f t) dt]
"""
a = (f0 * 1. / Q) ** 2
return (np.sqrt(np.pi / a)
* np.exp(-2j * np.pi * f * t0)
* np.exp(-np.pi ** 2 * (f - f0) ** 2 / a))
def sinegauss_PSD(f, t0, f0, Q): """Compute the PSD of the sine-gaussian function at frequency f
.. math::
PSD(f) = |H(f)|^2 + |H(-f)|^2
"""
a = (f0 * 1. / Q) ** 2
Pf = np.pi / a * np.exp(-2 * np.pi ** 2 * (f - f0) ** 2 / a)
Pmf = np.pi / a * np.exp(-2 * np.pi ** 2 * (-f - f0) ** 2 / a)
return Pf + Pmf
def wavelet_PSD(t, h, f0, Q=1.0): """Compute the wavelet PSD as a function of f0 and t
Parameters
----------
t : array_like
array of times, length N
h : array_like
array of observed values, length N
f0 : array_like
array of candidate frequencies, length Nf
Q : float
Q-parameter for wavelet
Returns
-------
PSD : ndarray
The 2-dimensional PSD, of shape (Nf, N), corresponding with
frequencies f0 and times t.
"""
t, h, f0 = map(np.asarray, (t, h, f0))
if (t.ndim != 1) or (t.shape != h.shape):
raise ValueError('t and h must be one dimensional and the same shape')
if f0.ndim != 1:
raise ValueError('f0 must be one dimensional')
Q = Q + np.zeros_like(f0)
f, H = FT_continuous(t, h)
W = np.conj(sinegauss_FT(f, 0, f0[:, None], Q[:, None]))
_, HW = IFT_continuous(f, H * W)
return abs(HW) ** 2