How to calculate a Fourier series in Numpy?

前端 未结 5 1299
遇见更好的自我
遇见更好的自我 2020-11-30 01:42

I have a periodic function of period T and would like to know how to obtain the list of the Fourier coefficients. I tried using fft module from numpy but it seems more dedic

5条回答
  •  失恋的感觉
    2020-11-30 02:27

    This is an old question, but since I had to code this, I am posting here the solution that uses the numpy.fft module, that is likely faster than other hand-crafted solutions.

    The DFT is the right tool for the job of calculating up to numerical precision the coefficients of the Fourier series of a function, defined as an analytic expression of the argument or as a numerical interpolating function over some discrete points.

    This is the implementation, which allows to calculate the real-valued coefficients of the Fourier series, or the complex valued coefficients, by passing an appropriate return_complex:

    def fourier_series_coeff_numpy(f, T, N, return_complex=False):
        """Calculates the first 2*N+1 Fourier series coeff. of a periodic function.
    
        Given a periodic, function f(t) with period T, this function returns the
        coefficients a0, {a1,a2,...},{b1,b2,...} such that:
    
        f(t) ~= a0/2+ sum_{k=1}^{N} ( a_k*cos(2*pi*k*t/T) + b_k*sin(2*pi*k*t/T) )
    
        If return_complex is set to True, it returns instead the coefficients
        {c0,c1,c2,...}
        such that:
    
        f(t) ~= sum_{k=-N}^{N} c_k * exp(i*2*pi*k*t/T)
    
        where we define c_{-n} = complex_conjugate(c_{n})
    
        Refer to wikipedia for the relation between the real-valued and complex
        valued coeffs at http://en.wikipedia.org/wiki/Fourier_series.
    
        Parameters
        ----------
        f : the periodic function, a callable like f(t)
        T : the period of the function f, so that f(0)==f(T)
        N_max : the function will return the first N_max + 1 Fourier coeff.
    
        Returns
        -------
        if return_complex == False, the function returns:
    
        a0 : float
        a,b : numpy float arrays describing respectively the cosine and sine coeff.
    
        if return_complex == True, the function returns:
    
        c : numpy 1-dimensional complex-valued array of size N+1
    
        """
        # From Shanon theoreom we must use a sampling freq. larger than the maximum
        # frequency you want to catch in the signal.
        f_sample = 2 * N
        # we also need to use an integer sampling frequency, or the
        # points will not be equispaced between 0 and 1. We then add +2 to f_sample
        t, dt = np.linspace(0, T, f_sample + 2, endpoint=False, retstep=True)
    
        y = np.fft.rfft(f(t)) / t.size
    
        if return_complex:
            return y
        else:
            y *= 2
            return y[0].real, y[1:-1].real, -y[1:-1].imag
    

    This is an example of usage:

    from numpy import ones_like, cos, pi, sin, allclose
    T = 1.5  # any real number
    
    def f(t):
        """example of periodic function in [0,T]"""
        n1, n2, n3 = 1., 4., 7.  # in Hz, or nondimensional for the matter.
        a0, a1, b4, a7 = 4., 2., -1., -3
        return a0 / 2 * ones_like(t) + a1 * cos(2 * pi * n1 * t / T) + b4 * sin(
            2 * pi * n2 * t / T) + a7 * cos(2 * pi * n3 * t / T)
    
    
    N_chosen = 10
    a0, a, b = fourier_series_coeff_numpy(f, T, N_chosen)
    
    # we have as expected that
    assert allclose(a0, 4)
    assert allclose(a, [2, 0, 0, 0, 0, 0, -3, 0, 0, 0])
    assert allclose(b, [0, 0, 0, -1, 0, 0, 0, 0, 0, 0])
    

    And the plot of the resulting a0,a1,...,a10,b1,b2,...,b10 coefficients: enter image description here

    This is an optional test for the function, for both modes of operation. You should run this after the example, or define a periodic function f and a period T before running the code.

    # #### test that it works with real coefficients:
    from numpy import linspace, allclose, cos, sin, ones_like, exp, pi, \
        complex64, zeros
    
    
    def series_real_coeff(a0, a, b, t, T):
        """calculates the Fourier series with period T at times t,
           from the real coeff. a0,a,b"""
        tmp = ones_like(t) * a0 / 2.
        for k, (ak, bk) in enumerate(zip(a, b)):
            tmp += ak * cos(2 * pi * (k + 1) * t / T) + bk * sin(
                2 * pi * (k + 1) * t / T)
        return tmp
    
    
    t = linspace(0, T, 100)
    f_values = f(t)
    a0, a, b = fourier_series_coeff_numpy(f, T, 52)
    # construct the series:
    f_series_values = series_real_coeff(a0, a, b, t, T)
    # check that the series and the original function match to numerical precision:
    assert allclose(f_series_values, f_values, atol=1e-6)
    
    # #### test similarly that it works with complex coefficients:
    
    def series_complex_coeff(c, t, T):
        """calculates the Fourier series with period T at times t,
           from the complex coeff. c"""
        tmp = zeros((t.size), dtype=complex64)
        for k, ck in enumerate(c):
            # sum from 0 to +N
            tmp += ck * exp(2j * pi * k * t / T)
            # sum from -N to -1
            if k != 0:
                tmp += ck.conjugate() * exp(-2j * pi * k * t / T)
        return tmp.real
    
    f_values = f(t)
    c = fourier_series_coeff_numpy(f, T, 7, return_complex=True)
    f_series_values = series_complex_coeff(c, t, T)
    assert allclose(f_series_values, f_values, atol=1e-6)
    

提交回复
热议问题