Python numpy.fft.fft2() Examples

def wiener_dft(im: np.ndarray, sigma: float) -> np.ndarray:
    """
    Adaptive Wiener filter applied to the 2D FFT of the image
    :param im: multidimensional array
    :param sigma: estimated noise power
    :return: filtered version of input im
    """
    noise_var = sigma ** 2
    h, w = im.shape

    im_noise_fft = fft2(im)
    im_noise_fft_mag = np.abs(im_noise_fft / (h * w) ** .5)

    im_noise_fft_mag_noise = wiener_adaptive(im_noise_fft_mag, noise_var)

    zeros_y, zeros_x = np.nonzero(im_noise_fft_mag == 0)

    im_noise_fft_mag[zeros_y, zeros_x] = 1
    im_noise_fft_mag_noise[zeros_y, zeros_x] = 0

    im_noise_fft_filt = im_noise_fft * im_noise_fft_mag_noise / im_noise_fft_mag
    im_noise_filt = np.real(ifft2(im_noise_fft_filt))

    return im_noise_filt.astype(np.float32) 

def crosscorr_2d(k1: np.ndarray, k2: np.ndarray) -> np.ndarray:
    """
    PRNU 2D cross-correlation
    :param k1: 2D matrix of size (h1,w1)
    :param k2: 2D matrix of size (h2,w2)
    :return: 2D matrix of size (max(h1,h2),max(w1,w2))
    """
    assert (k1.ndim == 2)
    assert (k2.ndim == 2)

    max_height = max(k1.shape[0], k2.shape[0])
    max_width = max(k1.shape[1], k2.shape[1])

    k1 -= k1.flatten().mean()
    k2 -= k2.flatten().mean()

    k1 = np.pad(k1, [(0, max_height - k1.shape[0]), (0, max_width - k1.shape[1])], mode='constant', constant_values=0)
    k2 = np.pad(k2, [(0, max_height - k2.shape[0]), (0, max_width - k2.shape[1])], mode='constant', constant_values=0)

    k1_fft = fft2(k1, )
    k2_fft = fft2(np.rot90(k2, 2), )

    return np.real(ifft2(k1_fft * k2_fft)).astype(np.float32) 

def test_expand(): 
    X = torch.randn(2,2,4,4).cuda().double()
    zeros = torch.zeros(2,2,4,4).cuda().double()
    r1, r2 = cfft.rfft2(X)
    c1, c2 = cfft.fft2(X, zeros)
    assert np.allclose(cfft.expand(r1).cpu().numpy(), c1.cpu().numpy())
    assert np.allclose(cfft.expand(r2, imag=True).cpu().numpy(), c2.cpu().numpy())
    r1, r2 = cfft.rfft3(X)
    c1, c2 = cfft.fft3(X, zeros)
    assert np.allclose(cfft.expand(r1).cpu().numpy(), c1.cpu().numpy())
    assert np.allclose(cfft.expand(r2, imag=True).cpu().numpy(), c2.cpu().numpy())

    X = torch.randn(2,2,5,5).cuda().double()
    zeros = torch.zeros(2,2,5,5).cuda().double()
    r1, r2 = cfft.rfft3(X)
    c1, c2 = cfft.fft3(X, zeros)
    assert np.allclose(cfft.expand(r1, odd=True).cpu().numpy(), c1.cpu().numpy())
    assert np.allclose(cfft.expand(r2, imag=True, odd=True).cpu().numpy(), c2.cpu().numpy()) 

def translation(im0, im1):
    
    # Convert images to grayscale
    im0 = cv2.cvtColor(im0,cv2.COLOR_BGR2GRAY)
    im1 = cv2.cvtColor(im1,cv2.COLOR_BGR2GRAY)
    
    shape = im0.shape
    f0 = fft2(im0)
    f1 = fft2(im1)
    ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1))))
    t0, t1 = np.unravel_index(np.argmax(ir), shape)
    if t0 > shape[0] // 2:
        t0 -= shape[0]
    if t1 > shape[1] // 2:
        t1 -= shape[1]
    return [t0, t1] 

def _applyImage(self,data):
        img = data['img']
        img_g = img
        if (np.ndim(img) == 3):
            img_g = (img[:, :, 0] + img[:, :, 1] + img[:, :, 2]) / 3
        return np.log10(np.abs(fftshift(fft2(img_g)))) * 255 

def linear_step(self):
        """Make one linear step dt forward in time"""
        # kinetic
        self.wf[:] = fft2(ifft2(self.wf) * self.T)

        # potential
        self.wf *= self.V 

def analyze(f, axes=(0, 1)):
        """
        Compute the Fourier Transform of the discretely sampled function f : T^2 -> C.

        Let f : T^2 -> C be a band-limited function on the torus.
        The samples f(theta_k, phi_l) correspond to points on a regular grid on the circle,
        as returned by spaces.T1.linspace:
        theta_k = phi_k = 2 pi k / N
        for k = 0, ..., N - 1 and l = 0, ..., N - 1

        This function computes
        \hat{f}_n = (1/N) \sum_{k=0}^{N-1} f(theta_k) e^{-i n theta_k}
        which, if f has band-limit less than N, is equal to:
        \hat{f}_n = \int_0^{2pi} f(theta) e^{-i n theta} dtheta / 2pi,
                  = <f(theta), e^{i n theta}>
        where dtheta / 2pi is the normalized Haar measure on T^1, and < , > denotes the inner product on Hilbert space,
        with respect to which this transform is unitary.

        The range of frequencies n is -floor(N/2) <= n <= ceil(N/2) - 1

        :param f:
        :param axis:
        :return:
        """
        # The numpy FFT returns coefficients in a different order than we want them,
        # and using a different normalization.
        f_hat = fft2(f, axes=axes)
        f_hat = fftshift(f_hat, axes=axes)
        size = np.prod([f.shape[ax] for ax in axes])
        return f_hat / size 

def fftd(I, dims=None):

    # Compute fft
    if dims is None:
        X = fftn(I)
    elif dims == 2:
        X = fft2(I, axes=(0, 1))
    else:
        X = fftn(I, axes=tuple(range(dims)))

    return X 

def Convolve(image1, image2, MinPad=True, pad=True):
    """
    Convolves image1 with image2.

    :param image1: 2D image array
    :param image2: 2D image array
    :param MinPad: whether to use minimal padding
    :param pad: whether to pad the array
    """
    #The size of the images:
    r1, c1 = image1.shape
    r2, c2 = image2.shape

    if MinPad:
        r = r1 + r2
        c = c1 + c2
    else:
        r = 2*max(r1,r2)
        c = 2*max(c1,c2)
    
    #or in power of two
    if pad:
        pr2 = int(m.log(r)/m.log(2.) + 1.)
        pc2 = int(m.log(c)/m.log(2.) + 1.)
        rOrig = r
        cOrig = c
        r = 2**pr2
        c = 2**pc2
    
    fftimage = fft2(image1, s=(r,c))*fft2(image2[::-1,::-1],s=(r,c))

    if pad:
        return (ifft2(fftimage))[:rOrig,:cOrig].real
    else:
        return (ifft2(fftimage)).real 

def fft2c(x, norm='ortho', axes=None):
    '''
    Centered fft
    Note: fft2 applies fft to last 2 axes by default
    :param x: 2D onwards. e.g: if its 3d, x.shape = (n,row,col). 4d:x.shape = (n,slice,row,col)
    :return:
    '''
    # axes = (len(x.shape)-2, len(x.shape)-1)  # get last 2 axes
    if not axes:
        axes = (-2, -1)  # get last 2 axes
    res = fftshift(fft2(ifftshift(x, axes=axes), norm=norm, axes=axes), axes=axes)
    return res 

def ifft2c(x, norm='ortho', axes=None):
    '''
    Centered ifft
    Note: fft2 applies fft to last 2 axes by default
    :param x: 2D onwards. e.g: if its 3d, x.shape = (n,row,col). 4d:x.shape = (n,slice,row,col)
    :return:
    '''
    if not axes:
        axes = (-2, -1)  # get last 2 axes
    res = fftshift(ifft2(ifftshift(x, axes=axes), norm=norm, axes=axes), axes=axes)
    return res 

def fft_convolve2d(x,y):
    """ 2D convolution, using FFT"""
    fr = fft2(x)
    fr2 = fft2(y)
    cc = np.real(ifft2(fr*fr2))
    cc = fftshift(cc)
    return cc 

def _checkSanity(self):
    '''
    Check that screen parameters are consistent.
    '''

    # sanity check: is screen large enough?
    assert np.ceil(self.nx*self.ips)+2 < np.min(self.nphi), \
          "screen is not large enough: {0} > {1}".\
          format(int(np.ceil(self.ips*self.nx)+2),np.min(self.nphi))

    # sanity check: is image square?
    assert self.model.shape[-1] == self.model.shape[-2], \
        'source image must be square'

    # sanity check: integer number of screen pixels / image pixel?
    assert self.ips % 1 == 0, 'image/screen pixels should be an integer'

    # is inner turbulence scale larger than r0?
    #assert 1./self.qmax > self.r0/self.screen_dx, 'turbulence inner scale < r0'
    assert self.r_inner > 1., 'turbulence inner scale < r0'

    # check image smoothness
    V = fft2(self.isrc)
    freq = fftfreq(self.nx,d=self.dx*radians(1.)/(3600*1e6))
    u = dot(transpose([np.ones(self.nx)]),[freq])
    v = dot(transpose([freq]),[ones(self.nx)])
    try:
        if max(abs(V[sqrt(u*u+v*v) > (1.+self.m)*self.r_inner*self.r0/self.wavelength])) / self.isrc.sum() > 0.01:
         self.logger.warning('image is not smooth enough:  {0:g} > 0.01'.format(max(abs(V[sqrt(u*u+v*v) > (1.+self.m)*self.r_inner*self.r0/self.wavelength])) / self.isrc.sum()))
    except ValueError:
        self.logger.warning('r_inner is too large to test smoothness')

    # is screen pixel smaller than inner turbulence scale?
    #assert 1./self.qmax >= 1, 'screen pixel > turbulence inner scale'
    assert self.r_inner*self.r0/self.screen_dx >= 1, 'screen pixel > turbulence inner scale'

    if (self.rf*self.rf/self.r0/(self.ips*self.screen_dx) < 3):
      self.logger.warning('WARNING: image resolution is approaching Refractive scale') 

def ft2(g, delta):
    return fftshift(fft2(fftshift(g))) * (delta ** 2) 

def crossCorrelate(self,a,b):
        """Computes the cross correlation of f* x g where f* is the conjugate of a function f and x represents the cross correlation"""
        if a.shape!=b.shape:
            print "WARNING in mapRecon.crossCorrelate - different shapes for convolution: %s %s"%(str(a.shape),str(b.shape))
        return fft.fft2(na.conjugate(fft.ifft2(a))*fft.ifft2(b)) 

def convolve(self,a,b):
        if a.shape!=b.shape:
            print "WARNING in mapRecon.convolve - different shapes for convolution: %s %s"%(str(a.shape),str(b.shape))
        return fft.ifft2(fft.fft2(a)*fft.fft2(b)) 

def load_ref(self,ref):
        ref1=self.initial_img(ref)
        if self.cor_mode == 1:
            ref1=numpy.conjugate(fft.fft2(ref1*self.wdw))
        return ref1 

def computeMask(subflag):
    """Compute the mask matrix... the input is array of 1 or 0 depending on whether subap is used."""
    subflag=numpy.array(subflag)
    subaps=subflag.flatten()
    dim=subaps.shape[0]
    M=numpy.zeros((dim*2,dim*2),numpy.float64)
    for i in range(dim):
        if subaps[i]:
            M[i,i]=1
            M[i+dim,i+dim]=1
    Mhat=fft.fft2(M)
    return M,Mhat 

def computeChatFromPokeMatrix2(pokemx,sigma,turbStrength,kappa):
    """Similar to computeChatFromPokeMatrix, except the pokemx here
    just has shape ncents, nmodes, removing the need for nlayers and
    nGS dimensions, this being swallowed into the pokemx.

    This effectively relies on the result that we are trying to
    compute Chat from GT C-1 G + Cphi-1.  And notes that GT G is the
    same as F G^* G^ F-1, so what we compute here is effectively G^*
    G^.

    After computing this, one can then make use of sparsity and get
    rid of some of it (use compactToBanded method).  Chat can probably
    be treated as banded (you should verify this), though may need
    some sparse matrix reordering first.

    Note, the pokemx is likely to be block diagonal, because poking
    one mirror does not affect the phase on other mirrors.

    We assume here that the noise on each subap is identical, ie equal
    to sigma, so that the noise covariance matrix is sigma^2 I with I
    being the identity matrix.

    Not really quite sure how this would plan out if using more than 1
    layer - maybe need to FFT the blocks of the pokemx, not the entire
    pokemx (todo).
    """
    turbStr=1./turbStrength**2
    sigmam2=1./(sigma*sigma)
    fftpokemx=numpy.fft.fft2(pokemx)
    fftpokemxCT=numpy.conjugate(numpy.transpose(fftpokemx))
    Chat=sigmam2*quick.dot(fftpokemxCT,fftpokemx)
    modesPerLayer=float(Chat.shape[0])/turbStrength.shape[0]
    if numpy.ceil(modesPerLayer)!=numpy.floor(modesPerLayer):
        print "pcg: warning - modes per layer not an integer."
    modesPerLayer=int(modesPerLayer)
    for i in range(turbStrength.shape[0]):#for each turbulent layer...
        p=i*modesPerLayer
        q=p+modesPerLayer
        Chat[p:q,p:q]+=turbStr[i]*kappa
    return Chat 

def define_kernel_pixels_fft(ref,target,rad,INNER_RADIUS=7,threshold=3.0):
    from numpy.fft import fft2, ifft2
    from astropy.stats import mad_std
    nx, ny = ref.image.shape
    x = np.concatenate((np.arange(nx/2),np.arange(-nx/2,0)))
    y = np.concatenate((np.arange(ny/2),np.arange(-ny/2,0)))
    fr = fft2(ref.image)
    ft = fft2(target.image)
    fk = ft/fr
    k = ifft2(fk)
    nk = k/k.max()
    std_nk = mad_std(nk)
    kp = np.where(np.abs(nk)>threshold)
    print 'kernel radius',rad
    crad = int(np.ceil(rad))
    rad2 = rad*rad
    inner_rad2 = INNER_RADIUS*INNER_RADIUS
    kCount = 1
    for p in range(kp[0].shape[0]):
        i = x[kp[0][p]]
        j = y[kp[1][p]]
        r2 = i*i + j*j
        if (r2 < rad2) and ((i,j) != (0,0)):
            if (r2 < inner_rad2):
                kCount += 1
            else:
                if (i/3 == i/3.0) and (j/3 == j/3.0):
                    kCount += 1
    kInd = np.zeros([kCount,2],dtype=np.int32)
    kExtended = np.zeros(kCount,dtype=np.int32)
    kInd[0] = [0,0]
    k = 1
    for p in range(kp[0].shape[0]):
        i = x[kp[0][p]]
        j = y[kp[1][p]]
        r2 = i*i + j*j
        if (r2 < rad2) and ((i,j) != (0,0)):
            if (r2 < inner_rad2):
                kInd[k] = [i,j]
                k += 1
            else:
                if (i/3 == i/3.0) and (j/3 == j/3.0):
                    kInd[k] = [i,j]
                    kExtended[k] = 1
                    k += 1
    n_extend = np.sum(kExtended)
    print kCount-n_extend,'modified delta basis functions'
    print n_extend,'extended basis functions'
    for k in range(kCount):
        print kInd[k]
    return kInd, kExtended 

def nextIter(self):
        """Perform a PCG iteration"""
        self.p=self.p.real
        self.xprev=self.x
        if self.options["diagonaliseA"]==1:#A has been diagonalised.
            self.q=self.diagonalA*self.p
        elif self.options["diagonaliseA"]==0.5:#sparse...
            self.q=numpy.zeros(self.p.shape,self.p.dtype)
            for key in self.sparseA.keys():
                if key==0:
                    self.q+=self.sparseA[key]*self.p
                elif key<0:
                    self.q[-key:,]+=self.sparseA[key]*self.p[:key]
                else:
                    self.q[:-key]+=self.sparseA[key]*self.p[key:,]
            
        else:#full implementation
            print "big mmx multiply for q=Ap"
            self.q=quick.dot(self.A,self.p)#dominant cost
        qp=quick.dot(self.q,self.p)#complex
        self.rz=quick.dot(self.r,self.z)
        if qp!=0:
            #print type(self.r),type(self.z),type(qp)#agbhome
            #print self.r.shape,self.z.shape,self.r.dtype.char,self.z.dtype.char
            #tmp=quick.dot(self.r,self.z)
            #print type(tmp)
            self.alpha=(self.rz/qp).real#complex
        else:
            self.alpha=0.#if q or p are zero, implies z must be, so r must be.
        self.x=(self.x+self.alpha*self.p).real#complex, but seems that x.imag tends to zero???  Can we ignore it right from the start?
        if self.options["removeHiFreqX"]:
            lsize=self.nact*self.nact
            for i in range(self.nLayers):
                fftx=numpy.fft.fft2(numpy.reshape(self.x[i*lsize:(i+1)*lsize],(self.nact,self.nact)))
            f=(self.nact-1)/2
            t=(self.nact+1)/2+1
            fftx[f:t,f:t]=numpy.average(fftx[f:t,f:t].flat)
            self.x[i*lsize:(i+1)*lsize]=numpy.fft.ifft2(fftx).ravel().real
            
        
        self.r-=self.alpha*self.q
        #self.z=quick.dot(self.Cm1,self.r)#dominant cost
        #Prev line actually computed in FFT space (next line...).
        self.z=self.computeinvChatr(self.invChatReordered,self.r,mode="diag").real
        if self.rz!=0:
            self.beta=(quick.dot(self.r,self.z)/self.rz).real
        else:
            self.beta=0.#if rz is zero, so is new r or z...
        #print "alpha,beta",self.alpha,self.beta
        self.p=(self.z+self.beta*self.p).real
        if self.options["removeHiFreqP"]:
            lsize=self.nact*self.nact
            for i in range(self.nLayers):
                fftp=numpy.fft.fft2(numpy.reshape(self.p[i*lsize:(i+1)*lsize],(self.nact,self.nact)))
            f=(self.nact-1)/2
            t=(self.nact+1)/2+1
            fftp[f:t,f:t]=numpy.average(fftp[f:t,f:t].flat)
            self.p[i*lsize:(i+1)*lsize]=numpy.fft.ifft2(fftp).ravel().real