Script to confirm that STFT can be invertible without the COLA constraint being satisfied

cola_test.py

"""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()

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 by allclose.