Python numpy.fft.fftshift() Examples

def ptf(self):
		"""
		Phase transfer function
		"""
		PSF = self.__psfcaculator__()
		PTF = __fftshift__(__fft2__(PSF))
		PTF = __np__.angle(PTF)
		b = 400
		R = (200)**2
		for i in range(b):
			for j in range(b):
				if (i-b/2)**2+(j-b/2)**2>R:
					PTF[i][j] = 0
		__plt__.imshow(abs(PTF),cmap=__cm__.rainbow)
		__plt__.colorbar()
		__plt__.show()
		return 0 

def average(dj, mask_count_threshold):
    vol_sum = None
    mask_sum = None
    for d in dj:
        v = IF.read_mrc_vol(d['subtomogram'])
        if (not N.all(N.isfinite(v))):
            raise Exception('error loading', d['subtomogram'])
        vm = IF.read_mrc_vol(d['mask'])
        v_r = GR.rotate_pad_mean(v, angle=d['angle'], loc_r=d['loc'])
        assert N.all(N.isfinite(v_r))
        vm_r = GR.rotate_mask(vm, angle=d['angle'])
        assert N.all(N.isfinite(vm_r))
        if (vol_sum is None):
            vol_sum = N.zeros(v_r.shape, dtype=N.float64, order='F')
        vol_sum += v_r
        if (mask_sum is None):
            mask_sum = N.zeros(vm_r.shape, dtype=N.float64, order='F')
        mask_sum += vm_r
    ind = (mask_sum >= mask_count_threshold)
    vol_sum_fft = NF.fftshift(NF.fftn(vol_sum))
    avg = N.zeros(vol_sum_fft.shape, dtype=N.complex)
    avg[ind] = (vol_sum_fft[ind] / mask_sum[ind])
    avg = N.real(NF.ifftn(NF.ifftshift(avg)))
    return {'v': avg, 'm': (mask_sum / len(dj)), } 

def impute_aligned_vols(t, v, vm, normalize=None):
    assert (normalize is not None)
    if normalize:
        v = ((v - v.mean()) / v.std())
        if (t is not None):
            t = ((t - t.mean()) / t.std())
    if (t is None):
        return v
    t_f = NF.fftshift(NF.fftn(t))
    v_f = NF.fftshift(NF.fftn(v))
    v_f[(vm == 0)] = t_f[(vm == 0)]
    v_f_if = N.real(NF.ifftn(NF.ifftshift(v_f)))
    if normalize:
        v_f_if = ((v_f_if - v_f_if.mean()) / v_f_if.std())
    if N.all(N.isfinite(v_f_if)):
        return v_f_if
    else:
        print('warning: imputation failed')
        return v 

def ift2(G, delta_f, FFT=None):
    """
    Wrapper for inverse fourier transform

    Parameters:
        G: data to transform
        delta_f: pixel seperation
        FFT (FFT object, optional): An accelerated FFT object
    """

    N = G.shape[0]

    if FFT:
        g = numpy.fft.fftshift(FFT(numpy.fft.fftshift(G))) * (N * delta_f) ** 2
    else:
        g = fft.ifftshift(fft.ifft2(fft.fftshift(G))) * (N * delta_f) ** 2

    return g 

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 __init__(self, images, masks, band_width_radius=1.0):

        im_f = {}
        for k in images:
            im = images[k]
            im = fftshift(fftn(im))
            im_f[k] = im

        self.im_f = im_f        # fft transformed images
        self.ms = masks         # masks
        self.ks = set()
        self.img_siz = im_f[k].shape
        self.set_fft_mid_co()
        self.set_rad()
        self.band_width_radius = band_width_radius 

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_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 particle_xcorr(ptcl, refmap_ft):
    r = util.euler2rot(*np.deg2rad(ptcl[star.Relion.ANGLES]))
    proj = vop.interpolate_slice_numba(refmap_ft, r)
    c = ctf.eval_ctf(s / apix, a, def1[i], def2[i], angast[i], phase[i],
                     kv[i], ac[i], cs[i], bf=0, lp=2 * apix)
    pshift = np.exp(-2 * np.pi * 1j * (-xshift[i] * sx + -yshift * sy))
    proj_ctf = proj * pshift * c

    with mrc.ZSliceReader(ptcl[star.Relion.IMAGE_NAME]) as f:
        exp_image_fft = rfft2(fftshift(f.read(i)))

    xcor_fft = exp_image_fft * proj_ctf
    xcor = fftshift(irfft2(xcor_fft))
    return xcor 

def subtract_outer(p1r, ptcl, submap_ft, refmap_ft, sx, sy, s, a, apix, coefs_method, r, nr, **kwargs):
    log = logging.getLogger('root')
    log.debug("%d@%s Exp %f +/- %f" % (ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX], ptcl[star.UCSF.IMAGE_ORIGINAL_PATH], np.mean(p1r), np.std(p1r)))
    ft = getattr(tls, 'ft', None)
    if ft is None:
        ft = rfft2(fftshift(p1r.copy()), threads=kwargs["fftthreads"],
                   planner_effort="FFTW_ESTIMATE",
                   overwrite_input=False,
                   auto_align_input=True,
                   auto_contiguous=True)
        tls.ft = ft
    if coefs_method >= 1:
        p1 = ft(p1r.copy(), np.zeros(ft.output_shape, dtype=ft.output_dtype)).copy()
    else:
        p1 = np.empty(ft.output_shape, ft.output_dtype)

    p1s = subtract(p1, submap_ft, refmap_ft, sx, sy, s, a, apix,
                   ptcl[star.Relion.DEFOCUSU], ptcl[star.Relion.DEFOCUSV], ptcl[star.Relion.DEFOCUSANGLE],
                   ptcl[star.Relion.PHASESHIFT], ptcl[star.Relion.VOLTAGE], ptcl[star.Relion.AC], ptcl[star.Relion.CS],
                   ptcl[star.Relion.ANGLEROT], ptcl[star.Relion.ANGLETILT], ptcl[star.Relion.ANGLEPSI],
                   ptcl[star.Relion.ORIGINX], ptcl[star.Relion.ORIGINY], coefs_method, r, nr, kwargs["pfac"])

    ift = getattr(tls, 'ift', None)
    if ift is None:
        ift = irfft2(p1s.copy(), threads=kwargs["fftthreads"],
                     planner_effort="FFTW_ESTIMATE",
                     auto_align_input=True,
                     auto_contiguous=True)
        tls.ift = ift
    p1sr = fftshift(ift(p1s.copy(), np.zeros(ift.output_shape, dtype=ift.output_dtype)).copy())
    log.debug("%d@%s Exp %f +/- %f, Sub %f +/- %f" % (ptcl[star.UCSF.IMAGE_ORIGINAL_INDEX], ptcl[star.UCSF.IMAGE_ORIGINAL_PATH], np.mean(p1r), np.std(p1r), np.mean(p1sr), np.std(p1sr)))
    new_image = p1r - p1sr
    if kwargs["crop"] is not None:
        orihalf = new_image.shape[0] // 2
        newhalf = kwargs["crop"] // 2
        x = orihalf - np.int(np.round(ptcl[star.Relion.ORIGINX]))
        y = orihalf - np.int(np.round(ptcl[star.Relion.ORIGINY]))
        new_image = new_image[y - newhalf:y + newhalf, x - newhalf:x + newhalf]
    return new_image 

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 ssnr_sequential___individual_data_collect(self, r, op):
    if (self is None):
        get_mrc_func = IV.get_mrc
    else:
        get_mrc_func = self.cache.get_mrc
    v = get_mrc_func(r['subtomogram'])
    if ('angle' in r):
        v = GR.rotate_pad_mean(v, angle=N.array(r['angle'], dtype=N.float), loc_r=N.array(r['loc'], dtype=N.float))
    if ((op is not None) and ('segmentation_tg' in op) and ('template' in r) and ('segmentation' in r['template'])):
        phi = IV.read_mrc_vol(r['template']['segmentation'])
        phi_m = (phi > 0.5)
        del phi
        (ang_inv, loc_inv) = AAL.reverse_transform_ang_loc(r['angle'], r['loc'])
        phi_mr = GR.rotate(phi_m, angle=ang_inv, loc_r=loc_inv, default_val=0)
        del phi_m
        del ang_inv, loc_inv
        import aitom.tomominer.pursuit.multi.util as PMU
        v_s = PMU.template_guided_segmentation(v=v, m=phi_mr, op=op['segmentation_tg'])
        del phi_mr
        if (v_s is not None):
            v_f = N.isfinite(v_s)
            if (v_f.sum() > 0):
                v_s[N.logical_not(v_f)] = v_s[v_f].mean()
                v = v_s
            del v_s
    v = NF.fftshift(NF.fftn(v))
    m = get_mrc_func(r['mask'])
    if ('angle' in r):
        m = GR.rotate_mask(m, angle=N.array(r['angle'], dtype=N.float))
    v[(m < op['mask_cutoff'])] = 0.0
    return {'v': v, 'm': m, } 

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_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 compute(frames):
    """ computes the bidirectional phase offset

    sometimes in line scanning there will be offsets between lines;
    if ops['do_bidiphase'], then bidiphase is computed and applied

    Parameters
    ----------
    frames : int16
        random subsample of frames in binary (frames x Ly x Lx)

    Returns
    -------
    bidiphase : int
        bidirectional phase offset in pixels

    """

    Ly = frames.shape[1]
    Lx = frames.shape[2]
    # lines scanned in 1 direction
    yr1 = np.arange(1, np.floor(Ly/2)*2, 2, int)
    # lines scanned in the other direction
    yr2 = np.arange(0, np.floor(Ly/2)*2, 2, int)

    # compute phase-correlation between lines in x-direction
    d1 = fft.fft(frames[:, yr1, :], axis=2)
    d2 = np.conj(fft.fft(frames[:, yr2, :], axis=2))
    d1 = d1 / (np.abs(d1) + 1e-5)
    d2 = d2 / (np.abs(d2) + 1e-5)

    #fhg =  gaussian_fft(1, int(np.floor(Ly/2)), Lx)
    cc = np.real(fft.ifft(d1 * d2 , axis=2))#* fhg[np.newaxis, :, :], axis=2))
    cc = cc.mean(axis=1).mean(axis=0)
    cc = fft.fftshift(cc)
    ix = np.argmax(cc[(np.arange(-10,11,1) + np.floor(Lx/2)).astype(int)])
    ix -= 10
    bidiphase = -1*ix

    return bidiphase 

def fourier_transform(v):
    return fftshift(fftn(v)) 

def apply_ctf(v, ctf):
    vc = N.real(     ifftn(  ifftshift(  ctf * fftshift(fftn(v)) ) )     )           # convolute v with ctf
    return vc 

def set_fft_mid_co(self):
        siz = self.img_siz
        assert(all(N.mod(siz, 1) == 0))
        assert(all(N.array(siz) > 0))

        mid_co = N.zeros(len(siz))

        # according to following code that uses numpy.fft.fftshift()
        for i in range(len(mid_co)):
            m = siz[i]
            mid_co[i] = N.floor(m/2)
        self.mid_co = mid_co 

def filter_given_curve(v, curve):
    grid = GV.grid_displacement_to_center(v.shape, GV.fft_mid_co(v.shape))
    rad = GV.grid_distance_to_center(grid)
    rad = N.round(rad).astype(N.int)
    b = N.zeros(rad.shape)
    for (i, a) in enumerate(curve):
        b[(rad == i)] = a
    vf = ifftn(ifftshift((fftshift(fftn(v)) * b)))
    vf = N.real(vf)
    return vf 

def set_fft_mid_co(self):
        siz = self.img_siz
        assert(N.all(N.mod(siz, 1) == 0))
        assert(N.all(N.array(siz) > 0))

        mid_co = N.zeros(len(siz))

        # according to following code that uses numpy.fft.fftshift()
        for i in range(len(mid_co)):
            m = siz[i]
            mid_co[i] = N.floor(m/2)
        self.mid_co = mid_co 

def fftnc(x, axes, ortho=True):
    tmp = fft.fftshift(x, axes=axes)
    tmp = fft.fftn(tmp, axes=axes, norm="ortho" if ortho else None)
    return fft.ifftshift(tmp, axes=axes) 

def ifftnc(x, axes, ortho=True):
    tmp = fft.fftshift(x, axes=axes)
    tmp = fft.ifftn(tmp, axes=axes, norm="ortho" if ortho else None)
    return fft.ifftshift(tmp, axes=axes) 

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_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_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 compute_response(self, **kargs):
        """Compute the window data frequency response

        :param norm: True by default. normalised the frequency data.
        :param int NFFT: total length of the final data sets( 2048 by default. 
            if less than data length, then NFFT is set to the data length*2).

        The response is stored in :attr:`response`.

        .. note:: Units are dB (20 log10) since we plot the frequency response)

        """
        from numpy.fft import fft, fftshift

        norm = kargs.get('norm', self.norm)

        # do some padding. Default is max(2048, data.len*2)
        NFFT = kargs.get('NFFT', 2048)
        if NFFT < len(self.data):
            NFFT = self.data.size * 2

        # compute the fft modulus
        A = fft(self.data, NFFT)
        mag = abs(fftshift(A))

        # do we want to normalise the data
        if norm is True:
            mag = mag / max(mag)

        response = 20. * stools.log10(mag) # factor 20 we are looking at the response
                                    # not the powe
        #response = clip(response,mindB,100)
        self.__response = response 

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,)))

        assert_array_almost_equal(fft.fftshift(freqs), shifted)
        assert_array_almost_equal(fft.ifftshift(shifted), freqs) 

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_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_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,)))

        assert_array_almost_equal(fft.fftshift(freqs), shifted)
        assert_array_almost_equal(fft.ifftshift(shifted), freqs)