import numpy as np
import numpy.fft as nf
import matplotlib.pyplot as mp
# 创建几个模拟信号
times = np.linspace(0, 2 * np.pi, 201)
sigs1 = (1 * np.pi) * np.sin(1 * times)
sigs2 = (3 * np.pi) * np.sin(3 * times)
sigs3 = (5 * np.pi) * np.sin(5 * times)
sigs4 = (7 * np.pi) * np.sin(7 * times)
sigs5 = (9 * np.pi) * np.sin(9 * times)
sigs6 = sigs1 + sigs2 + sigs3 + sigs4 + sigs5
"""对原来SIGS6 进行逆向傅里叶变化"""
# 需要指定频率的长度 即 等于 sin函数的长度
# 需要指定频率的间隔 即 等于 sin函数的间隔
freqs = nf.fftfreq(times.size, times[1] - times[0])
# 将sigs6 通过傅里叶变化转化成一个复数，该复数就是波形震动和幅度关系的一个向量表示
ffts = nf.fft(sigs6)
# 然后由于我们要画图，而复数不好展示，所以通过ABS函数求（绝对值） 即 复数的 模长
pows = np.abs(ffts)
# 由于逆傅里叶变化时会产生负值，但是在物理意义上没有意义所以需要real只保留实际真值即可
sigs7 = nf.ifft(ffts).real
"""找到最大的 sigs 并进行逆向变化"""
# 找到pows中最大的值的下标，并对应freqs频率中的值从而得到该最大频率的函数波形
max_bo = freqs[pows.argmax()]
# 将傅里叶 复数 copy一份
filter_ffts = ffts
# 如果 频率中不等于最大值的->max_bo 都从 震动能量层面将其 设置为0
noised_indices = np.where(np.abs(freqs) != max_bo)
filter_ffts[noised_indices] = 0
# 然后逆向傅里叶得到正弦波
sigs8 = nf.ifft(filter_ffts).real
"""画图，信号图"""
mp.figure('FFT', facecolor='lightgray')
mp.subplot(121)
mp.title('Time Domain', fontsize=16)
mp.xlabel('Time', fontsize=12)
mp.ylabel('Signal', fontsize=12)
mp.tick_params(labelsize=10)
mp.grid(linestyle=':')
mp.plot(times, sigs1, label='{:.4f}'.format(
1 / (2 * np.pi)))
mp.plot(times, sigs2, label='{:.4f}'.format(
3 / (2 * np.pi)))
mp.plot(times, sigs3, label='{:.4f}'.format(
5 / (2 * np.pi)))
mp.plot(times, sigs4, label='{:.4f}'.format(
7 / (2 * np.pi)))
mp.plot(times, sigs5, label='{:.4f}'.format(
9 / (2 * np.pi)))
mp.plot(times, sigs6, label='{:.4f}'.format(
1 / (2 * np.pi)))
mp.plot(times, sigs7, label='{:.4f}'.format(
1 / (2 * np.pi)), alpha=0.5, linewidth=6)
mp.plot(times, sigs8, alpha=0.5, linewidth=5)
mp.legend()
mp.subplot(122)
"""画图频率图"""
mp.title('Frequency Domain', fontsize=16)
mp.xlabel('Frequency', fontsize=12)
mp.ylabel('Power', fontsize=12)
mp.tick_params(labelsize=10)
mp.grid(linestyle=':')
# 由于负数是没有物理意义的，所以我们只取 大于等于0的数值即可
# 取freqs 中 大于等于0 的频率值，映射到freqs 和 pows 上对应的点
mp.plot(freqs[freqs >= 0], pows[freqs >= 0],
c='orangered', label='Frequency Spectrum')
mp.legend()
mp.tight_layout()
mp.show()
