<s> import numpy as np <EOL> import os . path as op <EOL> from numpy . testing import assert_array_almost_equal , assert_array_equal <EOL> from nose . tools import assert_true , assert_false , assert_equal , assert_raises <EOL> import mne <EOL> from mne import io , Epochs , read_events , pick_types , create_info , EpochsArray <EOL> from mne . utils import ( _TempDir , run_tests_if_main , slow_test , requires_h5py , <EOL> grand_average ) <EOL> from mne . time_frequency import single_trial_power <EOL> from mne . time_frequency . tfr import ( cwt_morlet , morlet , tfr_morlet , <EOL> _dpss_wavelet , tfr_multitaper , <EOL> AverageTFR , read_tfrs , write_tfrs , <EOL> combine_tfr , cwt ) <EOL> import matplotlib <EOL> matplotlib . use ( '<STR_LIT>' ) <EOL> raw_fname = op . join ( op . dirname ( __file__ ) , '<STR_LIT:..>' , '<STR_LIT:..>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:data>' , <EOL> '<STR_LIT>' ) <EOL> event_fname = op . join ( op . dirname ( __file__ ) , '<STR_LIT:..>' , '<STR_LIT:..>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT:data>' , '<STR_LIT>' ) <EOL> def test_morlet ( ) : <EOL> """<STR_LIT>""" <EOL> Wz = morlet ( <NUM_LIT:1000> , [ <NUM_LIT:10> ] , <NUM_LIT> , zero_mean = True ) <EOL> W = morlet ( <NUM_LIT:1000> , [ <NUM_LIT:10> ] , <NUM_LIT> , zero_mean = False ) <EOL> assert_true ( np . abs ( np . mean ( np . real ( Wz [ <NUM_LIT:0> ] ) ) ) < <NUM_LIT> ) <EOL> assert_true ( np . abs ( np . mean ( np . real ( W [ <NUM_LIT:0> ] ) ) ) > <NUM_LIT> ) <EOL> def test_time_frequency ( ) : <EOL> """<STR_LIT>""" <EOL> event_id = <NUM_LIT:1> <EOL> tmin = - <NUM_LIT> <EOL> tmax = <NUM_LIT> <EOL> raw = io . read_raw_fif ( raw_fname ) <EOL> events = read_events ( event_fname ) <EOL> include = [ ] <EOL> exclude = raw . info [ '<STR_LIT>' ] + [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> picks = pick_types ( raw . info , meg = '<STR_LIT>' , eeg = False , <EOL> stim = False , include = include , exclude = exclude ) <EOL> picks = picks [ : <NUM_LIT:2> ] <EOL> epochs = Epochs ( raw , events , event_id , tmin , tmax , picks = picks , <EOL> baseline = ( None , <NUM_LIT:0> ) ) <EOL> data = epochs . get_data ( ) <EOL> times = epochs . times <EOL> nave = len ( data ) <EOL> epochs_nopicks = Epochs ( raw , events , event_id , tmin , tmax , <EOL> baseline = ( None , <NUM_LIT:0> ) ) <EOL> freqs = np . arange ( <NUM_LIT:6> , <NUM_LIT:20> , <NUM_LIT:5> ) <EOL> n_cycles = freqs / <NUM_LIT> <EOL> power , itc = tfr_morlet ( epochs [ <NUM_LIT:0> ] , freqs = freqs , n_cycles = n_cycles , <EOL> use_fft = True , return_itc = True ) <EOL> evoked = epochs . average ( ) <EOL> power_evoked = tfr_morlet ( evoked , freqs , n_cycles , use_fft = True , <EOL> return_itc = False ) <EOL> assert_raises ( ValueError , tfr_morlet , evoked , freqs , <NUM_LIT:1.> , return_itc = True ) <EOL> power , itc = tfr_morlet ( epochs , freqs = freqs , n_cycles = n_cycles , <EOL> use_fft = True , return_itc = True ) <EOL> power_picks , itc_picks = tfr_morlet ( epochs_nopicks , freqs = freqs , <EOL> n_cycles = n_cycles , use_fft = True , <EOL> return_itc = True , picks = picks ) <EOL> assert_array_almost_equal ( power . data , power_picks . data ) <EOL> assert_array_almost_equal ( itc . data , itc_picks . data ) <EOL> assert_array_almost_equal ( power . data , power_evoked . data ) <EOL> print ( itc ) <EOL> print ( itc . ch_names ) <EOL> itc += power <EOL> itc -= power <EOL> power . apply_baseline ( baseline = ( - <NUM_LIT:0.1> , <NUM_LIT:0> ) , mode = '<STR_LIT>' ) <EOL> assert_true ( '<STR_LIT>' in power ) <EOL> assert_true ( '<STR_LIT>' in power ) <EOL> assert_false ( '<STR_LIT>' in power ) <EOL> assert_false ( '<STR_LIT>' in power ) <EOL> assert_equal ( power . nave , nave ) <EOL> assert_equal ( itc . nave , nave ) <EOL> assert_true ( power . data . shape == ( len ( picks ) , len ( freqs ) , len ( times ) ) ) <EOL> assert_true ( power . data . shape == itc . data . shape ) <EOL> assert_true ( np . sum ( itc . data >= <NUM_LIT:1> ) == <NUM_LIT:0> ) <EOL> assert_true ( np . sum ( itc . data <= <NUM_LIT:0> ) == <NUM_LIT:0> ) <EOL> itc2 = itc . copy ( ) <EOL> itc2 . info [ '<STR_LIT>' ] = [ itc2 . ch_names [ <NUM_LIT:0> ] ] <EOL> gave = grand_average ( [ itc2 , itc ] ) <EOL> assert_equal ( gave . data . shape , ( itc2 . data . shape [ <NUM_LIT:0> ] - <NUM_LIT:1> , <EOL> itc2 . data . shape [ <NUM_LIT:1> ] , <EOL> itc2 . data . shape [ <NUM_LIT:2> ] ) ) <EOL> assert_equal ( itc2 . ch_names [ <NUM_LIT:1> : ] , gave . ch_names ) <EOL> assert_equal ( gave . nave , <NUM_LIT:2> ) <EOL> itc2 . drop_channels ( itc2 . info [ "<STR_LIT>" ] ) <EOL> assert_array_almost_equal ( gave . data , itc2 . data ) <EOL> itc2 . data = np . ones ( itc2 . data . shape ) <EOL> itc . data = np . zeros ( itc . data . shape ) <EOL> itc2 . nave = <NUM_LIT:2> <EOL> itc . nave = <NUM_LIT:1> <EOL> itc . drop_channels ( [ itc . ch_names [ <NUM_LIT:0> ] ] ) <EOL> combined_itc = combine_tfr ( [ itc2 , itc ] ) <EOL> assert_array_almost_equal ( combined_itc . data , <EOL> np . ones ( combined_itc . data . shape ) * <NUM_LIT:2> / <NUM_LIT:3> ) <EOL> power , itc = tfr_morlet ( epochs , freqs = freqs , n_cycles = <NUM_LIT:2> , use_fft = False , <EOL> return_itc = True ) <EOL> assert_true ( power . data . shape == ( len ( picks ) , len ( freqs ) , len ( times ) ) ) <EOL> assert_true ( power . data . shape == itc . data . shape ) <EOL> assert_true ( np . sum ( itc . data >= <NUM_LIT:1> ) == <NUM_LIT:0> ) <EOL> assert_true ( np . sum ( itc . data <= <NUM_LIT:0> ) == <NUM_LIT:0> ) <EOL> Fs = raw . info [ '<STR_LIT>' ] <EOL> tfr = cwt_morlet ( data [ <NUM_LIT:0> ] , Fs , freqs , use_fft = True , n_cycles = <NUM_LIT:2> ) <EOL> assert_true ( tfr . shape == ( len ( picks ) , len ( freqs ) , len ( times ) ) ) <EOL> single_power = single_trial_power ( data , Fs , freqs , use_fft = False , <EOL> n_cycles = <NUM_LIT:2> ) <EOL> assert_array_almost_equal ( np . mean ( single_power ) , power . data ) <EOL> power_pick = power . pick_channels ( power . ch_names [ : <NUM_LIT:10> : <NUM_LIT:2> ] ) <EOL> assert_equal ( len ( power_pick . ch_names ) , len ( power . ch_names [ : <NUM_LIT:10> : <NUM_LIT:2> ] ) ) <EOL> assert_equal ( power_pick . data . shape [ <NUM_LIT:0> ] , len ( power . ch_names [ : <NUM_LIT:10> : <NUM_LIT:2> ] ) ) <EOL> power_drop = power . drop_channels ( power . ch_names [ <NUM_LIT:1> : <NUM_LIT:10> : <NUM_LIT:2> ] ) <EOL> assert_equal ( power_drop . ch_names , power_pick . ch_names ) <EOL> assert_equal ( power_pick . data . shape [ <NUM_LIT:0> ] , len ( power_drop . ch_names ) ) <EOL> mne . equalize_channels ( [ power_pick , power_drop ] ) <EOL> assert_equal ( power_pick . ch_names , power_drop . ch_names ) <EOL> assert_equal ( power_pick . data . shape , power_drop . data . shape ) <EOL> for decim in [ <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:8> , <NUM_LIT:9> ] : <EOL> for use_fft in [ True , False ] : <EOL> power , itc = tfr_morlet ( epochs , freqs = freqs , n_cycles = <NUM_LIT:2> , <EOL> use_fft = use_fft , return_itc = True , <EOL> decim = decim ) <EOL> assert_equal ( power . data . shape [ <NUM_LIT:2> ] , <EOL> np . ceil ( float ( len ( times ) ) / decim ) ) <EOL> freqs = range ( <NUM_LIT:50> , <NUM_LIT> ) <EOL> decim = <NUM_LIT:2> <EOL> _ , n_chan , n_time = data . shape <EOL> tfr = cwt_morlet ( data [ <NUM_LIT:0> , : , : ] , sfreq = epochs . info [ '<STR_LIT>' ] , <EOL> freqs = freqs , decim = decim ) <EOL> assert_equal ( tfr . shape , ( n_chan , len ( freqs ) , n_time // decim ) ) <EOL> Ws = morlet ( <NUM_LIT> , [ <NUM_LIT:10> , <NUM_LIT:20> ] , n_cycles = <NUM_LIT:2> ) <EOL> assert_raises ( ValueError , cwt , data [ <NUM_LIT:0> , : , : ] , Ws , mode = '<STR_LIT:foo>' ) <EOL> for use_fft in [ True , False ] : <EOL> for mode in [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] : <EOL> if ( not use_fft ) and mode == '<STR_LIT>' : <EOL> assert_raises ( ValueError , cwt , data [ <NUM_LIT:0> , : , : ] , Ws , <EOL> use_fft = use_fft , mode = mode ) <EOL> continue <EOL> cwt ( data [ <NUM_LIT:0> , : , : ] , Ws , use_fft = use_fft , mode = mode ) <EOL> def test_dpsswavelet ( ) : <EOL> """<STR_LIT>""" <EOL> freqs = np . arange ( <NUM_LIT:5> , <NUM_LIT> , <NUM_LIT:3> ) <EOL> Ws = _dpss_wavelet ( <NUM_LIT:1000> , freqs = freqs , n_cycles = freqs / <NUM_LIT> , <EOL> time_bandwidth = <NUM_LIT> , zero_mean = True ) <EOL> assert_true ( len ( Ws ) == <NUM_LIT:3> ) <EOL> assert_true ( np . abs ( np . mean ( np . real ( Ws [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] ) ) ) < <NUM_LIT> ) <EOL> assert_true ( len ( Ws [ <NUM_LIT:0> ] ) == len ( freqs ) ) <EOL> @ slow_test <EOL> def test_tfr_multitaper ( ) : <EOL> """<STR_LIT>""" <EOL> sfreq = <NUM_LIT> <EOL> ch_names = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> ch_types = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> info = create_info ( ch_names = ch_names , sfreq = sfreq , ch_types = ch_types ) <EOL> n_times = int ( sfreq ) <EOL> n_epochs = <NUM_LIT:3> <EOL> seed = <NUM_LIT> <EOL> rng = np . random . RandomState ( seed ) <EOL> noise = <NUM_LIT:0.1> * rng . randn ( n_epochs , len ( ch_names ) , n_times ) <EOL> t = np . arange ( n_times , dtype = np . float ) / sfreq <EOL> signal = np . sin ( np . pi * <NUM_LIT> * <NUM_LIT> * t ) <EOL> signal [ np . logical_or ( t < <NUM_LIT> , t > <NUM_LIT> ) ] = <NUM_LIT:0.> <EOL> on_time = np . logical_and ( t >= <NUM_LIT> , t <= <NUM_LIT> ) <EOL> signal [ on_time ] *= np . hanning ( on_time . sum ( ) ) <EOL> dat = noise + signal <EOL> reject = dict ( grad = <NUM_LIT> ) <EOL> events = np . empty ( ( n_epochs , <NUM_LIT:3> ) , int ) <EOL> first_event_sample = <NUM_LIT:100> <EOL> event_id = dict ( sin50hz = <NUM_LIT:1> ) <EOL> for k in range ( n_epochs ) : <EOL> events [ k , : ] = first_event_sample + k * n_times , <NUM_LIT:0> , event_id [ '<STR_LIT>' ] <EOL> epochs = EpochsArray ( data = dat , info = info , events = events , event_id = event_id , <EOL> reject = reject ) <EOL> freqs = np . arange ( <NUM_LIT:5> , <NUM_LIT:100> , <NUM_LIT:3> , dtype = np . float ) <EOL> power , itc = tfr_multitaper ( epochs , freqs = freqs , n_cycles = freqs / <NUM_LIT> , <EOL> time_bandwidth = <NUM_LIT> ) <EOL> picks = np . arange ( len ( ch_names ) ) <EOL> power_picks , itc_picks = tfr_multitaper ( epochs , freqs = freqs , <EOL> n_cycles = freqs / <NUM_LIT> , <EOL> time_bandwidth = <NUM_LIT> , picks = picks ) <EOL> power_evoked = tfr_multitaper ( epochs . average ( ) , freqs = freqs , <EOL> n_cycles = freqs / <NUM_LIT> , time_bandwidth = <NUM_LIT> , <EOL> return_itc = False ) <EOL> assert_array_almost_equal ( power . data , power_picks . data ) <EOL> assert_array_almost_equal ( itc . data , itc_picks . data ) <EOL> assert_array_equal ( power . data . shape , power_evoked . data . shape ) <EOL> assert_raises ( AssertionError , assert_array_almost_equal , <EOL> power . data , power_evoked . data ) <EOL> tmax = t [ np . argmax ( itc . data [ <NUM_LIT:0> , freqs == <NUM_LIT:50> , : ] ) ] <EOL> fmax = freqs [ np . argmax ( power . data [ <NUM_LIT:1> , : , t == <NUM_LIT:0.5> ] ) ] <EOL> assert_true ( tmax > <NUM_LIT> and tmax < <NUM_LIT> ) <EOL> assert_false ( np . any ( itc . data < <NUM_LIT:0.> ) ) <EOL> assert_true ( fmax > <NUM_LIT> and fmax < <NUM_LIT> ) <EOL> def test_crop ( ) : <EOL> """<STR_LIT>""" <EOL> data = np . zeros ( ( <NUM_LIT:3> , <NUM_LIT:2> , <NUM_LIT:3> ) ) <EOL> times = np . array ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> freqs = np . array ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> info = mne . create_info ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <NUM_LIT> , <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> tfr = AverageTFR ( info , data = data , times = times , freqs = freqs , <EOL> nave = <NUM_LIT:20> , comment = '<STR_LIT:test>' , method = '<STR_LIT>' ) <EOL> tfr . crop ( <NUM_LIT> , <NUM_LIT> ) <EOL> assert_array_equal ( tfr . times , [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> assert_equal ( tfr . data . shape [ - <NUM_LIT:1> ] , <NUM_LIT:2> ) <EOL> @ requires_h5py <EOL> def test_io ( ) : <EOL> """<STR_LIT>""" <EOL> tempdir = _TempDir ( ) <EOL> fname = op . join ( tempdir , '<STR_LIT>' ) <EOL> data = np . zeros ( ( <NUM_LIT:3> , <NUM_LIT:2> , <NUM_LIT:3> ) ) <EOL> times = np . array ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> freqs = np . array ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> info = mne . create_info ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <NUM_LIT> , <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> tfr = AverageTFR ( info , data = data , times = times , freqs = freqs , <EOL> nave = <NUM_LIT:20> , comment = '<STR_LIT:test>' , method = '<STR_LIT>' ) <EOL> tfr . save ( fname ) <EOL> tfr2 = read_tfrs ( fname , condition = '<STR_LIT:test>' ) <EOL> assert_array_equal ( tfr . data , tfr2 . data ) <EOL> assert_array_equal ( tfr . times , tfr2 . times ) <EOL> assert_array_equal ( tfr . freqs , tfr2 . freqs ) <EOL> assert_equal ( tfr . comment , tfr2 . comment ) <EOL> assert_equal ( tfr . nave , tfr2 . nave ) <EOL> assert_raises ( IOError , tfr . save , fname ) <EOL> tfr . comment = None <EOL> tfr . save ( fname , overwrite = True ) <EOL> assert_equal ( read_tfrs ( fname , condition = <NUM_LIT:0> ) . comment , tfr . comment ) <EOL> tfr . comment = '<STR_LIT>' <EOL> tfr2 . comment = '<STR_LIT>' <EOL> fname = op . join ( tempdir , '<STR_LIT>' ) <EOL> write_tfrs ( fname , [ tfr , tfr2 ] ) <EOL> tfr3 = read_tfrs ( fname , condition = '<STR_LIT>' ) <EOL> assert_equal ( tfr . comment , tfr3 . comment ) <EOL> assert_true ( isinstance ( tfr . info , mne . Info ) ) <EOL> tfrs = read_tfrs ( fname , condition = None ) <EOL> assert_equal ( len ( tfrs ) , <NUM_LIT:2> ) <EOL> tfr4 = tfrs [ <NUM_LIT:1> ] <EOL> assert_equal ( tfr2 . comment , tfr4 . comment ) <EOL> assert_raises ( ValueError , read_tfrs , fname , condition = '<STR_LIT>' ) <EOL> def test_plot ( ) : <EOL> """<STR_LIT>""" <EOL> import matplotlib . pyplot as plt <EOL> data = np . zeros ( ( <NUM_LIT:3> , <NUM_LIT:2> , <NUM_LIT:3> ) ) <EOL> times = np . array ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> freqs = np . array ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> info = mne . create_info ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <NUM_LIT> , <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> tfr = AverageTFR ( info , data = data , times = times , freqs = freqs , <EOL> nave = <NUM_LIT:20> , comment = '<STR_LIT:test>' , method = '<STR_LIT>' ) <EOL> tfr . plot ( [ <NUM_LIT:1> , <NUM_LIT:2> ] , title = '<STR_LIT:title>' ) <EOL> plt . close ( '<STR_LIT:all>' ) <EOL> ax = plt . subplot2grid ( ( <NUM_LIT:2> , <NUM_LIT:2> ) , ( <NUM_LIT:0> , <NUM_LIT:0> ) ) <EOL> ax2 = plt . subplot2grid ( ( <NUM_LIT:2> , <NUM_LIT:2> ) , ( <NUM_LIT:1> , <NUM_LIT:1> ) ) <EOL> ax3 = plt . subplot2grid ( ( <NUM_LIT:2> , <NUM_LIT:2> ) , ( <NUM_LIT:0> , <NUM_LIT:1> ) ) <EOL> tfr . plot ( picks = [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> ] , axes = [ ax , ax2 , ax3 ] ) <EOL> plt . close ( '<STR_LIT:all>' ) <EOL> tfr . plot_topo ( picks = [ <NUM_LIT:1> , <NUM_LIT:2> ] ) <EOL> plt . close ( '<STR_LIT:all>' ) <EOL> tfr . plot_topo ( picks = [ <NUM_LIT:1> , <NUM_LIT:2> ] ) <EOL> plt . close ( '<STR_LIT:all>' ) <EOL> def test_add_channels ( ) : <EOL> """<STR_LIT>""" <EOL> data = np . zeros ( ( <NUM_LIT:6> , <NUM_LIT:2> , <NUM_LIT:3> ) ) <EOL> times = np . array ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> freqs = np . array ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> info = mne . create_info ( <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <EOL> <NUM_LIT> , [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> tfr = AverageTFR ( info , data = data , times = times , freqs = freqs , <EOL> nave = <NUM_LIT:20> , comment = '<STR_LIT:test>' , method = '<STR_LIT>' ) <EOL> tfr_eeg = tfr . copy ( ) . pick_types ( meg = False , eeg = True ) <EOL> tfr_meg = tfr . copy ( ) . pick_types ( meg = True ) <EOL> tfr_stim = tfr . copy ( ) . pick_types ( meg = False , stim = True ) <EOL> tfr_eeg_meg = tfr . copy ( ) . pick_types ( meg = True , eeg = True ) <EOL> tfr_new = tfr_meg . copy ( ) . add_channels ( [ tfr_eeg , tfr_stim ] ) <EOL> assert_true ( all ( ch in tfr_new . ch_names <EOL> for ch in tfr_stim . ch_names + tfr_meg . ch_names ) ) <EOL> tfr_new = tfr_meg . copy ( ) . add_channels ( [ tfr_eeg ] ) <EOL> assert_true ( ch in tfr_new . ch_names for ch in tfr . ch_names ) <EOL> assert_array_equal ( tfr_new . data , tfr_eeg_meg . data ) <EOL> assert_true ( all ( ch not in tfr_new . ch_names <EOL> for ch in tfr_stim . ch_names ) ) <EOL> tfr_badsf = tfr_eeg . copy ( ) <EOL> tfr_badsf . info [ '<STR_LIT>' ] = <NUM_LIT> <EOL> tfr_eeg = tfr_eeg . crop ( - <NUM_LIT> , <NUM_LIT> ) <EOL> assert_raises ( RuntimeError , tfr_meg . add_channels , [ tfr_badsf ] ) <EOL> assert_raises ( AssertionError , tfr_meg . add_channels , [ tfr_eeg ] ) <EOL> assert_raises ( ValueError , tfr_meg . add_channels , [ tfr_meg ] ) <EOL> assert_raises ( AssertionError , tfr_meg . add_channels , tfr_badsf ) <EOL> run_tests_if_main ( ) </s>
