Romanchenko

Untitled

Dec 1st, 2020
542
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. def mFFT(arr, draw=False, name='a[n]', x1=0, x2=-1, y2=-1, norm=-1, draw_less=False, return_pair=False):
  2.     """
  3.    Parameters:
  4.        arr : list
  5.            Массив, для которого считается Фурье
  6.        draw : bool
  7.            Рисовать ли график
  8.        name : str, optional
  9.            Заголовок графика
  10.        x1 : double, optional
  11.            Левая граница по x на графике спектра
  12.        x2 : double, optional
  13.            Правая граница по x на графике спектра
  14.        y2 : double, optional
  15.            Верхняя граница по y на графике спектра
  16.        norm : double
  17.            Число для нормировки, стандартное - длина массива
  18.        draw_less : bool
  19.             Надо ли рисовать график последовательности вдобавок к спектру частот
  20.        return_pair : bool
  21.            Возвращать ли массив периодов (координаты по x)
  22.  
  23.    Returns:
  24.        Если не выставлен флаг energy, то возвращается массив абсолютных значений амплитуд спектра.
  25.        Если флаг выставлен, то возращается кортеж из четырех элементов - энергия амплитуды 27-дневного сигнала,
  26.        энергия амплитуды полугодового сигнала, их отношение и массив абсолютных значений амплитуд сигнала.
  27.    """
  28.     if x2 == -1:
  29.         x2 = len(arr) // 2
  30.     if norm == -1:
  31.         norm = len(arr)
  32.  
  33.     A = np.fft.rfft((arr - np.mean(arr)) / norm)
  34.     n = np.arange(len(arr))
  35.     n1 = len(arr) / n[1:]
  36.  
  37.     if draw:
  38.  
  39.         plt.rcParams['axes.grid'] = True
  40.         if not draw_less:
  41.             fig, ax = plt.subplots(2, figsize=(6, 4), dpi=150)
  42.             plt.tight_layout()
  43.  
  44.             plt.subplots_adjust(hspace=0.5)
  45.             ax[0].plot(n, arr, '.-')
  46.             ax[0].set_title(name)
  47.  
  48.             ax[1].set_title('$A$')
  49.             if y2 > 0:
  50.                 ax[1].set_ylim(0, y2)
  51.             ax[1].set_xlim(x1, x2)
  52.             ax[1].plot(n1[0: (len(arr) // 2 + 1)], np.abs(A), '-')
  53.             ax[1].set_xlabel('Период (в днях)')
  54.         else:
  55.             fig, ax = plt.subplots(1, figsize=(4, 2), dpi=150)
  56.             plt.tight_layout()
  57.             plt.subplots_adjust(hspace=0.5)
  58.             ax.set_title('$A$')
  59.  
  60.             if y2 > 0:
  61.                 ax.set_ylim(0, y2)
  62.             ax.set_xlim(x1, x2)
  63.             ax.plot(n1[0: (len(arr) // 2 + 1)], np.abs(A), '-')
  64.             ax.set_xlabel('Период (в днях)')
  65.  
  66.         plt.show()
  67.  
  68.     if return_pair:
  69.         return np.abs(A), n1[0: (len(arr) // 2 + 1)]
  70.     return np.abs(A)
RAW Paste Data