Fix for working with peak hold in usrp_fft.py
[debian/gnuradio] / gr-wxgui / src / python / fftsink2.py
1 #!/usr/bin/env python
2 #
3 # Copyright 2003,2004,2005,2006,2007 Free Software Foundation, Inc.
4
5 # This file is part of GNU Radio
6
7 # GNU Radio is free software; you can redistribute it and/or modify
8 # it under the terms of the GNU General Public License as published by
9 # the Free Software Foundation; either version 3, or (at your option)
10 # any later version.
11
12 # GNU Radio is distributed in the hope that it will be useful,
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 # GNU General Public License for more details.
16
17 # You should have received a copy of the GNU General Public License
18 # along with GNU Radio; see the file COPYING.  If not, write to
19 # the Free Software Foundation, Inc., 51 Franklin Street,
20 # Boston, MA 02110-1301, USA.
21
22
23 from gnuradio import gr, gru, window
24 from gnuradio.wxgui import stdgui2
25 import wx
26 import gnuradio.wxgui.plot as plot
27 import numpy
28 import threading
29 import math    
30
31 default_fftsink_size = (640,240)
32 default_fft_rate = gr.prefs().get_long('wxgui', 'fft_rate', 15)
33
34 class fft_sink_base(object):
35     def __init__(self, input_is_real=False, baseband_freq=0, y_per_div=10, ref_level=50,
36                  sample_rate=1, fft_size=512,
37                  fft_rate=default_fft_rate,
38                  average=False, avg_alpha=None, title='', peak_hold=False):
39
40         # initialize common attributes
41         self.baseband_freq = baseband_freq
42         self.y_divs = 8
43         self.y_per_div=y_per_div
44         self.ref_level = ref_level
45         self.sample_rate = sample_rate
46         self.fft_size = fft_size
47         self.fft_rate = fft_rate
48         self.average = average
49         if avg_alpha is None:
50             self.avg_alpha = 2.0 / fft_rate
51         else:
52             self.avg_alpha = avg_alpha
53         self.title = title
54         self.peak_hold = peak_hold
55         self.input_is_real = input_is_real
56         self.msgq = gr.msg_queue(2)         # queue that holds a maximum of 2 messages
57
58     def set_y_per_div(self, y_per_div):
59         self.y_per_div = y_per_div
60
61     def set_ref_level(self, ref_level):
62         self.ref_level = ref_level
63
64     def set_average(self, average):
65         self.average = average
66         if average:
67             self.avg.set_taps(self.avg_alpha)
68             self.set_peak_hold(False)
69         else:
70             self.avg.set_taps(1.0)
71
72     def set_peak_hold(self, enable):
73         self.peak_hold = enable
74         if enable:
75             self.set_average(False)
76         self.win.set_peak_hold(enable)
77
78     def set_avg_alpha(self, avg_alpha):
79         self.avg_alpha = avg_alpha
80
81     def set_baseband_freq(self, baseband_freq):
82         self.baseband_freq = baseband_freq
83
84     def set_sample_rate(self, sample_rate):
85         self.sample_rate = sample_rate
86         self._set_n()
87
88     def _set_n(self):
89         self.one_in_n.set_n(max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
90         
91
92 class fft_sink_f(gr.hier_block2, fft_sink_base):
93     def __init__(self, parent, baseband_freq=0,
94                  y_per_div=10, ref_level=50, sample_rate=1, fft_size=512,
95                  fft_rate=default_fft_rate, average=False, avg_alpha=None,
96                  title='', size=default_fftsink_size, peak_hold=False):
97
98         gr.hier_block2.__init__(self, "fft_sink_f",
99                                 gr.io_signature(1, 1, gr.sizeof_float),
100                                 gr.io_signature(0,0,0))
101
102         fft_sink_base.__init__(self, input_is_real=True, baseband_freq=baseband_freq,
103                                y_per_div=y_per_div, ref_level=ref_level,
104                                sample_rate=sample_rate, fft_size=fft_size,
105                                fft_rate=fft_rate,
106                                average=average, avg_alpha=avg_alpha, title=title,
107                                peak_hold=peak_hold)
108                                
109         self.s2p = gr.stream_to_vector(gr.sizeof_float, self.fft_size)
110         self.one_in_n = gr.keep_one_in_n(gr.sizeof_float * self.fft_size,
111                                          max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
112         
113         mywindow = window.blackmanharris(self.fft_size)
114         self.fft = gr.fft_vfc(self.fft_size, True, mywindow)
115         power = 0
116         for tap in mywindow:
117             power += tap*tap
118             
119         self.c2mag = gr.complex_to_mag(self.fft_size)
120         self.avg = gr.single_pole_iir_filter_ff(1.0, self.fft_size)
121
122         # FIXME  We need to add 3dB to all bins but the DC bin
123         self.log = gr.nlog10_ff(20, self.fft_size,
124                                -20*math.log10(self.fft_size)-10*math.log10(power/self.fft_size))
125         self.sink = gr.message_sink(gr.sizeof_float * self.fft_size, self.msgq, True)
126         self.connect(self, self.s2p, self.one_in_n, self.fft, self.c2mag, self.avg, self.log, self.sink)
127
128         self.win = fft_window(self, parent, size=size)
129         self.set_average(self.average)
130
131
132 class fft_sink_c(gr.hier_block2, fft_sink_base):
133     def __init__(self, parent, baseband_freq=0,
134                  y_per_div=10, ref_level=50, sample_rate=1, fft_size=512,
135                  fft_rate=default_fft_rate, average=False, avg_alpha=None,
136                  title='', size=default_fftsink_size, peak_hold=False):
137
138         gr.hier_block2.__init__(self, "fft_sink_c",
139                                 gr.io_signature(1, 1, gr.sizeof_gr_complex),
140                                 gr.io_signature(0,0,0))
141
142         fft_sink_base.__init__(self, input_is_real=False, baseband_freq=baseband_freq,
143                                y_per_div=y_per_div, ref_level=ref_level,
144                                sample_rate=sample_rate, fft_size=fft_size,
145                                fft_rate=fft_rate,
146                                average=average, avg_alpha=avg_alpha, title=title,
147                                peak_hold=peak_hold)
148
149         self.s2p = gr.stream_to_vector(gr.sizeof_gr_complex, self.fft_size)
150         self.one_in_n = gr.keep_one_in_n(gr.sizeof_gr_complex * self.fft_size,
151                                          max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
152         
153         mywindow = window.blackmanharris(self.fft_size)
154         self.fft = gr.fft_vcc(self.fft_size, True, mywindow)
155         power = 0
156         for tap in mywindow:
157             power += tap*tap
158             
159         self.c2mag = gr.complex_to_mag(self.fft_size)
160         self.avg = gr.single_pole_iir_filter_ff(1.0, self.fft_size)
161
162         # FIXME  We need to add 3dB to all bins but the DC bin
163         self.log = gr.nlog10_ff(20, self.fft_size,
164                                 -20*math.log10(self.fft_size)-10*math.log10(power/self.fft_size))
165         self.sink = gr.message_sink(gr.sizeof_float * self.fft_size, self.msgq, True)
166         self.connect(self, self.s2p, self.one_in_n, self.fft, self.c2mag, self.avg, self.log, self.sink)
167
168         self.win = fft_window(self, parent, size=size)
169         self.set_average(self.average)
170
171
172 # ------------------------------------------------------------------------
173
174 myDATA_EVENT = wx.NewEventType()
175 EVT_DATA_EVENT = wx.PyEventBinder (myDATA_EVENT, 0)
176
177
178 class DataEvent(wx.PyEvent):
179     def __init__(self, data):
180         wx.PyEvent.__init__(self)
181         self.SetEventType (myDATA_EVENT)
182         self.data = data
183
184     def Clone (self): 
185         self.__class__ (self.GetId())
186
187
188 class input_watcher (threading.Thread):
189     def __init__ (self, msgq, fft_size, event_receiver, **kwds):
190         threading.Thread.__init__ (self, **kwds)
191         self.setDaemon (1)
192         self.msgq = msgq
193         self.fft_size = fft_size
194         self.event_receiver = event_receiver
195         self.keep_running = True
196         self.start ()
197
198     def run (self):
199         while (self.keep_running):
200             msg = self.msgq.delete_head()  # blocking read of message queue
201             itemsize = int(msg.arg1())
202             nitems = int(msg.arg2())
203
204             s = msg.to_string()            # get the body of the msg as a string
205
206             # There may be more than one FFT frame in the message.
207             # If so, we take only the last one
208             if nitems > 1:
209                 start = itemsize * (nitems - 1)
210                 s = s[start:start+itemsize]
211
212             complex_data = numpy.fromstring (s, numpy.float32)
213             de = DataEvent (complex_data)
214             wx.PostEvent (self.event_receiver, de)
215             del de
216     
217
218 class fft_window (plot.PlotCanvas):
219     def __init__ (self, fftsink, parent, id = -1,
220                   pos = wx.DefaultPosition, size = wx.DefaultSize,
221                   style = wx.DEFAULT_FRAME_STYLE, name = ""):
222         plot.PlotCanvas.__init__ (self, parent, id, pos, size, style, name)
223
224         self.y_range = None
225         self.fftsink = fftsink
226         self.peak_hold = False
227         self.peak_vals = None
228         
229         self.SetEnableGrid (True)
230         # self.SetEnableZoom (True)
231         # self.SetBackgroundColour ('black')
232         
233         self.build_popup_menu()
234         self.set_baseband_freq(0.0)
235                 
236         EVT_DATA_EVENT (self, self.set_data)
237         wx.EVT_CLOSE (self, self.on_close_window)
238         self.Bind(wx.EVT_RIGHT_UP, self.on_right_click)
239         self.Bind(wx.EVT_MOTION, self.evt_motion)
240         
241         self.input_watcher = input_watcher(fftsink.msgq, fftsink.fft_size, self)
242
243     def set_scale(self, freq):
244         x = max(abs(self.fftsink.sample_rate), abs(self.fftsink.baseband_freq)) 
245         if x >= 1e9:
246             self._scale_factor = 1e-9
247             self._units = "GHz"
248             self._format = "%3.6f"
249         elif x >= 1e6:
250             self._scale_factor = 1e-6
251             self._units = "MHz"
252             self._format = "%3.3f"
253         else:
254             self._scale_factor = 1e-3
255             self._units = "kHz"
256             self._format = "%3.3f"
257
258     def set_baseband_freq(self, baseband_freq):
259         if self.peak_hold:
260             self.peak_vals = None
261         self.set_scale(baseband_freq)
262         self.fftsink.set_baseband_freq(baseband_freq)
263         
264     def on_close_window (self, event):
265         print "fft_window:on_close_window"
266         self.keep_running = False
267
268
269     def set_data (self, evt):
270         dB = evt.data
271         L = len (dB)
272
273         if self.peak_hold:
274             if self.peak_vals is None:
275                 self.peak_vals = dB
276             else:
277                 self.peak_vals = numpy.maximum(dB, self.peak_vals)
278                 dB = self.peak_vals
279
280         if self.fftsink.input_is_real:     # only plot 1/2 the points
281             x_vals = ((numpy.arange (L/2) * (self.fftsink.sample_rate 
282                        * self._scale_factor / L))
283                       + self.fftsink.baseband_freq * self._scale_factor)
284             self._points = numpy.zeros((len(x_vals), 2), numpy.float64)
285             self._points[:,0] = x_vals
286             self._points[:,1] = dB[0:L/2]
287         else:
288             # the "negative freqs" are in the second half of the array
289             x_vals = ((numpy.arange (-L/2, L/2)
290                        * (self.fftsink.sample_rate * self._scale_factor / L))
291                       + self.fftsink.baseband_freq * self._scale_factor)
292             self._points = numpy.zeros((len(x_vals), 2), numpy.float64)
293             self._points[:,0] = x_vals
294             self._points[:,1] = numpy.concatenate ((dB[L/2:], dB[0:L/2]))
295
296         lines = plot.PolyLine (self._points, colour='BLUE')
297         graphics = plot.PlotGraphics ([lines],
298                                       title=self.fftsink.title,
299                                       xLabel = self._units, yLabel = "dB")
300
301         self.Draw (graphics, xAxis=None, yAxis=self.y_range)
302         self.update_y_range ()
303
304
305     def set_peak_hold(self, enable):
306         self.peak_hold = enable
307         self.peak_vals = None
308
309     def update_y_range (self):
310         ymax = self.fftsink.ref_level
311         ymin = self.fftsink.ref_level - self.fftsink.y_per_div * self.fftsink.y_divs
312         self.y_range = self._axisInterval ('min', ymin, ymax)
313
314     def on_average(self, evt):
315         # print "on_average"
316         self.fftsink.set_average(evt.IsChecked())
317
318     def on_peak_hold(self, evt):
319         # print "on_peak_hold"
320         self.fftsink.set_peak_hold(evt.IsChecked())
321
322     def on_incr_ref_level(self, evt):
323         # print "on_incr_ref_level"
324         self.fftsink.set_ref_level(self.fftsink.ref_level
325                                    + self.fftsink.y_per_div)
326
327     def on_decr_ref_level(self, evt):
328         # print "on_decr_ref_level"
329         self.fftsink.set_ref_level(self.fftsink.ref_level
330                                    - self.fftsink.y_per_div)
331
332     def on_incr_y_per_div(self, evt):
333         # print "on_incr_y_per_div"
334         self.fftsink.set_y_per_div(next_up(self.fftsink.y_per_div, (1,2,5,10,20)))
335
336     def on_decr_y_per_div(self, evt):
337         # print "on_decr_y_per_div"
338         self.fftsink.set_y_per_div(next_down(self.fftsink.y_per_div, (1,2,5,10,20)))
339
340     def on_y_per_div(self, evt):
341         # print "on_y_per_div"
342         Id = evt.GetId()
343         if Id == self.id_y_per_div_1:
344             self.fftsink.set_y_per_div(1)
345         elif Id == self.id_y_per_div_2:
346             self.fftsink.set_y_per_div(2)
347         elif Id == self.id_y_per_div_5:
348             self.fftsink.set_y_per_div(5)
349         elif Id == self.id_y_per_div_10:
350             self.fftsink.set_y_per_div(10)
351         elif Id == self.id_y_per_div_20:
352             self.fftsink.set_y_per_div(20)
353
354     def on_right_click(self, event):
355         menu = self.popup_menu
356         for id, pred in self.checkmarks.items():
357             item = menu.FindItemById(id)
358             item.Check(pred())
359         self.PopupMenu(menu, event.GetPosition())
360
361     def evt_motion(self, event):
362         # Clip to plotted values
363         (ux, uy) = self.GetXY(event)      # Scaled position
364         x_vals = numpy.array(self._points[:,0])
365         if ux < x_vals[0] or ux > x_vals[-1]:
366             tip = self.GetToolTip()
367             if tip:
368                 tip.Enable(False)
369             return
370
371         # Get nearest X value (is there a better way)?
372         ind = numpy.argmin(numpy.abs(x_vals-ux))
373         x_val = x_vals[ind]
374         db_val = self._points[ind, 1]
375         text = (self._format+" %s dB=%3.3f") % (x_val, self._units, db_val)
376
377         # Display the tooltip
378         tip = wx.ToolTip(text)
379         tip.Enable(True)
380         tip.SetDelay(0)
381         self.SetToolTip(tip)
382         
383     def build_popup_menu(self):
384         self.id_incr_ref_level = wx.NewId()
385         self.id_decr_ref_level = wx.NewId()
386         self.id_incr_y_per_div = wx.NewId()
387         self.id_decr_y_per_div = wx.NewId()
388         self.id_y_per_div_1 = wx.NewId()
389         self.id_y_per_div_2 = wx.NewId()
390         self.id_y_per_div_5 = wx.NewId()
391         self.id_y_per_div_10 = wx.NewId()
392         self.id_y_per_div_20 = wx.NewId()
393         self.id_average = wx.NewId()
394         self.id_peak_hold = wx.NewId()
395         
396         self.Bind(wx.EVT_MENU, self.on_average, id=self.id_average)
397         self.Bind(wx.EVT_MENU, self.on_peak_hold, id=self.id_peak_hold)
398         self.Bind(wx.EVT_MENU, self.on_incr_ref_level, id=self.id_incr_ref_level)
399         self.Bind(wx.EVT_MENU, self.on_decr_ref_level, id=self.id_decr_ref_level)
400         self.Bind(wx.EVT_MENU, self.on_incr_y_per_div, id=self.id_incr_y_per_div)
401         self.Bind(wx.EVT_MENU, self.on_decr_y_per_div, id=self.id_decr_y_per_div)
402         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_1)
403         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_2)
404         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_5)
405         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_10)
406         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_20)
407         
408         # make a menu
409         menu = wx.Menu()
410         self.popup_menu = menu
411         menu.AppendCheckItem(self.id_average, "Average")
412         menu.AppendCheckItem(self.id_peak_hold, "Peak Hold")
413         menu.Append(self.id_incr_ref_level, "Incr Ref Level")
414         menu.Append(self.id_decr_ref_level, "Decr Ref Level")
415         # menu.Append(self.id_incr_y_per_div, "Incr dB/div")
416         # menu.Append(self.id_decr_y_per_div, "Decr dB/div")
417         menu.AppendSeparator()
418         # we'd use RadioItems for these, but they're not supported on Mac
419         menu.AppendCheckItem(self.id_y_per_div_1, "1 dB/div")
420         menu.AppendCheckItem(self.id_y_per_div_2, "2 dB/div")
421         menu.AppendCheckItem(self.id_y_per_div_5, "5 dB/div")
422         menu.AppendCheckItem(self.id_y_per_div_10, "10 dB/div")
423         menu.AppendCheckItem(self.id_y_per_div_20, "20 dB/div")
424
425         self.checkmarks = {
426             self.id_average : lambda : self.fftsink.average,
427             self.id_peak_hold : lambda : self.fftsink.peak_hold,
428             self.id_y_per_div_1 : lambda : self.fftsink.y_per_div == 1,
429             self.id_y_per_div_2 : lambda : self.fftsink.y_per_div == 2,
430             self.id_y_per_div_5 : lambda : self.fftsink.y_per_div == 5,
431             self.id_y_per_div_10 : lambda : self.fftsink.y_per_div == 10,
432             self.id_y_per_div_20 : lambda : self.fftsink.y_per_div == 20,
433             }
434
435
436 def next_up(v, seq):
437     """
438     Return the first item in seq that is > v.
439     """
440     for s in seq:
441         if s > v:
442             return s
443     return v
444
445 def next_down(v, seq):
446     """
447     Return the last item in seq that is < v.
448     """
449     rseq = list(seq[:])
450     rseq.reverse()
451
452     for s in rseq:
453         if s < v:
454             return s
455     return v
456
457
458 # ----------------------------------------------------------------
459 # Standalone test app
460 # ----------------------------------------------------------------
461
462 class test_app_block (stdgui2.std_top_block):
463     def __init__(self, frame, panel, vbox, argv):
464         stdgui2.std_top_block.__init__ (self, frame, panel, vbox, argv)
465
466         fft_size = 256
467
468         # build our flow graph
469         input_rate = 20.48e3
470
471         # Generate a complex sinusoid
472         #src1 = gr.sig_source_c (input_rate, gr.GR_SIN_WAVE, 2e3, 1)
473         src1 = gr.sig_source_c (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1)
474
475         # We add these throttle blocks so that this demo doesn't
476         # suck down all the CPU available.  Normally you wouldn't use these.
477         thr1 = gr.throttle(gr.sizeof_gr_complex, input_rate)
478
479         sink1 = fft_sink_c (panel, title="Complex Data", fft_size=fft_size,
480                             sample_rate=input_rate, baseband_freq=100e3,
481                             ref_level=0, y_per_div=20)
482         vbox.Add (sink1.win, 1, wx.EXPAND)
483
484         self.connect(src1, thr1, sink1)
485
486         #src2 = gr.sig_source_f (input_rate, gr.GR_SIN_WAVE, 2e3, 1)
487         src2 = gr.sig_source_f (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1)
488         thr2 = gr.throttle(gr.sizeof_float, input_rate)
489         sink2 = fft_sink_f (panel, title="Real Data", fft_size=fft_size*2,
490                             sample_rate=input_rate, baseband_freq=100e3,
491                             ref_level=0, y_per_div=20)
492         vbox.Add (sink2.win, 1, wx.EXPAND)
493
494         self.connect(src2, thr2, sink2)
495
496 def main ():
497     app = stdgui2.stdapp (test_app_block, "FFT Sink Test App")
498     app.MainLoop ()
499
500 if __name__ == '__main__':
501     main ()