Updated license from GPL version 2 or later to GPL version 3 or later.
[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         
235         EVT_DATA_EVENT (self, self.set_data)
236         wx.EVT_CLOSE (self, self.on_close_window)
237         self.Bind(wx.EVT_RIGHT_UP, self.on_right_click)
238
239         self.input_watcher = input_watcher(fftsink.msgq, fftsink.fft_size, self)
240
241
242     def on_close_window (self, event):
243         print "fft_window:on_close_window"
244         self.keep_running = False
245
246
247     def set_data (self, evt):
248         dB = evt.data
249         L = len (dB)
250
251         if self.peak_hold:
252             if self.peak_vals is None:
253                 self.peak_vals = dB
254             else:
255                 self.peak_vals = numpy.maximum(dB, self.peak_vals)
256                 dB = self.peak_vals
257
258         x = max(abs(self.fftsink.sample_rate), abs(self.fftsink.baseband_freq))
259         if x >= 1e9:
260             sf = 1e-9
261             units = "GHz"
262         elif x >= 1e6:
263             sf = 1e-6
264             units = "MHz"
265         else:
266             sf = 1e-3
267             units = "kHz"
268
269         if self.fftsink.input_is_real:     # only plot 1/2 the points
270             x_vals = ((numpy.arange (L/2)
271                        * (self.fftsink.sample_rate * sf / L))
272                       + self.fftsink.baseband_freq * sf)
273             points = numpy.zeros((len(x_vals), 2), numpy.float64)
274             points[:,0] = x_vals
275             points[:,1] = dB[0:L/2]
276         else:
277             # the "negative freqs" are in the second half of the array
278             x_vals = ((numpy.arange (-L/2, L/2)
279                        * (self.fftsink.sample_rate * sf / L))
280                       + self.fftsink.baseband_freq * sf)
281             points = numpy.zeros((len(x_vals), 2), numpy.float64)
282             points[:,0] = x_vals
283             points[:,1] = numpy.concatenate ((dB[L/2:], dB[0:L/2]))
284
285
286         lines = plot.PolyLine (points, colour='BLUE')
287
288         graphics = plot.PlotGraphics ([lines],
289                                       title=self.fftsink.title,
290                                       xLabel = units, yLabel = "dB")
291
292         self.Draw (graphics, xAxis=None, yAxis=self.y_range)
293         self.update_y_range ()
294
295     def set_peak_hold(self, enable):
296         self.peak_hold = enable
297         self.peak_vals = None
298
299     def update_y_range (self):
300         ymax = self.fftsink.ref_level
301         ymin = self.fftsink.ref_level - self.fftsink.y_per_div * self.fftsink.y_divs
302         self.y_range = self._axisInterval ('min', ymin, ymax)
303
304     def on_average(self, evt):
305         # print "on_average"
306         self.fftsink.set_average(evt.IsChecked())
307
308     def on_peak_hold(self, evt):
309         # print "on_peak_hold"
310         self.fftsink.set_peak_hold(evt.IsChecked())
311
312     def on_incr_ref_level(self, evt):
313         # print "on_incr_ref_level"
314         self.fftsink.set_ref_level(self.fftsink.ref_level
315                                    + self.fftsink.y_per_div)
316
317     def on_decr_ref_level(self, evt):
318         # print "on_decr_ref_level"
319         self.fftsink.set_ref_level(self.fftsink.ref_level
320                                    - self.fftsink.y_per_div)
321
322     def on_incr_y_per_div(self, evt):
323         # print "on_incr_y_per_div"
324         self.fftsink.set_y_per_div(next_up(self.fftsink.y_per_div, (1,2,5,10,20)))
325
326     def on_decr_y_per_div(self, evt):
327         # print "on_decr_y_per_div"
328         self.fftsink.set_y_per_div(next_down(self.fftsink.y_per_div, (1,2,5,10,20)))
329
330     def on_y_per_div(self, evt):
331         # print "on_y_per_div"
332         Id = evt.GetId()
333         if Id == self.id_y_per_div_1:
334             self.fftsink.set_y_per_div(1)
335         elif Id == self.id_y_per_div_2:
336             self.fftsink.set_y_per_div(2)
337         elif Id == self.id_y_per_div_5:
338             self.fftsink.set_y_per_div(5)
339         elif Id == self.id_y_per_div_10:
340             self.fftsink.set_y_per_div(10)
341         elif Id == self.id_y_per_div_20:
342             self.fftsink.set_y_per_div(20)
343
344         
345     def on_right_click(self, event):
346         menu = self.popup_menu
347         for id, pred in self.checkmarks.items():
348             item = menu.FindItemById(id)
349             item.Check(pred())
350         self.PopupMenu(menu, event.GetPosition())
351
352
353     def build_popup_menu(self):
354         self.id_incr_ref_level = wx.NewId()
355         self.id_decr_ref_level = wx.NewId()
356         self.id_incr_y_per_div = wx.NewId()
357         self.id_decr_y_per_div = wx.NewId()
358         self.id_y_per_div_1 = wx.NewId()
359         self.id_y_per_div_2 = wx.NewId()
360         self.id_y_per_div_5 = wx.NewId()
361         self.id_y_per_div_10 = wx.NewId()
362         self.id_y_per_div_20 = wx.NewId()
363         self.id_average = wx.NewId()
364         self.id_peak_hold = wx.NewId()
365
366         self.Bind(wx.EVT_MENU, self.on_average, id=self.id_average)
367         self.Bind(wx.EVT_MENU, self.on_peak_hold, id=self.id_peak_hold)
368         self.Bind(wx.EVT_MENU, self.on_incr_ref_level, id=self.id_incr_ref_level)
369         self.Bind(wx.EVT_MENU, self.on_decr_ref_level, id=self.id_decr_ref_level)
370         self.Bind(wx.EVT_MENU, self.on_incr_y_per_div, id=self.id_incr_y_per_div)
371         self.Bind(wx.EVT_MENU, self.on_decr_y_per_div, id=self.id_decr_y_per_div)
372         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_1)
373         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_2)
374         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_5)
375         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_10)
376         self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_20)
377
378
379         # make a menu
380         menu = wx.Menu()
381         self.popup_menu = menu
382         menu.AppendCheckItem(self.id_average, "Average")
383         menu.AppendCheckItem(self.id_peak_hold, "Peak Hold")
384         menu.Append(self.id_incr_ref_level, "Incr Ref Level")
385         menu.Append(self.id_decr_ref_level, "Decr Ref Level")
386         # menu.Append(self.id_incr_y_per_div, "Incr dB/div")
387         # menu.Append(self.id_decr_y_per_div, "Decr dB/div")
388         menu.AppendSeparator()
389         # we'd use RadioItems for these, but they're not supported on Mac
390         menu.AppendCheckItem(self.id_y_per_div_1, "1 dB/div")
391         menu.AppendCheckItem(self.id_y_per_div_2, "2 dB/div")
392         menu.AppendCheckItem(self.id_y_per_div_5, "5 dB/div")
393         menu.AppendCheckItem(self.id_y_per_div_10, "10 dB/div")
394         menu.AppendCheckItem(self.id_y_per_div_20, "20 dB/div")
395
396         self.checkmarks = {
397             self.id_average : lambda : self.fftsink.average,
398             self.id_peak_hold : lambda : self.fftsink.peak_hold,
399             self.id_y_per_div_1 : lambda : self.fftsink.y_per_div == 1,
400             self.id_y_per_div_2 : lambda : self.fftsink.y_per_div == 2,
401             self.id_y_per_div_5 : lambda : self.fftsink.y_per_div == 5,
402             self.id_y_per_div_10 : lambda : self.fftsink.y_per_div == 10,
403             self.id_y_per_div_20 : lambda : self.fftsink.y_per_div == 20,
404             }
405
406
407 def next_up(v, seq):
408     """
409     Return the first item in seq that is > v.
410     """
411     for s in seq:
412         if s > v:
413             return s
414     return v
415
416 def next_down(v, seq):
417     """
418     Return the last item in seq that is < v.
419     """
420     rseq = list(seq[:])
421     rseq.reverse()
422
423     for s in rseq:
424         if s < v:
425             return s
426     return v
427
428
429 # ----------------------------------------------------------------
430 # Standalone test app
431 # ----------------------------------------------------------------
432
433 class test_app_block (stdgui2.std_top_block):
434     def __init__(self, frame, panel, vbox, argv):
435         stdgui2.std_top_block.__init__ (self, frame, panel, vbox, argv)
436
437         fft_size = 256
438
439         # build our flow graph
440         input_rate = 20.48e3
441
442         # Generate a complex sinusoid
443         #src1 = gr.sig_source_c (input_rate, gr.GR_SIN_WAVE, 2e3, 1)
444         src1 = gr.sig_source_c (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1)
445
446         # We add these throttle blocks so that this demo doesn't
447         # suck down all the CPU available.  Normally you wouldn't use these.
448         thr1 = gr.throttle(gr.sizeof_gr_complex, input_rate)
449
450         sink1 = fft_sink_c (panel, title="Complex Data", fft_size=fft_size,
451                             sample_rate=input_rate, baseband_freq=100e3,
452                             ref_level=0, y_per_div=20)
453         vbox.Add (sink1.win, 1, wx.EXPAND)
454
455         self.connect(src1, thr1, sink1)
456
457         #src2 = gr.sig_source_f (input_rate, gr.GR_SIN_WAVE, 2e3, 1)
458         src2 = gr.sig_source_f (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1)
459         thr2 = gr.throttle(gr.sizeof_float, input_rate)
460         sink2 = fft_sink_f (panel, title="Real Data", fft_size=fft_size*2,
461                             sample_rate=input_rate, baseband_freq=100e3,
462                             ref_level=0, y_per_div=20)
463         vbox.Add (sink2.win, 1, wx.EXPAND)
464
465         self.connect(src2, thr2, sink2)
466
467 def main ():
468     app = stdgui2.stdapp (test_app_block, "FFT Sink Test App")
469     app.MainLoop ()
470
471 if __name__ == '__main__':
472     main ()