三角波的FFT演示 · 用 Python 做科學計算

# 三角波的FFT演示相關文檔： [_FFT演示程序_](fft_study.html) 本程序演示各種三角波形的FFT頻譜，用戶可以方便地修改三角波的各個參數，并立即看到其FFT頻譜的變化。 ![](_images/fft_study_04.swf) ``` # -*- coding: utf-8 -*- from enthought.traits.api import \ Str, Float, HasTraits, Property, cached_property, Range, Instance, on_trait_change, Enum from enthought.chaco.api import Plot, AbstractPlotData, ArrayPlotData, VPlotContainer from enthought.traits.ui.api import \ Item, View, VGroup, HSplit, ScrubberEditor, VSplit from enthought.enable.api import Component, ComponentEditor from enthought.chaco.tools.api import PanTool, ZoomTool import numpy as np # 鼠標拖動修改值的控件的樣式 scrubber = ScrubberEditor( hover_color = 0xFFFFFF, active_color = 0xA0CD9E, border_color = 0x808080 ) # 取FFT計算的結果freqs中的前n項進行合成，返回合成結果，計算loops個周期的波形 def fft_combine(freqs, n, loops=1): length = len(freqs) * loops data = np.zeros(length) index = loops * np.arange(0, length, 1.0) / length * (2 * np.pi) for k, p in enumerate(freqs[:n]): if k != 0: p *= 2 # 除去直流成分之外，其余的系數都*2 data += np.real(p) * np.cos(k*index) # 余弦成分的系數為實數部 data -= np.imag(p) * np.sin(k*index) # 正弦成分的系數為負的虛數部 return index, data class TriangleWave(HasTraits): # 指定三角波的最窄和最寬范圍，由于Range似乎不能將常數和traits名混用 # 所以定義這兩個不變的trait屬性 low = Float(0.02) hi = Float(1.0) # 三角波形的寬度 wave_width = Range("low", "hi", 0.5) # 三角波的頂點C的x軸坐標 length_c = Range("low", "wave_width", 0.5) # 三角波的定點的y軸坐標 height_c = Float(1.0) # FFT計算所使用的取樣點數，這里用一個Enum類型的屬性以供用戶從列表中選擇 fftsize = Enum( [(2**x) for x in range(6, 12)]) # FFT頻譜圖的x軸上限值 fft_graph_up_limit = Range(0, 400, 20) # 用于顯示FFT的結果 peak_list = Str # 采用多少個頻率合成三角波 N = Range(1, 40, 4) # 保存繪圖數據的對象 plot_data = Instance(AbstractPlotData) # 繪制波形圖的容器 plot_wave = Instance(Component) # 繪制FFT頻譜圖的容器 plot_fft = Instance(Component) # 包括兩個繪圖的容器 container = Instance(Component) # 設置用戶界面的視圖，注意一定要指定窗口的大小，這樣繪圖容器才能正常初始化 view = View( HSplit( VSplit( VGroup( Item("wave_width", editor = scrubber, label=u"波形寬度"), Item("length_c", editor = scrubber, label=u"最高點x坐標"), Item("height_c", editor = scrubber, label=u"最高點y坐標"), Item("fft_graph_up_limit", editor = scrubber, label=u"頻譜圖范圍"), Item("fftsize", label=u"FFT點數"), Item("N", label=u"合成波頻率數") ), Item("peak_list", style="custom", show_label=False, width=100, height=250) ), VGroup( Item("container", editor=ComponentEditor(size=(600,300)), show_label = False), orientation = "vertical" ) ), resizable = True, width = 800, height = 600, title = u"三角波FFT演示" ) # 創建繪圖的輔助函數，創建波形圖和頻譜圖有很多類似的地方，因此單獨用一個函數以 # 減少重復代碼 def _create_plot(self, data, name, type="line"): p = Plot(self.plot_data) p.plot(data, name=name, title=name, type=type) p.tools.append(PanTool(p)) zoom = ZoomTool(component=p, tool_mode="box", always_on=False) p.overlays.append(zoom) p.title = name return p def __init__(self): # 首先需要調用父類的初始化函數 super(TriangleWave, self).__init__() # 創建繪圖數據集，暫時沒有數據因此都賦值為空，只是創建幾個名字，以供Plot引用 self.plot_data = ArrayPlotData(x=[], y=[], f=[], p=[], x2=[], y2=[]) # 創建一個垂直排列的繪圖容器，它將頻譜圖和波形圖上下排列 self.container = VPlotContainer() # 創建波形圖，波形圖繪制兩條曲線：原始波形(x,y)和合成波形(x2,y2) self.plot_wave = self._create_plot(("x","y"), "Triangle Wave") self.plot_wave.plot(("x2","y2"), color="red") # 創建頻譜圖，使用數據集中的f和p self.plot_fft = self._create_plot(("f","p"), "FFT", type="scatter") # 將兩個繪圖容器添加到垂直容器中 self.container.add( self.plot_wave ) self.container.add( self.plot_fft ) # 設置 self.plot_wave.x_axis.title = "Samples" self.plot_fft.x_axis.title = "Frequency pins" self.plot_fft.y_axis.title = "(dB)" # 改變fftsize為1024，因為Enum的默認缺省值為枚舉列表中的第一個值 self.fftsize = 1024 # FFT頻譜圖的x軸上限值的改變事件處理函數，將最新的值賦值給頻譜圖的響應屬性 def _fft_graph_up_limit_changed(self): self.plot_fft.x_axis.mapper.range.high = self.fft_graph_up_limit def _N_changed(self): self.plot_sin_combine() # 多個trait屬性的改變事件處理函數相同時，可以用@on_trait_change指定 @on_trait_change("wave_width, length_c, height_c, fftsize") def update_plot(self): # 計算三角波 global y_data x_data = np.arange(0, 1.0, 1.0/self.fftsize) func = self.triangle_func() # 將func函數的返回值強制轉換成float64 y_data = np.cast["float64"](func(x_data)) # 計算頻譜 fft_parameters = np.fft.fft(y_data) / len(y_data) # 計算各個頻率的振幅 fft_data = np.clip(20*np.log10(np.abs(fft_parameters))[:self.fftsize/2+1], -120, 120) # 將計算的結果寫進數據集 self.plot_data.set_data("x", np.arange(0, self.fftsize)) # x坐標為取樣點 self.plot_data.set_data("y", y_data) self.plot_data.set_data("f", np.arange(0, len(fft_data))) # x坐標為頻率編號 self.plot_data.set_data("p", fft_data) # 合成波的x坐標為取樣點，顯示2個周期 self.plot_data.set_data("x2", np.arange(0, 2*self.fftsize)) # 更新頻譜圖x軸上限 self._fft_graph_up_limit_changed() # 將振幅大于-80dB的頻率輸出 peak_index = (fft_data > -80) peak_value = fft_data[peak_index][:20] result = [] for f, v in zip(np.flatnonzero(peak_index), peak_value): result.append("%s : %s" %(f, v) ) self.peak_list = "\n".join(result) # 保存現在的fft計算結果，并計算正弦合成波 self.fft_parameters = fft_parameters self.plot_sin_combine() # 計算正弦合成波，計算2個周期 def plot_sin_combine(self): index, data = fft_combine(self.fft_parameters, self.N, 2) self.plot_data.set_data("y2", data) # 返回一個ufunc計算指定參數的三角波 def triangle_func(self): c = self.wave_width c0 = self.length_c hc = self.height_c def trifunc(x): x = x - int(x) # 三角波的周期為1，因此只取x坐標的小數部分進行計算 if x >= c: r = 0.0 elif x < c0: r = x / c0 * hc else: r = (c-x) / (c-c0) * hc return r # 用trifunc函數創建一個ufunc函數，可以直接對數組進行計算, 不過通過此函數 # 計算得到的是一個Object數組，需要進行類型轉換 return np.frompyfunc(trifunc, 1, 1) if __name__ == "__main__": triangle = TriangleWave() triangle.configure_traits() ```