Script to confirm that STFT can be invertible without the COLA constraint being satisfied
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""Framework to confirm that STFT is invertible even if window is not COLA | |
For the theory behind this, see: https://gauss256.github.io/blog/cola.html""" | |
import numpy as np | |
from scipy import signal | |
def stft(x, w, nperseg, nhop): | |
"""Forward STFT""" | |
X = np.array( | |
[np.fft.fft(w * x[i:i + nperseg]) for i in range(0, 1 + len(x) - nperseg, nhop)] | |
) | |
return X | |
def istft(X, w, nperseg, nhop): | |
"""Inverse STFT""" | |
x = np.zeros(X.shape[0] * nhop) | |
wsum = np.zeros(X.shape[0] * nhop) | |
for n, i in enumerate(range(0, 1 + len(x) - nperseg, nhop)): | |
x[i:i + nperseg] += np.real(np.fft.ifft(X[n])) * w | |
wsum[i:i + nperseg] += w ** 2. | |
pos = wsum != 0 | |
x[pos] /= wsum[pos] | |
return x | |
def main(): | |
"""Change the parameters and window below to test other scenarios""" | |
length = 4096 | |
nperseg = 256 | |
nhop = 64 | |
noverlap = nperseg - nhop | |
w = np.sqrt(signal.hann(nperseg)) | |
pad = nperseg | |
# The theory assumes that the signal has been zero-padded to infinity, | |
# so some attention needs to be paid to the boundaries. The following | |
# is not elegant but it serves the purpose. | |
x = np.random.randn(length + 2 * pad) | |
if pad > 0: | |
x[:pad] = 0 | |
x[-pad:] = 0 | |
X = stft(x, w, nperseg, nhop) | |
y = istft(X, w, nperseg, nhop)[pad:pad + length] | |
x_hat = x[pad:pad + length] | |
print(f"check_COLA: {signal.check_COLA(w, nperseg, noverlap)}") | |
print(f"allclose : {np.allclose(x_hat, y)}") | |
if __name__ == "__main__": | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
The window used is square root Hann. It is not COLA, as confirmed by
check_COLA
. But the STFT that uses it is invertible, as confirmed byallclose
.