Skip to content

Instantly share code, notes, and snippets.

@fasiha
Last active June 23, 2022 12:44
Show Gist options
  • Select an option

  • Save fasiha/5702e57a54e9280d117dec9ea1a50e78 to your computer and use it in GitHub Desktop.

Select an option

Save fasiha/5702e57a54e9280d117dec9ea1a50e78 to your computer and use it in GitHub Desktop.
Short-time Fourier transform with inverse in Python/Numpy, see https://stackoverflow.com/q/51655119/500207
import numpy as np
import numpy.fft as fft
def stft(x, Nwin, Nfft=None):
"""
Short-time Fourier transform: convert a 1D vector to a 2D array
The short-time Fourier transform (STFT) breaks a long vector into disjoint
chunks (no overlap) and runs an FFT (Fast Fourier Transform) on each chunk.
The resulting 2D array can
Parameters
----------
x : array_like
Input signal (expected to be real)
Nwin : int
Length of each window (chunk of the signal). Should be ≪ `len(x)`.
Nfft : int, optional
Zero-pad each chunk to this length before FFT. Should be ≥ `Nwin`,
(usually with small prime factors, for fastest FFT). Default: `Nwin`.
Returns
-------
out : complex ndarray
`len(x) // Nwin` by `Nfft` complex array representing the STFT of `x`.
See also
--------
istft : inverse function (convert a STFT array back to a data vector)
stftbins : time and frequency bins corresponding to `out`
"""
Nfft = Nfft or Nwin
Nwindows = x.size // Nwin
# reshape into array `Nwin` wide, and as tall as possible. This is
# optimized for C-order (row-major) layouts.
arr = np.reshape(x[:Nwindows * Nwin], (-1, Nwin))
stft = fft.rfft(arr, Nfft)
return stft
def stftbins(x, Nwin, Nfft=None, d=1.0):
"""
Time and frequency bins corresponding to short-time Fourier transform.
Call this with the same arguments as `stft`, plus one extra argument: `d`
sample spacing, to get the time and frequency axes that the output of
`stft` correspond to.
Parameters
----------
x : array_like
same as `stft`
Nwin : int
same as `stft`
Nfft : int, optional
same as `stft`
d : float, optional
Sample spacing of `x` (or 1 / sample frequency), units of seconds.
Default: 1.0.
Returns
-------
t : ndarray
Array of length `len(x) // Nwin`, in units of seconds, corresponding to
the first dimension (height) of the output of `stft`.
f : ndarray
Array of length `Nfft`, in units of Hertz, corresponding to the second
dimension (width) of the output of `stft`.
"""
Nfft = Nfft or Nwin
Nwindows = x.size // Nwin
t = np.arange(Nwindows) * (Nwin * d)
f = fft.rfftfreq(Nfft, d)
return t, f
def istft(stftArr, Nwin):
"""
Inverse short-time Fourier transform (ISTFT)
Given an array representing the output of `stft`, convert it back to the
original samples.
Parameters
----------
stftArr : ndarray
Output of `stft` (or something the same size)
Nwin : int
Same input as `stft`: length of each chunk that the STFT was calculated
over.
Returns
-------
y : ndarray
Data samples corresponding to STFT data.
See also:
stft : the forward transform
"""
arr = fft.irfft(stftArr)[:, :Nwin]
return np.reshape(arr, -1)
if __name__ == '__main__':
sampleRate = 100.0 # Hertz
N = 1024
Nwin = 64
# Generate a chirp: start frequency at 5 Hz and going down at 2 Hz/s
time = np.arange(N) / sampleRate # seconds
x = np.cos(2 * np.pi * time * (5 - 2 * 0.5 * time))
# Test with Nfft bigger than Nwin
Nfft = Nwin * 2
s = stft(x, Nwin, Nfft=Nfft)
y = istft(s, Nwin)
# Make sure the stft and istft are inverses. Caveat: `x` and `y` won't be
# the same length if `N/Nwin` isn't integral!
maxerr = np.max(np.abs(x - y))
assert (maxerr < np.spacing(1) * 10)
# Test `stftbins`
t, f = stftbins(x, Nwin, Nfft=Nfft, d=1 / sampleRate)
assert (len(t) == s.shape[0])
assert (len(f) == s.shape[1])
try:
import pylab as plt
plt.imshow(np.abs(s), aspect="auto", extent=[f[0], f[-1], t[-1], t[0]])
plt.xlabel('frequency (Hertz)')
plt.ylabel('time (seconds (start of chunk))')
plt.title('STFT with chirp example')
plt.grid()
plt.show()
except ModuleNotFoundError:
pass
@GChims

GChims commented Jun 23, 2020

Copy link
Copy Markdown

Can I use the same implementation on real sensor data for example with 2 columns 'datetime' and 'value' columns. Thanks

@fasiha

fasiha commented Jun 23, 2020

Copy link
Copy Markdown
Author

These functions take just raw Numpy 1D arrays so yes, give stft the value column and it it will work.

@nehat005

Copy link
Copy Markdown

Thanks for the code!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment