Python fft.fftfreq函数代码示例

def get_spectrum_1d(data_reg,x_reg,y_reg):
    """Compute the 1d power spectrum.
    """
    # remove the mean and squarize
    data_reg-=data_reg.mean()
    jpj,jpi = data_reg.shape
    msize = min(jpj,jpi)
    data_reg = data_reg[:msize-1,:msize-1]
    x_reg = x_reg[:msize-1,:msize-1]
    y_reg = y_reg[:msize-1,:msize-1]
    # wavenumber vector
    x1dreg,y1dreg = x_reg[0,:],y_reg[:,0]
    Ni,Nj = msize-1,msize-1
    dx=npy.int(npy.ceil(x1dreg[1]-x1dreg[0]))
    k_max  = npy.pi / dx
    kx = fft.fftshift(fft.fftfreq(Ni, d=1./(2.*k_max)))
    ky = fft.fftshift(fft.fftfreq(Nj, d=1./(2.*k_max)))
    kkx, kky = npy.meshgrid( ky, kx )
    Kh = npy.sqrt(kkx**2 + kky**2)
    Nmin  = min(Ni,Nj)
    leng  = Nmin/2+1
    kstep = npy.zeros(leng)
    kstep[0] =  k_max / Nmin
    for ind in range(1, leng):
        kstep[ind] = kstep[ind-1] + 2*k_max/Nmin
    norm_factor = 1./( (Nj*Ni)**2 )
    # tukey windowing = tapered cosine window
    cff_tukey = 0.25
    yw=npy.linspace(0, 1, Nj)
    wdw_j = npy.ones(yw.shape)
    xw=npy.linspace(0, 1, Ni)
    wdw_i= npy.ones(xw.shape)
    first_conditioni = xw<cff_tukey/2
    first_conditionj = yw<cff_tukey/2
    wdw_i[first_conditioni] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (xw[first_conditioni] - cff_tukey/2) ))
    wdw_j[first_conditionj] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (yw[first_conditionj] - cff_tukey/2) ))
    third_conditioni = xw>=(1 - cff_tukey/2)
    third_conditionj = yw>=(1 - cff_tukey/2)
    wdw_i[third_conditioni] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (xw[third_conditioni] - 1 + cff_tukey/2)))
    wdw_j[third_conditionj] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (yw[third_conditionj] - 1 + cff_tukey/2)))
    wdw_ii, wdw_jj = npy.meshgrid(wdw_j, wdw_i, sparse=True)
    wdw = wdw_ii * wdw_jj
    data_reg*=wdw
    #2D spectrum
    cff  = norm_factor
    tempconj=fft.fft2(data_reg).conj()
    tempamp=cff * npy.real(tempconj*fft.fft2(data_reg))
    spec_2d=fft.fftshift(tempamp)
    #1D spectrum
    leng    = len(kstep)
    spec_1d = npy.zeros(leng)
    krange     = Kh <= kstep[0]
    spec_1d[0] = spec_2d[krange].sum()
    for ind in range(1, leng):
        krange = (kstep[ind-1] < Kh) & (Kh <= kstep[ind])
        spec_1d[ind] = spec_2d[krange].sum()
    spec_1d[0] /= kstep[0]
    for ind in range(1, leng):
        spec_1d[ind] /= kstep[ind]-kstep[ind-1]
    return spec_1d, kstep

def B_ell(flat_map, component, bin_size, window=None):
    """
    Return the binned beam function of the given flat map component,
    using ell bins of bin_size. If a window name is given, multiply
    the map component by the window function before taking the 2-D
    DFT.

    The returned beam function is the absolute value of the DFT, so it
    has the same units as the map component. The returned bins array
    contains the left edges of the bins corresponding to the returned
    data: for all i in range(bins.size),
    bins[i] <= data[i] < bins[i+1]
    """
    ell_x, ell_y= np.meshgrid(fftshift(fftfreq(flat_map.x.size, flat_map.dx()/(2*np.pi))),
                              fftshift(fftfreq(flat_map.y.size, flat_map.dy()/(2*np.pi))))
    ell_x = ell_x.T
    ell_y = ell_y.T
    ell_r = np.sqrt(ell_x**2 + ell_y**2)
    ell_bins = np.arange(0, np.max(ell_r), bin_size)
    beam = flat_map[component]
    if window is not None:
        beam *= get_window(window, flat_map.y.size)
        beam *= get_window(window, flat_map.x.size)[:, np.newaxis]
    dft = fftshift(fft2(beam))
    # Shift the indices down by 1 so that they correspond to the data
    # indices. These indices should always be valid indices for the
    # DFT data because r_ell has a lower bound at 0 and max(r_ell)
    # will have index ell_bins.size.
    bin_indices = np.digitize(ell_r.flatten(), ell_bins) - 1
    binned = np.zeros(ell_bins.size) # dtype?
    for i in range(binned.size):
        binned[i] = np.sqrt(np.mean(abs(dft.flatten()[i==bin_indices])**2))
    return ell_bins, binned, dft

    def applyPreconditioner(self, arr):

        ftMap = fft2(arr)

        Ny = (arr.shape)[0]
        Nx = (arr.shape)[1]

        pixScaleX = (
            numpy.abs(self.ra1 - self.ra0)
            / Nx
            * numpy.pi
            / 180.0
            * numpy.cos(numpy.pi / 180.0 * 0.5 * (self.dec0 + self.dec1))
        )
        pixScaleY = numpy.abs(self.dec1 - self.dec0) / Ny * numpy.pi / 180.0

        lx = 2 * numpy.pi * fftfreq(Nx, d=pixScaleX)
        ly = 2 * numpy.pi * fftfreq(Ny, d=pixScaleY)

        ix = numpy.mod(numpy.arange(Nx * Ny), Nx)
        iy = numpy.arange(Nx * Ny) / Nx

        lMap = ftMap * 0.0

        lMap[iy, ix] = numpy.sqrt(lx[ix] ** 2 + ly[iy] ** 2)

        lOverl0 = lMap / self.l0

        Fl = self.A * ((lOverl0) ** self.alpha) / (1.0 + lOverl0 ** self.alpha)

        filteredFtMap = ftMap / Fl

        filteredFtMap[0, 0] = 0.0  # avoids nan and makes mean 0

        arr[:, :] = (numpy.real(ifft2(filteredFtMap)))[:, :]

    def icwt2d(self, da=0.25):
        '''
        Inverse bi-dimensional continuous wavelet transform as in Wang and
        Lu (2010), equation [5].

        Parameters
        ----------
        da : float, optional
            Spacing in the frequency axis.
        '''
        if self.Wf is None:
            raise TypeError("Run cwt2D before icwt2D")
        m0, l0, k0 = self.Wf.shape

        if m0 != self.scales.size:
            raise Warning('Scale parameter array shape does not match\
                           wavelet transform array shape.')
        # Calculates the zonal and meridional wave numters.
        L, K = 2 ** int(np.ceil(np.log2(l0))), 2 ** int(np.ceil(np.log2(k0)))
        # Calculates the zonal and meridional wave numbers.
        l, k = fftfreq(L, self.dy), fftfreq(K, self.dx)
        # Creates empty inverse wavelet transform array and fills it for every
        # discrete scale using the convolution theorem.
        self.iWf = np.zeros((m0, L, K), 'complex')
        for i, an in enumerate(self.scales):
            psi_ft_bar = an * self.wavelet.psi_ft(an * k, an * l)
            W_ft = fftn(self.Wf[i, :, :], s=(L, K))
            self.iWf[i, :, :] = ifftn(W_ft * psi_ft_bar, s=(L, K)) *\
                da / an ** 2.

        self.iWf = self.iWf[:, :l0, :k0].real.sum(axis=0) / self.wavelet.cpsi

        return self

    def gvecs(self):
        """
        Array with the reduced coordinates of the G-vectors.

        .. note::

            These are the G-vectors of the FFT box and should be used
            when we compute quantities on the FFT mesh.
            These vectors differ from the gvecs stored in |GSphere| that
            are k-centered and enclosed by a sphere whose radius is defined by ecut.
        """
        gx_list = np.rint(fftfreq(self.nx) * self.nx)
        gy_list = np.rint(fftfreq(self.ny) * self.ny)
        gz_list = np.rint(fftfreq(self.nz) * self.nz)
        #print(gz_list, gy_list, gx_list)

        gvecs = np.empty((self.size, 3), dtype=np.int)

        idx = -1
        for gx in gx_list:
            for gy in gy_list:
                for gz in gz_list:
                    idx += 1
                    gvecs[idx, :] = gx, gy, gz

        return gvecs

 def test_definition(self):
     x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
     assert_array_almost_equal(9*fft.fftfreq(9), x)
     assert_array_almost_equal(9*pi*fft.fftfreq(9, pi), x)
     x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
     assert_array_almost_equal(10*fft.fftfreq(10), x)
     assert_array_almost_equal(10*pi*fft.fftfreq(10, pi), x)

def subpixel_shift(img_arr, dx=0.5, inplace=True):
    '''
    Shifts an image by a fraction of a pixel in a random direction,
    with circular boundary conditions.

    Arguments:
        img_arr (2D ndarray): array to subpixel shift in a random direction
        dx (float): distance in pixel to shift by
        inplace (bool): if True, will modify the array inplace

    '''

    f = fft.fft2(img_arr, axes=(0, 1))
    g = np.meshgrid(
        fft.fftfreq(img_arr.shape[-3]),
        fft.fftfreq(img_arr.shape[-2])
    )

    u = npr.normal(0, 1, size=2)
    fk = g[0] * u[0] + g[1] * u[1]
    phi = np.exp(2j * np.pi * dx * fk)

    out = np.real(fft.ifft2(phi[..., np.newaxis] * f, axes=(0, 1)))
    if inplace:
        img_arr[:] = out[:]
    else:
        return out

    def __call__(self, r):
        """Calculate the characteristic function from the transform of the BaseModel."""

        # Determine q-values.
        # We want very fine r-spacing so we can properly normalize f(r). This
        # equates to having a large qmax so that the Fourier transform is
        # finely spaced. We also want the calculation to be fast, so we pick
        # qmax such that the number of q-points is a power of 2. This allows us
        # to use the fft. 
        # 
        # The initial dr is somewhat arbitrary, but using dr = 0.01 allows for
        # the f(r) calculated from a particle of diameter 50, over r =
        # arange(1, 60, 0.1) to agree with the sphericalCF with Rw < 1e-4%.
        #
        # We also have to make a q-spacing small enough to compute out to at
        # least the size of the signal. 
        dr = min(0.01, r[1] - r[0])
        ed = 2 * self._model.calculate_ER()

        # Check for nans. If we find any, then return zeros.
        if numpy.isnan(ed).any():
            y = numpy.zeros_like(r)
            return y

        rmax = max(ed, 2 * r[-1])
        dq = pi / rmax
        qmax = pi / dr
        numpoints = int(2**(ceil(log2(qmax/dq))))
        qmax = dq * numpoints

        # Calculate F(q) = q * I(q) from model
        q = fftfreq(int(qmax/dq)) * qmax
        fq = q * self._model.evalDistribution(q)

        # Calculate g(r) and the effective r-points
        rp = fftfreq(numpoints) * 2 * pi / dq
        # Note sine transform = imaginary part of ifft
        gr = ifft(fq).imag

        # Calculate full-fr for normalization
        assert (rp[0] == 0.0)
        frp = numpy.zeros_like(gr)
        frp[1:] = gr[1:] / rp[1:]

        # Inerpolate onto requested grid, do not use data after jump in rp
        assert (numpoints % 2 == 0)
        nhalf = numpoints / 2
        fr = numpy.interp(r, rp[:nhalf], gr[:nhalf])
        vmask = (r != 0)
        fr[vmask] /= r[vmask]

        # Normalize. We approximate fr[0] by using the fact that f(r) is linear
        # at low r. By definition, fr[0] should equal 1.
        fr0 = 2*frp[2] - frp[1]
        fr /= fr0

        # Fix potential divide-by-zero issue, fr is 1 at r == 0
        fr[~vmask] = 1

        return fr

def _configure(infiles, z, df):
    conf.infiles = infiles
    img_hdr = fits.getheader(conf.infiles[0])
    conf.nx = img_hdr['naxis1']
    conf.ny = img_hdr['naxis2']
    conf.nf = MWA_FREQ_EOR_ALL_80KHZ.size
    conf.dx = np.abs(img_hdr['cdelt1']) * np.pi / 180
    conf.dy = np.abs(img_hdr['cdelt2']) * np.pi / 180
    conf.df = df
    conf.du = 1 / (conf.nx * conf.dx)
    conf.dv = 1 / (conf.ny * conf.dy)
    conf.deta = 1 / (conf.nf * conf.df)
    conf.freq = MWA_FREQ_EOR_ALL_80KHZ
    conf.u = fftshift(fftfreq(conf.nx, conf.dx))
    conf.v = fftshift(fftfreq(conf.ny, conf.dy))
    conf.eta = fftshift(fftfreq(conf.nf, conf.df))
    conf.z = z
    conf.cosmo_d = Cosmo.comoving_transverse_distance(conf.z).value
    conf.cosmo_e = np.sqrt(Cosmo.Om0 * (1 + conf.z) ** 3 + Cosmo.Ok0 *
                           (1 + conf.z) ** 2 + Cosmo.Ode0)
    conf.cosmo_h0 = Cosmo.H0.value * 1e3
    conf.cosmo_c = const.si.c.value
    conf.kx = conf.u * 2 * np.pi / conf.cosmo_d
    conf.ky = conf.v * 2 * np.pi / conf.cosmo_d
    conf.dkx = conf.du * 2 * np.pi / conf.cosmo_d
    conf.dky = conf.dv * 2 * np.pi / conf.cosmo_d
    conf.k_perp = np.sqrt(conf.kx ** 2 + conf.ky[np.newaxis, ...].T ** 2)
    conf.k_par = conf.eta * 2 * np.pi * conf.cosmo_h0 * _F21 * \
        conf.cosmo_e / (conf.cosmo_c * (1 + conf.z) ** 2)

def fftFromLiteMap(liteMap,applySlepianTaper = False,nresForSlepian=3.0):
    """
    @brief Creates an fft2D object out of a liteMap
    @param liteMap The map whose fft is being taken
    @param applySlepianTaper If True applies the lowest order taper (to minimize edge-leakage)
    @param nresForSlepian If above is True, specifies the resolution of the taeper to use.
    """
    ft = fft2D()
        
    ft.Nx = liteMap.Nx
    ft.Ny = liteMap.Ny
    flTrace.issue("flipper.fftTools",1, "Taking FFT of map with (Ny,Nx)= (%f,%f)"%(ft.Ny,ft.Nx))
    
    ft.pixScaleX = liteMap.pixScaleX 
                
    
    ft.pixScaleY = liteMap.pixScaleY
    
    
    lx =  2*np.pi  * fftfreq( ft.Nx, d = ft.pixScaleX )
    ly =  2*np.pi  * fftfreq( ft.Ny, d = ft.pixScaleY )
    
    ix = np.mod(np.arange(ft.Nx*ft.Ny),ft.Nx)
    # This was dividing ints, so change below needed to maintain same behaviour in python3
    iy = np.array(np.arange(ft.Nx*ft.Ny)/ft.Nx, dtype = int)     
    
    modLMap = np.zeros([ft.Ny,ft.Nx])
    modLMap[iy,ix] = np.sqrt(lx[ix]**2 + ly[iy]**2)
    
    ft.modLMap  =  modLMap
    
    ft.lx = lx
    ft.ly = ly
    ft.ix = ix
    ft.iy = iy
    ft.thetaMap = np.zeros([ft.Ny,ft.Nx])
    ft.thetaMap[iy[:],ix[:]] = np.arctan2(ly[iy[:]],lx[ix[:]])
    ft.thetaMap *=180./np.pi
    
    map = liteMap.data.copy()
    #map = map0.copy()
    #map[:,:] =map0[::-1,:]
    taper = map.copy()*0.0 + 1.0

    if (applySlepianTaper) :
        try:
            path_to_taper = os.path.join(__FLIPPER_DIR__, "tapers", 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny,ft.Nx,nresForSlepian))
            f = open(path_to_taper)
            taper = pickle.load(f)
            f.close()
        except:
            taper = slepianTaper00(ft.Nx,ft.Ny,nresForSlepian)
            path_to_taper = os.path.join(__FLIPPER_DIR__, "tapers", 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny,ft.Nx,nresForSlepian))
            f = open(path_to_taper, mode="w")
            pickle.dump(taper,f)
            f.close()
    
    ft.kMap = fftfast.fft(map*taper,axes=[-2,-1])
    del map, modLMap, lx, ly
    return ft

def high_pass_filter(img, filtersize=10):
    """
    A FFT implmentation of high pass filter from pyKLIP.
    Args:
        img: a 2D image
        filtersize: size in Fourier space of the size of the space. In image space, size=img_size/filtersize
    Returns:
        filtered: the filtered image
    """
    if filtersize == 0:
        return img
    # mask NaNs
    nan_index = np.where(np.isnan(img))
    img[nan_index] = 0

    transform = fft.fft2(img)

    # coordinate system in FFT image
    u,v = np.meshgrid(fft.fftfreq(transform.shape[1]), fft.fftfreq(transform.shape[0]))
    # scale u,v so it has units of pixels in FFT space
    rho = np.sqrt((u*transform.shape[1])**2 + (v*transform.shape[0])**2)
    # scale rho up so that it has units of pixels in FFT space
    # rho *= transform.shape[0]
    # create the filter
    filt = 1. - np.exp(-(rho**2/filtersize**2))

    filtered = np.real(fft.ifft2(transform*filt))

    # restore NaNs
    filtered[nan_index] = np.nan
    img[nan_index] = np.nan

    return filtered

    def __init__(self, cube, header, phys_units=False):
        super(VCS, self).__init__()

        self.header = header
        self.cube = cube
        self.fftcube = None
        self.correlated_cube = None
        self.ps1D = None
        self.phys_units = phys_units

        if np.isnan(self.cube).any():
            self.cube[np.isnan(self.cube)] = 0
            # Feel like this should be more specific
            self.good_pixel_count = np.sum(self.cube.max(axis=0) != 0)
        else:
            self.good_pixel_count = float(self.cube.shape[1] * self.cube.shape[2])

        # Lazy check to make sure we have units of km/s
        if np.abs(self.header["CDELT3"]) > 1:
            self.vel_to_pix = np.abs(self.header["CDELT3"]) / 1000.0
        else:
            self.vel_to_pix = np.abs(self.header["CDELT3"])

        self.vel_channels = np.arange(1, self.cube.shape[0], 1)

        if self.phys_units:
            self.vel_freqs = np.abs(fftfreq(self.cube.shape[0])) / self.vel_to_pix
        else:
            self.vel_freqs = np.abs(fftfreq(self.cube.shape[0]))

    def source_terms(self, mara, retphi=False):
        from numpy.fft import fftfreq, fftn, ifftn
        ng = mara.number_guard_zones()
        G = self.G
        L = 1.0
        Nx, Ny, Nz = mara.fluid.shape
        Nx -= 2*ng
        Ny -= 2*ng
        Nz -= 2*ng
        P = mara.fluid.primitive[ng:-ng,ng:-ng,ng:-ng]
        rho = P[...,0]
        vx = P[...,2]
        vy = P[...,3]
        vz = P[...,4]

        K = [fftfreq(Nx)[:,np.newaxis,np.newaxis] * (2*np.pi*Nx/L),
             fftfreq(Ny)[np.newaxis,:,np.newaxis] * (2*np.pi*Ny/L),
             fftfreq(Nz)[np.newaxis,np.newaxis,:] * (2*np.pi*Nz/L)]
        delsq = -(K[0]**2 + K[1]**2 + K[2]**2)
        delsq[0,0,0] = 1.0 # prevent division by 0

        rhohat = fftn(rho)
        phihat = (4*np.pi*G) * rhohat / delsq
        fx = -ifftn(1.j * K[0] * phihat).real
        fy = -ifftn(1.j * K[1] * phihat).real
        fz = -ifftn(1.j * K[2] * phihat).real

        S = np.zeros(mara.fluid.shape + (5,))
        S[ng:-ng,ng:-ng,ng:-ng,0] = 0.0
        S[ng:-ng,ng:-ng,ng:-ng,1] = rho * (fx*vx + fy*vy + fz*vz)
        S[ng:-ng,ng:-ng,ng:-ng,2] = rho * fx
        S[ng:-ng,ng:-ng,ng:-ng,3] = rho * fy
        S[ng:-ng,ng:-ng,ng:-ng,4] = rho * fz
        return (S, ifftn(phihat).real) if retphi else S

def compare_interpolated_spectrum():
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111)

    out = fft(ifftshift(f_full))
    freqs = fftfreq(len(f_full), d=0.01) # spacing, Ang
    sfreqs = fftshift(freqs)
    taper = gauss_taper(freqs, sigma=0.0496) #Ang, corresponds to 2.89 km/s at 5150A.
    tout = out * taper

    ax.plot(sfreqs, fftshift(tout))

    wl_h, fl_h = np.abs(np.load("PH6.8kms_0.01ang.npy"))
    wl_l, fl_l = np.abs(np.load("PH2.5kms.npy"))

    #end edges
    wl_he = wl_h[200:-200]
    fl_he = fl_h[200:-200]
    interp = Sinc_w(wl_l, fl_l, a=5, window='kaiser')
    fl_hi = interp(wl_he)

    d = wl_he[1] - wl_he[0]
    out = fft(ifftshift(fl_hi))
    freqs = fftfreq(len(out), d=d)
    ax.plot(fftshift(freqs), fftshift(out))

    plt.show()

def icwt2d(W, a, dx=0.25, dy=0.25, da=0.25, wavelet=Mexican_hat()):
    """
    Inverse bi-dimensional continuous wavelet transform as in Wang and
    Lu (2010), equation [5].
    PARAMETERS
        W (array like):
            Wavelet transform, the result of the cwt2d function.
        a (array like, optional):
            Scale parameter array.
        w (class, optional) :
            Mother wavelet class. Default is Mexican_hat()
    RETURNS
        iW (array like) :
            Inverse wavelet transform.
    EXAMPLE
    """
    m0, l0, k0 = W.shape
    if m0 != a.size:
        raise Warning('Scale parameter array shape does not match wavelet' \
                       ' transform array shape.')
    # Calculates the zonal and meridional wave numters.
    L, K = 2 ** int(ceil(log2(l0))), 2 ** int(ceil(log2(k0)))
    # Calculates the zonal and meridional wave numbers.
    l, k = fftfreq(L, dy), fftfreq(K, dx)
    # Creates empty inverse wavelet transform array and fills it for every 
    # discrete scale using the convolution theorem.
    iW = zeros((m0, L, K), 'complex')
    for i, an in enumerate(a):
        psi_ft_bar = an * wavelet.psi_ft(an * k, an * l)
        W_ft = fft2(W[i, :, :], s=(L, K))
        iW[i, :, :] = ifft2(W_ft * psi_ft_bar, s=(L, K)) * da / an ** 2.

    return iW[:, :l0, :k0].real.sum(axis=0) / wavelet.cpsi