<s> import numpy as np <EOL> from scipy import linalg <EOL> from numpy . testing import assert_almost_equal , assert_array_almost_equal <EOL> from nose . tools import assert_true <EOL> from mne . time_frequency . stft import stft , istft , stftfreq , stft_norm2 <EOL> def test_stft ( ) : <EOL> "<STR_LIT>" <EOL> sfreq = <NUM_LIT> <EOL> f = <NUM_LIT> <EOL> for T in [ <NUM_LIT> , <NUM_LIT> ] : <EOL> t = np . arange ( T ) . astype ( np . float ) <EOL> x = np . sin ( <NUM_LIT:2> * np . pi * f * t / sfreq ) <EOL> x = np . array ( [ x , x + <NUM_LIT:1.> ] ) <EOL> wsize = <NUM_LIT> <EOL> tstep = <NUM_LIT:4> <EOL> X = stft ( x , wsize , tstep ) <EOL> xp = istft ( X , tstep , Tx = T ) <EOL> freqs = stftfreq ( wsize , sfreq = <NUM_LIT:1000> ) <EOL> max_freq = freqs [ np . argmax ( np . sum ( np . abs ( X [ <NUM_LIT:0> ] ) ** <NUM_LIT:2> , axis = <NUM_LIT:1> ) ) ] <EOL> assert_true ( X . shape [ <NUM_LIT:1> ] == len ( freqs ) ) <EOL> assert_true ( np . all ( freqs >= <NUM_LIT:0.> ) ) <EOL> assert_true ( np . abs ( max_freq - f ) < <NUM_LIT:1.> ) <EOL> assert_array_almost_equal ( x , xp , decimal = <NUM_LIT:6> ) <EOL> assert_almost_equal ( np . sqrt ( stft_norm2 ( X ) ) , <EOL> [ linalg . norm ( xx ) for xx in x ] , decimal = <NUM_LIT:6> ) <EOL> x = np . random . randn ( <NUM_LIT:2> , T ) <EOL> wsize = <NUM_LIT:16> <EOL> tstep = <NUM_LIT:8> <EOL> X = stft ( x , wsize , tstep ) <EOL> xp = istft ( X , tstep , Tx = T ) <EOL> freqs = stftfreq ( wsize , sfreq = <NUM_LIT:1000> ) <EOL> max_freq = freqs [ np . argmax ( np . sum ( np . abs ( X [ <NUM_LIT:0> ] ) ** <NUM_LIT:2> , axis = <NUM_LIT:1> ) ) ] <EOL> assert_true ( X . shape [ <NUM_LIT:1> ] == len ( freqs ) ) <EOL> assert_true ( np . all ( freqs >= <NUM_LIT:0.> ) ) <EOL> assert_array_almost_equal ( x , xp , decimal = <NUM_LIT:6> ) <EOL> assert_almost_equal ( np . sqrt ( stft_norm2 ( X ) ) , <EOL> [ linalg . norm ( xx ) for xx in x ] , <EOL> decimal = <NUM_LIT:6> ) <EOL> x = np . zeros ( ( <NUM_LIT:0> , T ) ) <EOL> X = stft ( x , wsize , tstep ) <EOL> xp = istft ( X , tstep , T ) <EOL> assert_true ( xp . shape == x . shape ) </s>
