FFT_Stock_Prediction.py 文件源码

def fourierEx(x, n_predict, harmonics):

    n = len(x)                  # number of input samples
    n_harmonics = harmonics            # f, 2*f, 3*f, .... n_harmonics  ( 1,2, )
    t = np.arange(0, n)         # place to store data
    p = np.polyfit(t, x, 1)     # find trend
    x_no_trend = x - p[0] * t 
    x_frequency_domains = fft.fft(x_no_trend)
    f = np.fft.fftfreq(n)       # frequencies
    indexes = list(range(n))
    indexes.sort(key=lambda i: np.absolute(f[i]))

    t = np.arange(0, n + n_predict)
    restored_signal = np.zeros(t.size)
    for i in indexes[:1 + n_harmonics * 2]:
        amplitude = np.absolute(x_frequency_domains[i] / n)
        phase = np.angle(x_frequency_domains[i])
        restored_signal += amplitude * np.cos(2 * np.pi * f[i] * t + phase)

    return restored_signal + p[0] * t