MatplotlibFourierDemo

   1 import numpy as np
   2 import wx
   3 
   4 import matplotlib
   5 matplotlib.interactive(False)
   6 matplotlib.use('WXAgg')
   7 from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg
   8 from matplotlib.figure import Figure
   9 from matplotlib.pyplot import gcf, setp
  10 
  11 
  12 class Knob:
  13     """
  14     Knob - simple class with a "setKnob" method.  
  15     A Knob instance is attached to a Param instance, e.g. param.attach(knob)
  16     Base class is for documentation purposes.
  17     """
  18     def setKnob(self, value):
  19         pass
  20 
  21 
  22 class Param:
  23     """
  24     The idea of the "Param" class is that some parameter in the GUI may have
  25     several knobs that both control it and reflect the parameter's state, e.g.
  26     a slider, text, and dragging can all change the value of the frequency in
  27     the waveform of this example.  
  28     The class allows a cleaner way to update/"feedback" to the other knobs when 
  29     one is being changed.  Also, this class handles min/max constraints for all
  30     the knobs.
  31     Idea - knob list - in "set" method, knob object is passed as well
  32       - the other knobs in the knob list have a "set" method which gets
  33         called for the others.
  34     """
  35     def __init__(self, initialValue=None, minimum=0., maximum=1.):
  36         self.minimum = minimum
  37         self.maximum = maximum
  38         if initialValue != self.constrain(initialValue):
  39             raise ValueError('illegal initial value')
  40         self.value = initialValue
  41         self.knobs = []
  42         
  43     def attach(self, knob):
  44         self.knobs += [knob]
  45         
  46     def set(self, value, knob=None):
  47         self.value = value
  48         self.value = self.constrain(value)
  49         for feedbackKnob in self.knobs:
  50             if feedbackKnob != knob:
  51                 feedbackKnob.setKnob(self.value)
  52         return self.value
  53 
  54     def constrain(self, value):
  55         if value <= self.minimum:
  56             value = self.minimum
  57         if value >= self.maximum:
  58             value = self.maximum
  59         return value
  60 
  61 
  62 class SliderGroup(Knob):
  63     def __init__(self, parent, label, param):
  64         self.sliderLabel = wx.StaticText(parent, label=label)
  65         self.sliderText = wx.TextCtrl(parent, -1, style=wx.TE_PROCESS_ENTER)
  66         self.slider = wx.Slider(parent, -1)
  67         self.slider.SetMax(param.maximum*1000)
  68         self.setKnob(param.value)
  69         
  70         sizer = wx.BoxSizer(wx.HORIZONTAL)
  71         sizer.Add(self.sliderLabel, 0, wx.EXPAND | wx.ALIGN_CENTER | wx.ALL, border=2)
  72         sizer.Add(self.sliderText, 0, wx.EXPAND | wx.ALIGN_CENTER | wx.ALL, border=2)
  73         sizer.Add(self.slider, 1, wx.EXPAND)
  74         self.sizer = sizer
  75 
  76         self.slider.Bind(wx.EVT_SLIDER, self.sliderHandler)
  77         self.sliderText.Bind(wx.EVT_TEXT_ENTER, self.sliderTextHandler)
  78 
  79         self.param = param
  80         self.param.attach(self)
  81 
  82     def sliderHandler(self, evt):
  83         value = evt.GetInt() / 1000.
  84         self.param.set(value)
  85         
  86     def sliderTextHandler(self, evt):
  87         value = float(self.sliderText.GetValue())
  88         self.param.set(value)
  89         
  90     def setKnob(self, value):
  91         self.sliderText.SetValue('%g'%value)
  92         self.slider.SetValue(value*1000)
  93 
  94 
  95 class FourierDemoFrame(wx.Frame):
  96     def __init__(self, *args, **kwargs):
  97         wx.Frame.__init__(self, *args, **kwargs)
  98 
  99         self.fourierDemoWindow = FourierDemoWindow(self)
 100         self.frequencySliderGroup = SliderGroup(self, label='Frequency f0:', \
 101             param=self.fourierDemoWindow.f0)
 102         self.amplitudeSliderGroup = SliderGroup(self, label=' Amplitude a:', \
 103             param=self.fourierDemoWindow.A)
 104 
 105         sizer = wx.BoxSizer(wx.VERTICAL)
 106         sizer.Add(self.fourierDemoWindow, 1, wx.EXPAND)
 107         sizer.Add(self.frequencySliderGroup.sizer, 0, \
 108             wx.EXPAND | wx.ALIGN_CENTER | wx.ALL, border=5)
 109         sizer.Add(self.amplitudeSliderGroup.sizer, 0, \
 110             wx.EXPAND | wx.ALIGN_CENTER | wx.ALL, border=5)
 111         self.SetSizer(sizer)
 112         
 113 
 114 class FourierDemoWindow(wx.Window, Knob):
 115     def __init__(self, *args, **kwargs):
 116         wx.Window.__init__(self, *args, **kwargs)
 117         self.lines = []
 118         self.figure = Figure()
 119         self.canvas = FigureCanvasWxAgg(self, -1, self.figure)
 120         self.canvas.callbacks.connect('button_press_event', self.mouseDown)
 121         self.canvas.callbacks.connect('motion_notify_event', self.mouseMotion)
 122         self.canvas.callbacks.connect('button_release_event', self.mouseUp)
 123         self.state = ''
 124         self.mouseInfo = (None, None, None, None)
 125         self.f0 = Param(2., minimum=0., maximum=6.)
 126         self.A = Param(1., minimum=0.01, maximum=2.)
 127         self.draw()
 128         
 129         # Not sure I like having two params attached to the same Knob,
 130         # but that is what we have here... it works but feels kludgy -
 131         # although maybe it's not too bad since the knob changes both params
 132         # at the same time (both f0 and A are affected during a drag)
 133         self.f0.attach(self)
 134         self.A.attach(self)
 135         self.Bind(wx.EVT_SIZE, self.sizeHandler)
 136        
 137     def sizeHandler(self, *args, **kwargs):
 138         self.canvas.SetSize(self.GetSize())
 139         
 140     def mouseDown(self, evt):
 141         if self.lines[0] in self.figure.hitlist(evt):
 142             self.state = 'frequency'
 143         elif self.lines[1] in self.figure.hitlist(evt):
 144             self.state = 'time'
 145         else:
 146             self.state = ''
 147         self.mouseInfo = (evt.xdata, evt.ydata, max(self.f0.value, .1), self.A.value)
 148 
 149     def mouseMotion(self, evt):
 150         if self.state == '':
 151             return
 152         x, y = evt.xdata, evt.ydata
 153         if x is None:  # outside the axes
 154             return
 155         x0, y0, f0Init, AInit = self.mouseInfo
 156         self.A.set(AInit+(AInit*(y-y0)/y0), self)
 157         if self.state == 'frequency':
 158             self.f0.set(f0Init+(f0Init*(x-x0)/x0))
 159         elif self.state == 'time':
 160             if (x-x0)/x0 != -1.:
 161                 self.f0.set(1./(1./f0Init+(1./f0Init*(x-x0)/x0)))
 162                     
 163     def mouseUp(self, evt):
 164         self.state = ''
 165 
 166     def draw(self):
 167         if not hasattr(self, 'subplot1'):
 168             self.subplot1 = self.figure.add_subplot(211)
 169             self.subplot2 = self.figure.add_subplot(212)
 170         x1, y1, x2, y2 = self.compute(self.f0.value, self.A.value)
 171         color = (1., 0., 0.)
 172         self.lines += self.subplot1.plot(x1, y1, color=color, linewidth=2)
 173         self.lines += self.subplot2.plot(x2, y2, color=color, linewidth=2)
 174         #Set some plot attributes
 175         self.subplot1.set_title("Click and drag waveforms to change frequency and amplitude", fontsize=12)
 176         self.subplot1.set_ylabel("Frequency Domain Waveform X(f)", fontsize = 8)
 177         self.subplot1.set_xlabel("frequency f", fontsize = 8)
 178         self.subplot2.set_ylabel("Time Domain Waveform x(t)", fontsize = 8)
 179         self.subplot2.set_xlabel("time t", fontsize = 8)
 180         self.subplot1.set_xlim([-6, 6])
 181         self.subplot1.set_ylim([0, 1])
 182         self.subplot2.set_xlim([-2, 2])
 183         self.subplot2.set_ylim([-2, 2])
 184         self.subplot1.text(0.05, .95, r'$X(f) = \mathcal{F}\{x(t)\}$', \
 185             verticalalignment='top', transform = self.subplot1.transAxes)
 186         self.subplot2.text(0.05, .95, r'$x(t) = a \cdot \cos(2\pi f_0 t) e^{-\pi t^2}$', \
 187             verticalalignment='top', transform = self.subplot2.transAxes)
 188 
 189     def compute(self, f0, A):
 190         f = np.arange(-6., 6., 0.02)
 191         t = np.arange(-2., 2., 0.01)
 192         x = A*np.cos(2*np.pi*f0*t)*np.exp(-np.pi*t**2)
 193         X = A/2*(np.exp(-np.pi*(f-f0)**2) + np.exp(-np.pi*(f+f0)**2))
 194         return f, X, t, x
 195 
 196     def repaint(self):
 197         self.canvas.draw()
 198 
 199     def setKnob(self, value):
 200         # Note, we ignore value arg here and just go by state of the params
 201         x1, y1, x2, y2 = self.compute(self.f0.value, self.A.value)
 202         setp(self.lines[0], xdata=x1, ydata=y1)
 203         setp(self.lines[1], xdata=x2, ydata=y2)
 204         self.repaint()
 205 
 206 
 207 class App(wx.App):
 208     def OnInit(self):
 209         self.frame1 = FourierDemoFrame(parent=None, title="Fourier Demo", size=(640, 480))
 210         self.frame1.Show()
 211         return True
 212         
 213 app = App()
 214 app.MainLoop()