#NEWLINE# def fftconvolve( in1 , in2 , mode = 'full' , axis = None ) : #NEWLINE# #INDENT# ' Convolve two N-dimensional arrays using FFT. See convolve.\n\n    This is a fix of scipy.signal.fftconvolve, adding an axis argument and\n    importing locally the stuff only needed for this function\n\n    ' #NEWLINE# s1 = np . array ( in1 . shape ) #NEWLINE# s2 = np . array ( in2 . shape ) #NEWLINE# complex_result = ( np . issubdtype ( in1 . dtype , np . complex ) or np . issubdtype ( in2 . dtype , np . complex ) ) #NEWLINE# if ( axis is None ) : #NEWLINE# #INDENT# size = ( ( s1 + s2 ) - 1 ) #NEWLINE# fslice = tuple ( [ slice ( 0 , int ( sz ) ) for sz in size ] ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# equal_shapes = ( s1 == s2 ) #NEWLINE# equal_shapes [ axis ] = True #NEWLINE# assert equal_shapes . all ( ) , 'Shape mismatch on non-convolving axes' #NEWLINE# size = ( ( s1 [ axis ] + s2 [ axis ] ) - 1 ) #NEWLINE# fslice = [ slice ( l ) for l in s1 ] #NEWLINE# fslice [ axis ] = slice ( 0 , int ( size ) ) #NEWLINE# fslice = tuple ( fslice ) #NEWLINE# #UNINDENT# fsize = ( 2 ** int ( np . ceil ( np . log2 ( size ) ) ) ) #NEWLINE# if ( axis is None ) : #NEWLINE# #INDENT# IN1 = fftpack . fftn ( in1 , fsize ) #NEWLINE# IN1 *= fftpack . fftn ( in2 , fsize ) #NEWLINE# ret = fftpack . ifftn ( IN1 ) [ fslice ] . copy ( ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# IN1 = fftpack . fft ( in1 , fsize , axis = axis ) #NEWLINE# IN1 *= fftpack . fft ( in2 , fsize , axis = axis ) #NEWLINE# ret = fftpack . ifft ( IN1 , axis = axis ) [ fslice ] . copy ( ) #NEWLINE# #UNINDENT# del IN1 #NEWLINE# if ( not complex_result ) : #NEWLINE# #INDENT# ret = ret . real #NEWLINE# #UNINDENT# if ( mode == 'full' ) : #NEWLINE# #INDENT# return ret #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# if ( mode == 'same' ) : #NEWLINE# #INDENT# if ( np . product ( s1 , axis = 0 ) > np . product ( s2 , axis = 0 ) ) : #NEWLINE# #INDENT# osize = s1 #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# osize = s2 #NEWLINE# #UNINDENT# return signaltools . _centered ( ret , s1 ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# if ( mode == 'valid' ) : #NEWLINE# #INDENT# return signaltools . _centered ( ret , ( abs ( ( s2 - s1 ) ) + 1 ) )