python类fftshift()的实例源码

def create_kgrid(nx, ny, nz, lx=2*pi, ly=2*pi, lz=2*pi):
   """
   Create a 3D k grid for Fourier space calculations
   """

   print lx, ly, lz

   kx = nf.fftshift(nf.fftfreq(nx))*nx*2*pi/lx
   ky = nf.fftshift(nf.fftfreq(ny))*ny*2*pi/ly
   kz = nf.fftshift(nf.fftfreq(nz))*nz*2*pi/lz

   mg = np.meshgrid(kx,ky,kz)

   km = np.sqrt(np.sum((m**2 for m in mg)))

   return kx[:,nna,nna], ky[nna,:,nna], kz[nna,nna,:], km

#==================================================

def correlate(self, imgfft):
        #Very much related to the convolution theorem, the cross-correlation
        #theorem states that the Fourier transform of the cross-correlation of
        #two functions is equal to the product of the individual Fourier
        #transforms, where one of them has been complex conjugated:  


        if self.imgfft is not 0 or imgfft.imgfft is not 0:
            imgcj = np.conjugate(self.imgfft)
            imgft = imgfft.imgfft

            prod = deepcopy(imgcj)
            for x in range(imgcj.shape[0]):
                for y in range(imgcj.shape[0]):
                    prod[x][y] = imgcj[x][y] * imgft[x][y]

            cc = Corr( np.real(fft.ifft2(fft.fftshift(prod)))) # real image of the correlation

            # adjust to center
            cc.data = np.roll(cc.data, int(cc.data.shape[0] / 2), axis = 0)
            cc.data = np.roll(cc.data, int(cc.data.shape[1] / 2), axis = 1)
        else:
            raise FFTnotInit()
        return cc

def kfilter(ar, kf, lx=2*pi, ly=2*pi, lz=2*pi):
   """
   Function to filter a 3D array by zeroing out larger k values
   """

   while len(ar.shape) < 3:
      ar = ar.reshape(ar.shape + (1,))

   #  COMPUTE THE ARRAY SIZE 
   kx, ky, kz, km = create_kgrid(*ar.shape, lx=lx, ly=ly, lz=lz)

   #  FOURIER TRANSFORM THE ARRAY 
   far = nf.fftshift(nf.fftn(ar))

   #  SET VALUES ABOVE kf AS 0+0i
   far = (np.sign(km - kf) - 1.)/(-2.)*far

   #  BACK TRANSFORM TO REAL SPACE
   arf = np.real(nf.ifftn(nf.ifftshift(far)))
   return arf

#==================================================

def kdiv(arx,ary,arz,kf=None,lx=2*pi,ly=2*pi,lz=2*pi):
   """
   Function to compute divergence in Fourier space
   """

   #  COMPUTE THE ARRAY SIZE
   kx, ky, kz, km = create_kgrid(*np.shape(arx), lx=lx, ly=ly, lz=lz)

   #  FOURIER TRANSFORM THE ARRAY
   far = [nf.fftshift(nf.fftn(a)) for a in (arx, ary, arz)]

   # COMPUTE div=i*(kx*ax+ky*ay+kz*az)
   mg = np.meshgrid(kx,ky,kz)

   divarf = 1.j*reduce(operator.add, [a*b for a,b in zip(mg, far)])

   #  SET VALUES ABOVE kf AS 0+0i if kf non zero
   if kf is not None:
      divarf = (np.sign(km - kf) - 1.)/(-2.)*divarf

   divar = np.real(nf.ifftn(nf.ifftshift(divarf)))

   return divar

def kurl(arx,ary,arz, kf, lx=2*np.pi, ly=2*np.pi, lz=2*np.pi):
   kx, ky, kz, km = create_kgrid(*arx.shape, lx=lx, ly=ly, lz=lz)

   #  FOURIER TRANSFORM THE ARRAY
   farx = nf.fftshift(nf.fftn(arx))
   fary = nf.fftshift(nf.fftn(ary))
   farz = nf.fftshift(nf.fftn(arz))

   #  SET VALUES ABOVE kf AS 0+0i
   farx = (np.sign(km - kf) - 1.)/(-2.)*farx
   fary = (np.sign(km - kf) - 1.)/(-2.)*fary
   farz = (np.sign(km - kf) - 1.)/(-2.)*farz

   #  COMPUTE VORTICITY
   axf = eye*(ky*farz-kz*fary)
   ayf = eye*(kz*farx-kx*farz)
   azf = eye*(kx*fary-ky*farx)

   #  BACK TRANSFORM TO REAL SPACE
   wx  = np.real(nf.ifftn(nf.ifftshift(axf)))
   wy  = np.real(nf.ifftn(nf.ifftshift(ayf)))
   wz  = np.real(nf.ifftn(nf.ifftshift(azf)))
   return wx,wy,wz

def crystal():
    # look at crystal and their ffts
    crystal3D = np.zeros((201,201,201), dtype= np.int32)
    crystal3D_fourier = np.zeros((201,201,201), dtype= np.complex64)

    dx = 1

    for row in range(60,140,dx):
        for col in range(80,120,dx):
            for time in range(90,110,dx):
                crystal3D[row,col,time] = 1

    crystal3D_fourier = fft.fftshift(fft.fftn(crystal3D))
    #del crystal3D
    diffPattern3D = (abs(crystal3D_fourier)**2)
    del crystal3D_fourier
    return diffPattern3D

def skew_image():
    image_skewed = np.zeros((image.shape[0], ceil((image.shape[1]/np.cos(theta)))+ 50  ))

    for i in range(0,image.shape[0]):
        for j in range(0,image.shape[1]):
            xs = ceil(j / (1 + np.tan(theta)**2) + i*np.tan(theta)/ ( 1 + np.tan(theta)**2)  )
            #np.disp(xs)
            ys = i
            image_skewed[ys,xs] = image[i,j] 

    fft_image_skewed = fft.fftshift(fft.fft2(image_skewed))
    return image_skewed, fft_image_skewed
# image in reciprocal space
#fft_image = fft.fftshift(fft.fft2(image))


# Create crystal: skewed and unskewed. 2D and 3D.
# create FFTs of these. Filter out one peak.
################################################

def skewed_crystal_diffPatterns(crystal, theta):    
    dx_filter = 1            
    crystal_skewed = np.zeros((crystal.shape[0], np.floor((crystal.shape[1]/np.cos(theta))) + 50  ))
    crystal_filter2D_skewed =  np.zeros(( crystal_skewed.shape[0], crystal_skewed.shape[1] ))
    # simulate skewed diffraction patterns
    theta = 10*np.pi/180

    for i in range(0,crystal.shape[0]-6):
        for j in range(0,crystal.shape[1]-6):
            xs = ceil(j / (1 + np.tan(theta)**2) + i*np.tan(theta)/ ( 1 + np.tan(theta)**2)  )
            ys = i
            crystal_skewed[ys,xs] = crystal[i,j] 

    for row in range(80,150,dx_filter):
        for col in range(80,160,dx_filter):            
            crystal_filter2D_skewed[row,col] = 1

    diffPattern = abs( fft.fftshift(fft.fft2(crystal_skewed)))**2
    return (diffPattern, crystal_filter2D_skewed)

#diffPattern_skewed, crystal_filter2D_skewed = skewed_crystal_diffPatterns(crystal, theta )


# call shrinkwrap for filtered skewed crystal
# yel = shrinkwrap(diffPattern_skewed*crystal_filter2D_skewed)

def cross_correlation_using_fft(self, x, y):
    f1 = fft(x)
    f2 = fft(np.flipud(y))
        cc = np.real(ifft(f1 * f2))
    return fftshift(cc)



    # clean up
   # def close(self):
      # close serial
   #   self.ser.flush()
   #   self.ser.close() 



# Main

def ift(self):
      return MyImage(np.real(fft.ifft2(fft.fftshift(self.ft))))

def ft(self):
        self.imgfft = fft.fftshift(fft.fft2(self.img.data))

def ift(self):
        self.imgifft = MyImage(np.real(fft.ifft2(fft.fftshift(self.imgfft))))

def test_definition(self):
        x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
        y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
        assert_array_almost_equal(fft.fftshift(x), y)
        assert_array_almost_equal(fft.ifftshift(y), x)
        x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
        y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
        assert_array_almost_equal(fft.fftshift(x), y)
        assert_array_almost_equal(fft.ifftshift(y), x)

def test_inverse(self):
        for n in [1, 4, 9, 100, 211]:
            x = np.random.random((n,))
            assert_array_almost_equal(fft.ifftshift(fft.fftshift(x)), x)

def test_axes_keyword(self):
        freqs = [[0, 1, 2], [3, 4, -4], [-3, -2, -1]]
        shifted = [[-1, -3, -2], [2, 0, 1], [-4, 3, 4]]
        assert_array_almost_equal(fft.fftshift(freqs, axes=(0, 1)), shifted)
        assert_array_almost_equal(fft.fftshift(freqs, axes=0),
                fft.fftshift(freqs, axes=(0,)))
        assert_array_almost_equal(fft.ifftshift(shifted, axes=(0, 1)), freqs)
        assert_array_almost_equal(fft.ifftshift(shifted, axes=0),
                fft.ifftshift(shifted, axes=(0,)))

def cross_corr(img1,img2,mask=None):
    '''Compute the autocorrelation of two images.
        Right now does not take mask into account.
        todo: take mask into account (requires extra calculations)
        input: 
            img1: first image
            img2: second image
            mask: a mask array
        output:
            the autocorrelation of the two images (same shape as the correlated images)

    '''
    #if(mask is not None):
    #   img1 *= mask
    #  img2 *= mask

    #img1_mean = np.mean( img1.flat )
    #img2_mean = np.mean( img2.flat )

    # imgc = fftshift( ifft2(     
    #        fft2(img1/img1_mean -1.0 )*np.conj(fft2( img2/img2_mean -1.0 ))).real )

    #imgc = fftshift( ifft2(     
    #        fft2(  img1/img1_mean  )*np.conj(fft2(  img2/img2_mean   ))).real )

    imgc = fftshift( ifft2(     
            fft2(  img1  )*np.conj(fft2(  img2  ))).real )

    #imgc /= (img1.shape[0]*img1.shape[1])**2
    if(mask is not None):
        maskc = cross_corr(mask,mask)        
        imgc /= np.maximum( 1, maskc )


    return imgc

def test_definition(self):
        x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
        y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
        assert_array_almost_equal(fft.fftshift(x), y)
        assert_array_almost_equal(fft.ifftshift(y), x)
        x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
        y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
        assert_array_almost_equal(fft.fftshift(x), y)
        assert_array_almost_equal(fft.ifftshift(y), x)

def test_inverse(self):
        for n in [1, 4, 9, 100, 211]:
            x = np.random.random((n,))
            assert_array_almost_equal(fft.ifftshift(fft.fftshift(x)), x)

def test_axes_keyword(self):
        freqs = [[0, 1, 2], [3, 4, -4], [-3, -2, -1]]
        shifted = [[-1, -3, -2], [2, 0, 1], [-4, 3, 4]]
        assert_array_almost_equal(fft.fftshift(freqs, axes=(0, 1)), shifted)
        assert_array_almost_equal(fft.fftshift(freqs, axes=0),
                fft.fftshift(freqs, axes=(0,)))
        assert_array_almost_equal(fft.ifftshift(shifted, axes=(0, 1)), freqs)
        assert_array_almost_equal(fft.ifftshift(shifted, axes=0),
                fft.ifftshift(shifted, axes=(0,)))

def test_definition(self):
        x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
        y = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
        assert_array_almost_equal(fft.fftshift(x), y)
        assert_array_almost_equal(fft.ifftshift(y), x)
        x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
        y = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
        assert_array_almost_equal(fft.fftshift(x), y)
        assert_array_almost_equal(fft.ifftshift(y), x)

def test_inverse(self):
        for n in [1, 4, 9, 100, 211]:
            x = np.random.random((n,))
            assert_array_almost_equal(fft.ifftshift(fft.fftshift(x)), x)

def test_axes_keyword(self):
        freqs = [[0, 1, 2], [3, 4, -4], [-3, -2, -1]]
        shifted = [[-1, -3, -2], [2, 0, 1], [-4, 3, 4]]
        assert_array_almost_equal(fft.fftshift(freqs, axes=(0, 1)), shifted)
        assert_array_almost_equal(fft.fftshift(freqs, axes=0),
                fft.fftshift(freqs, axes=(0,)))
        assert_array_almost_equal(fft.ifftshift(shifted, axes=(0, 1)), freqs)
        assert_array_almost_equal(fft.ifftshift(shifted, axes=0),
                fft.ifftshift(shifted, axes=(0,)))

def xcorr(imageA, imageB):
    FimageA = _fft.fft2(imageA)
    CFimageB = _np.conj(_fft.fft2(imageB))
    return _fft.fftshift(_np.real(_fft.ifft2((FimageA * CFimageB)))) / _np.sqrt(imageA.size)

def test_CenteredFFT(backend, M, N, K, B ):
    from numpy.fft import fftshift, ifftshift, fftn, ifftn

    b = backend()
    A = b.FFTc( (M,N,K), dtype=np.dtype('complex64') )

    # forward
    ax = (0,1,2)
    x = b.rand_array( (M*N*K,B) )
    y = b.rand_array( (M*N*K,B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.eval(y, x)

    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    y_exp = fftshift( fftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

    # adjoint
    x = b.rand_array( (M*N*K,B) )
    y = b.rand_array( (M*N*K,B) )
    x_h = x.to_host().reshape( (M,N,K,B), order='F' )

    A.H.eval(y, x)

    y_act = y.to_host().reshape( (M,N,K,B), order='F' )
    y_exp = fftshift( ifftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
    npt.assert_allclose(y_act, y_exp, rtol=1e-2)

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkwl(img):
    """Calculate wavelength  given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkwl(img):
    """Calculate wavelength  given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire