具有傅里叶分析的SpanSelector

时间:2015-05-19 09:17:08

标签: python matplotlib fft

我正在尝试使用matplotlib小部件SpanSelector以某些用户点击的间隔查找绘图的频率。我正在尝试使用matplotlib小部件SpanSelector来做这个,但我不确定如何做到这一点。我尝试修改给定http://matplotlib.org/examples/widgets/span_selector.html的示例 但它不起作用,仍然只放大选定的区域。这是我正在尝试使用的代码:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import SpanSelector

fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(211, axisbg='#FFFFCC')

def make_f(start,end):
    sum = 0
    for n in range(start,end+1):
        sum += np.cos(n * np.pi *time)
    return sum

time   = np.linspace(0,10,2000)
signal = make_f(1,10)

ax.plot(time, signal, '-')
#ax.set_ylim(-2,2)
ax.set_title('Press left mouse button and drag to test')


fourier = np.fft.rfft(signal)
n = signal.size
rate = time[1]-time[0]

ax2 = fig.add_subplot(212, axisbg='#FFFFCC')
line2, = ax2.plot(np.fft.rfftfreq(n,rate), 20*np.log10(fourier), '-')


def onselect(xmin, xmax):
    indmin, indmax = np.searchsorted(time, (xmin, xmax))
    indmax = min(len(time)-1, indmax)

    thisx = time[indmin:indmax]
    thisy = signal[indmin:indmax]
    line2.set_data(thisx, thisy)
    ax2.set_xlim(thisx[0], thisx[-1])
    ax2.set_ylim(thisy.min(), thisy.max())
    fig.canvas.draw()

# set useblit True on gtkagg for enhanced performance
def make_fourier():
    signal = SpanSelector(ax, onselect, 'horizontal', useblit=True,
                    rectprops=dict(alpha=0.5, facecolor='red') )
    fourier = np.fft.rfft(signal)
    n = signal.size
    rate = time[1]-time[0]
    return fourier

span = make_fourier()
plt.show()

1 个答案:

答案 0 :(得分:0)

我发现最好的方法是在onselect函数中进行更改。所以新的源代码将是:

# -*- coding: utf-8 -*-
"""
Created on Tue May 19 10:45:47 2015

@author: greenthumbtack
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import SpanSelector


fig,axarr = plt.subplots(2,3)
#fig = plt.figure(figsize=(10,8))

def make_f(start,end):
    sum = 0
    for n in range(start,end+1):
        sum += np.cos(2 * n * np.pi *time) + 2* np.random.random()
    return sum

time   = np.linspace(0,10,1001)
signal = make_f(1,10)

axarr[0,0].plot(time, signal, '-')
#ax.set_ylim(-2,2)
axarr[0,0].set_title('Press left mouse button and drag to test')


fourier = np.fft.rfft(signal)
n = signal.size
rate = time[1]-time[0]

axarr[1,0].plot(np.fft.rfftfreq(n,rate), 20*np.log10(fourier), '-')
axarr[0,1].plot(time, signal)
axarr[1,1].plot(np.fft.rfftfreq(n,rate), 20*np.log10(fourier), '-')
axarr[0,2].plot(time, signal)
axarr[1,2].plot(np.fft.rfftfreq(n,rate), 20*np.log10(fourier), '-')


def onselect(xmin, xmax):
    indmin, indmax = np.searchsorted(time, (xmin, xmax))
    indmax = min(len(time)-1, indmax)

    thisx = time[indmin:indmax]
    thisy = signal[indmin:indmax]
    thisy_fourier = np.fft.rfft(thisy)
    thisx_n = thisx.size
    rate = thisx[1] - thisx[0]
    axarr[1,0].clear()
    axarr[1,0].plot(np.fft.rfftfreq(thisx_n,rate), 20*np.log10(thisy_fourier))
    axarr[1,0].set_xlim(0, 11)
    #ax2.set_ylim(0, thisy.max())
    fig.canvas.draw()

def onselect2(xmin, xmax):
    indmin, indmax = np.searchsorted(time, (xmin, xmax))
    indmax = min(len(time)-1, indmax)

    thisx = time[indmin:indmax]
    thisy = signal[indmin:indmax]
    thisy_fourier = np.fft.rfft(thisy)
    thisx_n = thisx.size
    rate = thisx[1] - thisx[0]
    axarr[1,1].clear()
    axarr[1,1].plot(np.fft.rfftfreq(thisx_n,rate), 20*np.log10(thisy_fourier))
    axarr[1,1].set_xlim(0, 11)
    #ax2.set_ylim(0, thisy.max())
    fig.canvas.draw()

def onselect3(xmin, xmax):
    indmin, indmax = np.searchsorted(time, (xmin, xmax))
    indmax = min(len(time)-1, indmax)

    thisx = time[indmin:indmax]
    thisy = signal[indmin:indmax]
    thisy_fourier = np.fft.rfft(thisy)
    thisx_n = thisx.size
    rate = thisx[1] - thisx[0]
    axarr[1,2].clear()
    axarr[1,2].plot(np.fft.rfftfreq(thisx_n,rate), 20*np.log10(thisy_fourier))
    axarr[1,2].set_xlim(0, 11)
    #ax2.set_ylim(0, thisy.max())
    fig.canvas.draw()

span = SpanSelector(axarr[0,0], onselect, 'horizontal', useblit=True,
                    rectprops=dict(alpha=0.5, facecolor='red') )
span2 = SpanSelector(axarr[0,1], onselect2, 'horizontal', useblit=True,
                    rectprops=dict(alpha=0.5, facecolor='red') )
span3 = SpanSelector(axarr[0,2], onselect3, 'horizontal', useblit=True,
                    rectprops=dict(alpha=0.5, facecolor='red') )

我也做了一些其他的修改。它现在有一个2x3网格,其中有3个图表可以执行fft,每个都可以单独完成。一旦我意识到需要在onselect函数中进行任何修改,它就非常简单。