<s> """<STR_LIT>""" <EOL> from __future__ import print_function <EOL> import numpy as np <EOL> import numpy . fft as fft <EOL> from scipy import signal <EOL> from statsmodels . sandbox . archive . linalg_decomp_1 import OneTimeProperty <EOL> from statsmodels . tsa . arima_process import ArmaProcess <EOL> class ArmaFft ( ArmaProcess ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , ar , ma , n ) : <EOL> super ( ArmaFft , self ) . __init__ ( ar , ma ) <EOL> self . ar = np . asarray ( ar ) <EOL> self . ma = np . asarray ( ma ) <EOL> self . nobs = n <EOL> self . arpoly = np . polynomial . Polynomial ( ar ) <EOL> self . mapoly = np . polynomial . Polynomial ( ma ) <EOL> self . nar = len ( ar ) <EOL> self . nma = len ( ma ) <EOL> def padarr ( self , arr , maxlag , atend = True ) : <EOL> '''<STR_LIT>''' <EOL> if atend : <EOL> return np . r_ [ arr , np . zeros ( maxlag - len ( arr ) ) ] <EOL> else : <EOL> return np . r_ [ np . zeros ( maxlag - len ( arr ) ) , arr ] <EOL> def pad ( self , maxlag ) : <EOL> '''<STR_LIT>''' <EOL> arpad = np . r_ [ self . ar , np . zeros ( maxlag - self . nar ) ] <EOL> mapad = np . r_ [ self . ma , np . zeros ( maxlag - self . nma ) ] <EOL> return arpad , mapad <EOL> def fftar ( self , n = None ) : <EOL> '''<STR_LIT>''' <EOL> if n is None : <EOL> n = len ( self . ar ) <EOL> return fft . fft ( self . padarr ( self . ar , n ) ) <EOL> def fftma ( self , n ) : <EOL> '''<STR_LIT>''' <EOL> if n is None : <EOL> n = len ( self . ar ) <EOL> return fft . fft ( self . padarr ( self . ma , n ) ) <EOL> def fftarma ( self , n = None ) : <EOL> '''<STR_LIT>''' <EOL> if n is None : <EOL> n = self . nobs <EOL> return ( self . fftma ( n ) / self . fftar ( n ) ) <EOL> def spd ( self , npos ) : <EOL> '''<STR_LIT>''' <EOL> n = npos <EOL> w = fft . fftfreq ( <NUM_LIT:2> * n ) * <NUM_LIT:2> * np . pi <EOL> hw = self . fftarma ( <NUM_LIT:2> * n ) <EOL> return ( hw * hw . conj ( ) ) . real * <NUM_LIT:0.5> / np . pi , w <EOL> def spdshift ( self , n ) : <EOL> '''<STR_LIT>''' <EOL> mapadded = self . padarr ( self . ma , n ) <EOL> arpadded = self . padarr ( self . ar , n ) <EOL> hw = fft . fft ( fft . fftshift ( mapadded ) ) / fft . fft ( fft . fftshift ( arpadded ) ) <EOL> w = fft . fftfreq ( n ) * <NUM_LIT:2> * np . pi <EOL> wslice = slice ( n // <NUM_LIT:2> - <NUM_LIT:1> , None , None ) <EOL> return ( hw * hw . conj ( ) ) . real , w <EOL> def spddirect ( self , n ) : <EOL> '''<STR_LIT>''' <EOL> hw = fft . fft ( self . ma , n ) / fft . fft ( self . ar , n ) <EOL> w = fft . fftfreq ( n ) * <NUM_LIT:2> * np . pi <EOL> wslice = slice ( None , n // <NUM_LIT:2> , None ) <EOL> return ( np . abs ( hw ) ** <NUM_LIT:2> ) * <NUM_LIT:0.5> / np . pi , w <EOL> def _spddirect2 ( self , n ) : <EOL> '''<STR_LIT>''' <EOL> hw = ( fft . fft ( np . r_ [ self . ma [ : : - <NUM_LIT:1> ] , self . ma ] , n ) <EOL> / fft . fft ( np . r_ [ self . ar [ : : - <NUM_LIT:1> ] , self . ar ] , n ) ) <EOL> return ( hw * hw . conj ( ) ) <EOL> def spdroots ( self , w ) : <EOL> '''<STR_LIT>''' <EOL> return self . spdroots_ ( self . arroots , self . maroots , w ) <EOL> def spdroots_ ( self , arroots , maroots , w ) : <EOL> '''<STR_LIT>''' <EOL> w = np . atleast_2d ( w ) . T <EOL> cosw = np . cos ( w ) <EOL> maroots = <NUM_LIT:1.> / maroots <EOL> arroots = <NUM_LIT:1.> / arroots <EOL> num = <NUM_LIT:1> + maroots ** <NUM_LIT:2> - <NUM_LIT:2> * maroots * cosw <EOL> den = <NUM_LIT:1> + arroots ** <NUM_LIT:2> - <NUM_LIT:2> * arroots * cosw <EOL> hw = <NUM_LIT:0.5> / np . pi * num . prod ( - <NUM_LIT:1> ) / den . prod ( - <NUM_LIT:1> ) <EOL> return np . squeeze ( hw ) , w . squeeze ( ) <EOL> def spdpoly ( self , w , nma = <NUM_LIT:50> ) : <EOL> '''<STR_LIT>''' <EOL> mpoly = np . polynomial . Polynomial ( self . arma2ma ( nma ) ) <EOL> hw = mpoly ( np . exp ( <NUM_LIT> * w ) ) <EOL> spd = np . real_if_close ( hw * hw . conj ( ) * <NUM_LIT:0.5> / np . pi ) <EOL> return spd , w <EOL> def filter ( self , x ) : <EOL> '''<STR_LIT>''' <EOL> n = x . shape [ <NUM_LIT:0> ] <EOL> if n == self . fftarma : <EOL> fftarma = self . fftarma <EOL> else : <EOL> fftarma = self . fftma ( n ) / self . fftar ( n ) <EOL> tmpfft = fftarma * fft . fft ( x ) <EOL> return fft . ifft ( tmpfft ) <EOL> def filter2 ( self , x , pad = <NUM_LIT:0> ) : <EOL> '''<STR_LIT>''' <EOL> from statsmodels . tsa . filters import fftconvolve3 <EOL> if not pad : <EOL> pass <EOL> elif pad == '<STR_LIT>' : <EOL> x = self . padarr ( x , x . shape [ <NUM_LIT:0> ] + <NUM_LIT:2> * ( self . nma + self . nar ) , atend = False ) <EOL> else : <EOL> x = self . padarr ( x , x . shape [ <NUM_LIT:0> ] + int ( pad ) , atend = False ) <EOL> return fftconvolve3 ( x , self . ma , self . ar ) <EOL> def acf2spdfreq ( self , acovf , nfreq = <NUM_LIT:100> , w = None ) : <EOL> '''<STR_LIT>''' <EOL> if w is None : <EOL> w = np . linspace ( <NUM_LIT:0> , np . pi , nfreq ) [ : , None ] <EOL> nac = len ( acovf ) <EOL> hw = <NUM_LIT:0.5> / np . pi * ( acovf [ <NUM_LIT:0> ] + <EOL> <NUM_LIT:2> * ( acovf [ <NUM_LIT:1> : ] * np . cos ( w * np . arange ( <NUM_LIT:1> , nac ) ) ) . sum ( <NUM_LIT:1> ) ) <EOL> return hw <EOL> def invpowerspd ( self , n ) : <EOL> '''<STR_LIT>''' <EOL> hw = self . fftarma ( n ) <EOL> return np . real_if_close ( fft . ifft ( hw * hw . conj ( ) ) , tol = <NUM_LIT:200> ) [ : n ] <EOL> def spdmapoly ( self , w , twosided = False ) : <EOL> '''<STR_LIT>''' <EOL> if w is None : <EOL> w = np . linspace ( <NUM_LIT:0> , np . pi , nfreq ) <EOL> return <NUM_LIT:0.5> / np . pi * self . mapoly ( np . exp ( w * <NUM_LIT> ) ) <EOL> def plot4 ( self , fig = None , nobs = <NUM_LIT:100> , nacf = <NUM_LIT:20> , nfreq = <NUM_LIT:100> ) : <EOL> rvs = self . generate_sample ( nsample = <NUM_LIT:100> , burnin = <NUM_LIT> ) <EOL> acf = self . acf ( nacf ) [ : nacf ] <EOL> pacf = self . pacf ( nacf ) <EOL> w = np . linspace ( <NUM_LIT:0> , np . pi , nfreq ) <EOL> spdr , wr = self . spdroots ( w ) <EOL> if fig is None : <EOL> import matplotlib . pyplot as plt <EOL> fig = plt . figure ( ) <EOL> ax = fig . add_subplot ( <NUM_LIT:2> , <NUM_LIT:2> , <NUM_LIT:1> ) <EOL> ax . plot ( rvs ) <EOL> ax . set_title ( '<STR_LIT>' % ( self . ar , self . ma ) ) <EOL> ax = fig . add_subplot ( <NUM_LIT:2> , <NUM_LIT:2> , <NUM_LIT:2> ) <EOL> ax . plot ( acf ) <EOL> ax . set_title ( '<STR_LIT>' % ( self . ar , self . ma ) ) <EOL> ax = fig . add_subplot ( <NUM_LIT:2> , <NUM_LIT:2> , <NUM_LIT:3> ) <EOL> ax . plot ( wr , spdr ) <EOL> ax . set_title ( '<STR_LIT>' % ( self . ar , self . ma ) ) <EOL> ax = fig . add_subplot ( <NUM_LIT:2> , <NUM_LIT:2> , <NUM_LIT:4> ) <EOL> ax . plot ( pacf ) <EOL> ax . set_title ( '<STR_LIT>' % ( self . ar , self . ma ) ) <EOL> return fig <EOL> def spdar1 ( ar , w ) : <EOL> if np . ndim ( ar ) == <NUM_LIT:0> : <EOL> rho = ar <EOL> else : <EOL> rho = - ar [ <NUM_LIT:1> ] <EOL> return <NUM_LIT:0.5> / np . pi / ( <NUM_LIT:1> + rho * rho - <NUM_LIT:2> * rho * np . cos ( w ) ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> def maxabs ( x , y ) : <EOL> return np . max ( np . abs ( x - y ) ) <EOL> nobs = <NUM_LIT:200> <EOL> ar = [ <NUM_LIT:1> , <NUM_LIT:0.0> ] <EOL> ma = [ <NUM_LIT:1> , <NUM_LIT:0.0> ] <EOL> ar2 = np . zeros ( nobs ) <EOL> ar2 [ : <NUM_LIT:2> ] = [ <NUM_LIT:1> , - <NUM_LIT> ] <EOL> uni = np . zeros ( nobs ) <EOL> uni [ <NUM_LIT:0> ] = <NUM_LIT:1.> <EOL> arcomb = np . convolve ( ar , ar2 , mode = '<STR_LIT>' ) <EOL> marep = signal . lfilter ( ma , arcomb , uni ) <EOL> print ( marep [ : <NUM_LIT:10> ] ) <EOL> mafr = fft . fft ( marep ) <EOL> rvs = np . random . normal ( size = nobs ) <EOL> datafr = fft . fft ( rvs ) <EOL> y = fft . ifft ( mafr * datafr ) <EOL> print ( np . corrcoef ( np . c_ [ y [ <NUM_LIT:2> : ] , y [ <NUM_LIT:1> : - <NUM_LIT:1> ] , y [ : - <NUM_LIT:2> ] ] , rowvar = <NUM_LIT:0> ) ) <EOL> arrep = signal . lfilter ( [ <NUM_LIT:1> ] , marep , uni ) <EOL> print ( arrep [ : <NUM_LIT:20> ] ) <EOL> arfr = fft . fft ( arrep ) <EOL> yfr = fft . fft ( y ) <EOL> x = fft . ifft ( arfr * yfr ) . real <EOL> print ( x [ : <NUM_LIT:5> ] ) <EOL> print ( rvs [ : <NUM_LIT:5> ] ) <EOL> print ( np . corrcoef ( np . c_ [ x [ <NUM_LIT:2> : ] , x [ <NUM_LIT:1> : - <NUM_LIT:1> ] , x [ : - <NUM_LIT:2> ] ] , rowvar = <NUM_LIT:0> ) ) <EOL> arcombp = np . zeros ( nobs ) <EOL> arcombp [ : len ( arcomb ) ] = arcomb <EOL> map_ = np . zeros ( nobs ) <EOL> map_ [ : len ( ma ) ] = ma <EOL> ar0fr = fft . fft ( arcombp ) <EOL> ma0fr = fft . fft ( map_ ) <EOL> y2 = fft . ifft ( ma0fr / ar0fr * datafr ) <EOL> print ( y2 [ : <NUM_LIT:10> ] ) <EOL> print ( y [ : <NUM_LIT:10> ] ) <EOL> print ( maxabs ( y , y2 ) ) <EOL> ar = [ <NUM_LIT:1> , - <NUM_LIT> ] <EOL> ma = [ <NUM_LIT:1> , <NUM_LIT> ] <EOL> arma1 = ArmaFft ( [ <NUM_LIT:1> , - <NUM_LIT:0.5> , <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT> , - <NUM_LIT> , <NUM_LIT> ] , [ <NUM_LIT:1> , <NUM_LIT> ] , nobs ) <EOL> nfreq = nobs <EOL> w = np . linspace ( <NUM_LIT:0> , np . pi , nfreq ) <EOL> w2 = np . linspace ( <NUM_LIT:0> , <NUM_LIT:2> * np . pi , nfreq ) <EOL> import matplotlib . pyplot as plt <EOL> plt . close ( '<STR_LIT:all>' ) <EOL> plt . figure ( ) <EOL> spd1 , w1 = arma1 . spd ( <NUM_LIT:2> ** <NUM_LIT:10> ) <EOL> print ( spd1 . shape ) <EOL> _ = plt . plot ( spd1 ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> plt . figure ( ) <EOL> spd2 , w2 = arma1 . spdshift ( <NUM_LIT:2> ** <NUM_LIT:10> ) <EOL> print ( spd2 . shape ) <EOL> _ = plt . plot ( w2 , spd2 ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> plt . figure ( ) <EOL> spd3 , w3 = arma1 . spddirect ( <NUM_LIT:2> ** <NUM_LIT:10> ) <EOL> print ( spd3 . shape ) <EOL> _ = plt . plot ( w3 , spd3 ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> plt . figure ( ) <EOL> spd3b = arma1 . _spddirect2 ( <NUM_LIT:2> ** <NUM_LIT:10> ) <EOL> print ( spd3b . shape ) <EOL> _ = plt . plot ( spd3b ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> plt . figure ( ) <EOL> spdr , wr = arma1 . spdroots ( w ) <EOL> print ( spdr . shape ) <EOL> plt . plot ( w , spdr ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> plt . figure ( ) <EOL> spdar1_ = spdar1 ( arma1 . ar , w ) <EOL> print ( spdar1_ . shape ) <EOL> _ = plt . plot ( w , spdar1_ ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> plt . figure ( ) <EOL> wper , spdper = arma1 . periodogram ( nfreq ) <EOL> print ( spdper . shape ) <EOL> _ = plt . plot ( w , spdper ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> startup = <NUM_LIT:1000> <EOL> rvs = arma1 . generate_sample ( startup + <NUM_LIT> ) [ startup : ] <EOL> import matplotlib . mlab as mlb <EOL> plt . figure ( ) <EOL> sdm , wm = mlb . psd ( x ) <EOL> print ( '<STR_LIT>' , sdm . shape ) <EOL> sdm = sdm . ravel ( ) <EOL> plt . plot ( wm , sdm ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> from nitime . algorithms import LD_AR_est <EOL> wnt , spdnt = LD_AR_est ( rvs , <NUM_LIT:10> , <NUM_LIT> ) <EOL> plt . figure ( ) <EOL> print ( '<STR_LIT>' , spdnt . shape ) <EOL> _ = plt . plot ( spdnt . ravel ( ) ) <EOL> print ( spdnt [ : <NUM_LIT:10> ] ) <EOL> plt . title ( '<STR_LIT>' ) <EOL> fig = plt . figure ( ) <EOL> arma1 . plot4 ( fig ) </s>
