[numpy] Decompose signal via fft

import pandas as pd
import dask.dataframe as dd
import numpy as np
import os
from sklearn.datasets import make_classification

import matplotlib.pyplot as plt

# parameters
n_samples = 20000

def decompose_into_n_signals(srs, n):
    fourier = np.fft.rfft(srs)
    frequencies = np.fft.rfftfreq(srs.size, d=2e-2/srs.size)
    out = []
    for vals in np.array_split(frequencies, n):
        ft_threshed = fourier.copy()
        ft_threshed[(vals.min() > frequencies)] = 0
        ft_threshed[(vals.max() < frequencies)] = 0
        out.append(np.fft.irfft(ft_threshed))
    return out

def cut_into_n_bands(srs, n):
    fourier = np.fft.rfft(srs)
    frequencies = np.fft.rfftfreq(srs.size, d=2e-2/srs.size)
    out = []
    for vals in np.array_split(frequencies, n):
        ft_threshed = fourier.copy()
        ft_threshed[(vals.min() > frequencies)] = 0
        ft_threshed[(vals.max() < frequencies)] = 0
        out.append(ft_threshed.real)
    return out

def plot_n_signals(sigs):
    fig, axarr = plt.subplots(len(sigs), figsize=(12, 20))
    for i, sig in enumerate(sigs):
        plt.sca(axarr[i])
        plt.plot(sig)
    plt.gcf().suptitle(f"Decomposition of signal into {len(sigs)} frequency bands", fontsize=24)

# features
df = make_classification(n_samples = n_samples, n_features = 10)
df_raw = pd.DataFrame(df[0], columns = ['var1', 'var2', 'var3', 'var4', 'var5', 'var6', 'var7', 'var8', 'var9', 'var10'])

# class
df_raw['class'] = df[1]
del df

n_signals = 10
plot_n_signals(cut_into_n_bands(df_raw['var1'].values, n_signals))

signal_decomposed = decompose_into_n_signals(df_raw['var1'].values, n_signals)
plot_n_signals(signal_decomposed)
plt.figure(figsize=(20,10))
plt.plot(signal_decomposed[0])