numpy.fft 实现 czt (Chirp Z-transform)

numpy.fft 实现 czt (Chirp Z-transform)

动机

如果对$L^2(\mathbb{R})$上做Fourier变换，直接用离散FFT是不行的。需要用CZT。用于数值计算的numpy没有提供CZT，需要重新实现。本文用FFT实现CZT。

f ^ ( ω ) = ∫ R f ( x ) e − 2 π i ω x d x ∼ ∫ [ a , b ] f ( x ) e − 2 π i ω x d x , − a , b ≫ 1 ∼ b − a N ( ∑ l = 0 N − 1 f ( l N ) e − 2 π i ω ( a + ( b − a ) l / N ) + f ( b ) e − 2 π i ω b − f ( a ) e − 2 π i ω a 2 ) = b − a N e − 2 π i ω a ( ∑ l = 0 N − 1 f ( l N ) e − 2 π i ( b − a ) ω l / N + f ( b ) e 2 π i ω ( a − b ) − f ( a ) 2 ) ⟹ f ^ ( c + k d − c M ) = b − a N e − 2 π i a ( c + k d − c M ) ( ∑ l = 0 N − 1 f ( l N ) e − 2 π i b − a N c l e − 2 π i ( b − a ) ( d − c ) M N k l + . . . ) = b − a N e − 2 π i a ( c + k d − c M ) ( C Z T ( { f l } , M , e − 2 π i ( b − a ) ( d − c ) M N , e 2 π i b − a N c ) + . . . ) \hat{f}(\omega)=\int_\R f(x)e^{-2\pi i\omega x} \mathrm{d}x\\ \sim \int_{[a, b]} f(x)e^{-2\pi i\omega x} \mathrm{d}x, -a, b\gg 1 \\ \sim \frac{b-a}{N}(\sum_{l=0}^{N-1} f(\frac{l}{N})e^{-2\pi i \omega (a+(b-a)l/N)} + \frac{f(b)e^{-2\pi i\omega b}-f(a)e^{-2\pi i\omega a}}{2})\\ = \frac{b-a}{N}e^{-2\pi i \omega a}(\sum_{l=0}^{N-1} f(\frac{l}{N})e^{-2\pi i (b-a) \omega l/N} + \frac{f(b)e^{2\pi i\omega (a-b)}-f(a)}{2})\\ \Longrightarrow\\ \hat{f}(c+k\frac{d-c}{M})=\frac{b-a}{N}e^{-2\pi i a(c+k\frac{d-c}{M})}(\sum_{l=0}^{N-1} f(\frac{l}{N})e^{-2\pi i \frac{b-a}{N} cl}e^{-2\pi i \frac{(b-a)(d-c)}{MN}k l} + ...)\\ =\frac{b-a}{N}e^{-2\pi i a(c+k\frac{d-c}{M})}(CZT(\{f_l\}, M, e^{-2\pi i \frac{(b-a)(d-c)}{MN}}, e^{2\pi i \frac{b-a}{N} c})+...) f^​(ω)=∫R​f(x)e−2πiωxdx∼∫[a,b]​f(x)e−2πiωxdx,−a,b≫1∼Nb−a​(l=0∑N−1​f(Nl​)e−2πiω(a+(b−a)l/N)+2f(b)e−2πiωb−f(a)e−2πiωa​)=Nb−a​e−2πiωa(l=0∑N−1​f(Nl​)e−2πi(b−a)ωl/N+2f(b)e2πiω(a−b)−f(a)​)⟹f^​(c+kMd−c​)=Nb−a​e−2πia(c+kMd−c​)(l=0∑N−1​f(Nl​)e−2πiNb−a​cle−2πiMN(b−a)(d−c)​kl+...)=Nb−a​e−2πia(c+kMd−c​)(CZT({fl​},M,e−2πiMN(b−a)(d−c)​,e2πiNb−a​c)+...)

原理

基于Bluestein算法
参考 Chirp Z-transform 及其中引文

源码

def czt(x, m=None, w=None, a=1):
    """chirp z transform: (Based on FFT)

        z = A * W.^(-(0:M-1)), g_i = sum_jz^{-ij}x_j
    Inputs m, w and a must be scalars.
    
    Arguments:
        x {array} -- the original signal
    
    Keyword Arguments:
        m {number} -- (default: {None})
        w {complex number} -- base of the series (default: {np.exp(-2j * np.pi / m)})
        a {number} -- (default: {1})
    
    Returns:
        array -- czt of x
    """
    n = x.shape[0]
    if k is None:
        m = n
    if w is None:
        w = np.exp(-2j * np.pi / m)

    if not isinstance(x, np.ndarray):
        x = np.asarray(x, dtype=np.complex)

    # Length for power-of-two fft.
    nfft = 2 ** nextpow2(n+k-1)
    # Premultiply data.
    k = np.arange(-n+1, max((m, n))) ** 2 / 2
    ww = w ** k   # Chirp filter is 1./ww
    # nn = n-1 : 2*n-1
    aa = a ** (-np.arange(0, n)) * ww[n-1:2*n-1]
    y = x * aa

    # Fast convolution via FFT.
    fy = np.fft.fft(y, nfft) * np.fft.fft(1 / ww[:m+n-1], nfft)
    g = np.fft.ifft(fy)
    # Finally multiply Chirp filter.
    return g[n-1:n+m-1] * ww[n-1:n+m-1]

Fourier 变换实现

有了czt就可以实现FT了

def cft(f, trange=(0,1), Nt=None, wrange=(-10,10), Nw=101, padding='c'):
    """continuous Fourier transform
       
    Assertion: wsize * tsize <= Nt;
    Reference: https://en.wikipedia.org/wiki/Chirp_Z-transform
    
    Arguments:
        f {function|array} -- function or values of functions
    
    Keyword Arguments:
        trange {tuple} -- the range of time (default: {(0,1)})
        Nt {[type]} -- number of samples on time-domain (default: {None})
        wrange {tuple} -- the range of frequency (default: {(-10,10)})
        Nw {number} -- number of samples on frequency-domain (default: {101})
        padding {str} -- Used in feature (default: {'c'})
    
    Returns:
        tuple -- function values and frequencies
    """

    a, b = trange
    c, d = wrange
    dt = (b-a)/ Nt; dw = (d-c)/ Nw
    
    if callable(f):
        if Nt is None:
            Nt = 100
        f = f(np.linspace(a, b, Nt))
    else:
        if Nt is None:
            Nt = f.shape[0]
        else:
            f = f[:Nt]
    F = czt(f, Nw-1, np.exp(-2j*np.pi*dt*dw), np.exp(2j*np.pi*dt*c))
    ws = np.linspace(c, d, Nw)
    if padding=='c':
        p = f[-1]
    elif padding=='c1':
        p = 2*f[-1]-f[-2]
    F = dt * np.exp(-2j*np.pi*a*ws[:-1]) * (F + (p * np.exp(2j*np.pi*(a-b) * ws[:-1]) - f[0])/2)
    return np.concatenate([F, [F[-1]]]), ws