9d97c4e3c5c1021bed1c5464848a2dcbef2b0a7f
[debian/gnuradio] / gr-wxgui / src / python / waterfallsink_nongl.py
1 #!/usr/bin/env python
2 #
3 # Copyright 2003,2004,2005,2007,2008 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 os
29 import threading
30 import math    
31
32 default_fftsink_size = (640,240)
33 default_fft_rate = gr.prefs().get_long('wxgui', 'fft_rate', 15)
34
35 class waterfall_sink_base(object):
36     def __init__(self, input_is_real=False, baseband_freq=0,
37                  sample_rate=1, fft_size=512,
38                  fft_rate=default_fft_rate,
39                  average=False, avg_alpha=None, title=''):
40
41         # initialize common attributes
42         self.baseband_freq = baseband_freq
43         self.sample_rate = sample_rate
44         self.fft_size = fft_size
45         self.fft_rate = fft_rate
46         self.average = average
47         if avg_alpha is None:
48             self.avg_alpha = 2.0 / fft_rate
49         else:
50             self.avg_alpha = avg_alpha
51         self.title = title
52         self.input_is_real = input_is_real
53         self.msgq = gr.msg_queue(2)         # queue up to 2 messages
54
55     def set_average(self, average):
56         self.average = average
57         if average:
58             self.avg.set_taps(self.avg_alpha)
59         else:
60             self.avg.set_taps(1.0)
61
62     def set_avg_alpha(self, avg_alpha):
63         self.avg_alpha = avg_alpha
64
65     def set_baseband_freq(self, baseband_freq):
66         self.baseband_freq = baseband_freq
67
68     def set_sample_rate(self, sample_rate):
69         self.sample_rate = sample_rate
70         self._set_n()
71
72     def _set_n(self):
73         self.one_in_n.set_n(max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
74         
75 class waterfall_sink_f(gr.hier_block2, waterfall_sink_base):
76     def __init__(self, parent, baseband_freq=0,
77                  y_per_div=10, ref_level=50, sample_rate=1, fft_size=512,
78                  fft_rate=default_fft_rate, average=False, avg_alpha=None,
79                  title='', size=default_fftsink_size):
80
81         gr.hier_block2.__init__(self, "waterfall_sink_f",
82                                 gr.io_signature(1, 1, gr.sizeof_float),
83                                 gr.io_signature(0,0,0))
84
85         waterfall_sink_base.__init__(self, input_is_real=True, baseband_freq=baseband_freq,
86                                sample_rate=sample_rate, fft_size=fft_size,
87                                fft_rate=fft_rate,
88                                average=average, avg_alpha=avg_alpha, title=title)
89                                
90         self.s2p = gr.serial_to_parallel(gr.sizeof_float, self.fft_size)
91         self.one_in_n = gr.keep_one_in_n(gr.sizeof_float * self.fft_size,
92                                          max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
93         
94         mywindow = window.blackmanharris(self.fft_size)
95         self.fft = gr.fft_vfc(self.fft_size, True, mywindow)
96         self.c2mag = gr.complex_to_mag(self.fft_size)
97         self.avg = gr.single_pole_iir_filter_ff(1.0, self.fft_size)
98         self.log = gr.nlog10_ff(20, self.fft_size, -20*math.log10(self.fft_size))
99         self.sink = gr.message_sink(gr.sizeof_float * self.fft_size, self.msgq, True)
100         self.connect(self, self.s2p, self.one_in_n, self.fft, self.c2mag, self.avg, self.log, self.sink)
101
102         self.win = waterfall_window(self, parent, size=size)
103         self.set_average(self.average)
104
105
106 class waterfall_sink_c(gr.hier_block2, waterfall_sink_base):
107     def __init__(self, parent, baseband_freq=0,
108                  y_per_div=10, ref_level=50, sample_rate=1, fft_size=512,
109                  fft_rate=default_fft_rate, average=False, avg_alpha=None, 
110                  title='', size=default_fftsink_size):
111
112         gr.hier_block2.__init__(self, "waterfall_sink_f",
113                                 gr.io_signature(1, 1, gr.sizeof_gr_complex),
114                                 gr.io_signature(0,0,0))
115
116         waterfall_sink_base.__init__(self, input_is_real=False, baseband_freq=baseband_freq,
117                                      sample_rate=sample_rate, fft_size=fft_size,
118                                      fft_rate=fft_rate,
119                                      average=average, avg_alpha=avg_alpha, title=title)
120
121         self.s2p = gr.serial_to_parallel(gr.sizeof_gr_complex, self.fft_size)
122         self.one_in_n = gr.keep_one_in_n(gr.sizeof_gr_complex * self.fft_size,
123                                          max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
124         
125         mywindow = window.blackmanharris(self.fft_size)
126         self.fft = gr.fft_vcc(self.fft_size, True, mywindow)
127         self.c2mag = gr.complex_to_mag(self.fft_size)
128         self.avg = gr.single_pole_iir_filter_ff(1.0, self.fft_size)
129         self.log = gr.nlog10_ff(20, self.fft_size, -20*math.log10(self.fft_size))
130         self.sink = gr.message_sink(gr.sizeof_float * self.fft_size, self.msgq, True)
131         self.connect(self, self.s2p, self.one_in_n, self.fft, self.c2mag, self.avg, self.log, self.sink)
132
133         self.win = waterfall_window(self, parent, size=size)
134         self.set_average(self.average)
135
136
137 # ------------------------------------------------------------------------
138
139 myDATA_EVENT = wx.NewEventType()
140 EVT_DATA_EVENT = wx.PyEventBinder (myDATA_EVENT, 0)
141
142
143 class DataEvent(wx.PyEvent):
144     def __init__(self, data):
145         wx.PyEvent.__init__(self)
146         self.SetEventType (myDATA_EVENT)
147         self.data = data
148
149     def Clone (self): 
150         self.__class__ (self.GetId())
151
152
153 class input_watcher (threading.Thread):
154     def __init__ (self, msgq, fft_size, event_receiver, **kwds):
155         threading.Thread.__init__ (self, **kwds)
156         self.setDaemon (1)
157         self.msgq = msgq
158         self.fft_size = fft_size
159         self.event_receiver = event_receiver
160         self.keep_running = True
161         self.start ()
162
163     def run (self):
164         while (self.keep_running):
165             msg = self.msgq.delete_head()  # blocking read of message queue
166             itemsize = int(msg.arg1())
167             nitems = int(msg.arg2())
168
169             s = msg.to_string()            # get the body of the msg as a string
170
171             # There may be more than one FFT frame in the message.
172             # If so, we take only the last one
173             if nitems > 1:
174                 start = itemsize * (nitems - 1)
175                 s = s[start:start+itemsize]
176
177             complex_data = numpy.fromstring (s, numpy.float32)
178             de = DataEvent (complex_data)
179             wx.PostEvent (self.event_receiver, de)
180             del de
181     
182
183 class waterfall_window (wx.Panel):
184     def __init__ (self, fftsink, parent, id = -1,
185                   pos = wx.DefaultPosition, size = wx.DefaultSize,
186                   style = wx.DEFAULT_FRAME_STYLE, name = ""):
187         wx.Panel.__init__(self, parent, id, pos, size, style, name)
188         self.set_baseband_freq = fftsink.set_baseband_freq
189         self.fftsink = fftsink
190         self.bm = wx.EmptyBitmap(self.fftsink.fft_size, 300, -1)
191
192         self.scale_factor = 5.0           # FIXME should autoscale, or set this
193         
194         dc1 = wx.MemoryDC()
195         dc1.SelectObject(self.bm)
196         dc1.Clear()
197
198         self.pens = self.make_pens()
199
200         wx.EVT_PAINT( self, self.OnPaint )
201         wx.EVT_CLOSE (self, self.on_close_window)
202         EVT_DATA_EVENT (self, self.set_data)
203         
204         self.build_popup_menu()
205         
206         wx.EVT_CLOSE (self, self.on_close_window)
207         self.Bind(wx.EVT_RIGHT_UP, self.on_right_click)
208
209         self.input_watcher = input_watcher(fftsink.msgq, fftsink.fft_size, self)
210
211
212     def on_close_window (self, event):
213         print "waterfall_window: on_close_window"
214         self.keep_running = False
215
216     def const_list(self,const,len):
217         return [const] * len
218
219     def make_colormap(self):
220         r = []
221         r.extend(self.const_list(0,96))
222         r.extend(range(0,255,4))
223         r.extend(self.const_list(255,64))
224         r.extend(range(255,128,-4))
225         
226         g = []
227         g.extend(self.const_list(0,32))
228         g.extend(range(0,255,4))
229         g.extend(self.const_list(255,64))
230         g.extend(range(255,0,-4))
231         g.extend(self.const_list(0,32))
232         
233         b = range(128,255,4)
234         b.extend(self.const_list(255,64))
235         b.extend(range(255,0,-4))
236         b.extend(self.const_list(0,96))
237         return (r,g,b)
238
239     def make_pens(self):
240         (r,g,b) = self.make_colormap()
241         pens = []
242         for i in range(0,256):
243             colour = wx.Colour(r[i], g[i], b[i])
244             pens.append( wx.Pen(colour, 2, wx.SOLID))
245         return pens
246         
247     def OnPaint(self, event):
248         dc = wx.PaintDC(self)
249         self.DoDrawing(dc)
250
251     def DoDrawing(self, dc=None):
252         if dc is None:
253             dc = wx.ClientDC(self)
254         dc.DrawBitmap(self.bm, 0, 0, False )
255     
256
257     def const_list(self,const,len):
258         a = [const]
259         for i in range(1,len):
260             a.append(const)
261         return a
262
263
264     def set_data (self, evt):
265         dB = evt.data
266         L = len (dB)
267
268         dc1 = wx.MemoryDC()
269         dc1.SelectObject(self.bm)
270         dc1.Blit(0,1,self.fftsink.fft_size,300,dc1,0,0,wx.COPY,False,-1,-1)
271
272         x = max(abs(self.fftsink.sample_rate), abs(self.fftsink.baseband_freq))
273         if x >= 1e9:
274             sf = 1e-9
275             units = "GHz"
276         elif x >= 1e6:
277             sf = 1e-6
278             units = "MHz"
279         else:
280             sf = 1e-3
281             units = "kHz"
282
283
284         if self.fftsink.input_is_real:     # only plot 1/2 the points
285             d_max = L/2
286             p_width = 2
287         else:
288             d_max = L/2
289             p_width = 1
290
291         scale_factor = self.scale_factor
292         if self.fftsink.input_is_real:     # real fft
293            for x_pos in range(0, d_max):
294                value = int(dB[x_pos] * scale_factor)
295                value = min(255, max(0, value))
296                dc1.SetPen(self.pens[value])
297                dc1.DrawRectangle(x_pos*p_width, 0, p_width, 2) 
298         else:                               # complex fft
299            for x_pos in range(0, d_max):    # positive freqs
300                value = int(dB[x_pos] * scale_factor)
301                value = min(255, max(0, value))
302                dc1.SetPen(self.pens[value])
303                dc1.DrawRectangle(x_pos*p_width + d_max, 0, p_width, 2) 
304            for x_pos in range(0 , d_max):   # negative freqs
305                value = int(dB[x_pos+d_max] * scale_factor)
306                value = min(255, max(0, value))
307                dc1.SetPen(self.pens[value])
308                dc1.DrawRectangle(x_pos*p_width, 0, p_width, 2) 
309
310         del dc1
311         self.DoDrawing (None)
312
313     def on_average(self, evt):
314         # print "on_average"
315         self.fftsink.set_average(evt.IsChecked())
316
317     def on_right_click(self, event):
318         menu = self.popup_menu
319         for id, pred in self.checkmarks.items():
320             item = menu.FindItemById(id)
321             item.Check(pred())
322         self.PopupMenu(menu, event.GetPosition())
323
324
325     def build_popup_menu(self):
326         self.id_incr_ref_level = wx.NewId()
327         self.id_decr_ref_level = wx.NewId()
328         self.id_incr_y_per_div = wx.NewId()
329         self.id_decr_y_per_div = wx.NewId()
330         self.id_y_per_div_1 = wx.NewId()
331         self.id_y_per_div_2 = wx.NewId()
332         self.id_y_per_div_5 = wx.NewId()
333         self.id_y_per_div_10 = wx.NewId()
334         self.id_y_per_div_20 = wx.NewId()
335         self.id_average = wx.NewId()
336
337         self.Bind(wx.EVT_MENU, self.on_average, id=self.id_average)
338         #self.Bind(wx.EVT_MENU, self.on_incr_ref_level, id=self.id_incr_ref_level)
339         #self.Bind(wx.EVT_MENU, self.on_decr_ref_level, id=self.id_decr_ref_level)
340         #self.Bind(wx.EVT_MENU, self.on_incr_y_per_div, id=self.id_incr_y_per_div)
341         #self.Bind(wx.EVT_MENU, self.on_decr_y_per_div, id=self.id_decr_y_per_div)
342         #self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_1)
343         #self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_2)
344         #self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_5)
345         #self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_10)
346         #self.Bind(wx.EVT_MENU, self.on_y_per_div, id=self.id_y_per_div_20)
347
348
349         # make a menu
350         menu = wx.Menu()
351         self.popup_menu = menu
352         menu.AppendCheckItem(self.id_average, "Average")
353         # menu.Append(self.id_incr_ref_level, "Incr Ref Level")
354         # menu.Append(self.id_decr_ref_level, "Decr Ref Level")
355         # menu.Append(self.id_incr_y_per_div, "Incr dB/div")
356         # menu.Append(self.id_decr_y_per_div, "Decr dB/div")
357         # menu.AppendSeparator()
358         # we'd use RadioItems for these, but they're not supported on Mac
359         #menu.AppendCheckItem(self.id_y_per_div_1, "1 dB/div")
360         #menu.AppendCheckItem(self.id_y_per_div_2, "2 dB/div")
361         #menu.AppendCheckItem(self.id_y_per_div_5, "5 dB/div")
362         #menu.AppendCheckItem(self.id_y_per_div_10, "10 dB/div")
363         #menu.AppendCheckItem(self.id_y_per_div_20, "20 dB/div")
364
365         self.checkmarks = {
366             self.id_average : lambda : self.fftsink.average
367             #self.id_y_per_div_1 : lambda : self.fftsink.y_per_div == 1,
368             #self.id_y_per_div_2 : lambda : self.fftsink.y_per_div == 2,
369             #self.id_y_per_div_5 : lambda : self.fftsink.y_per_div == 5,
370             #self.id_y_per_div_10 : lambda : self.fftsink.y_per_div == 10,
371             #self.id_y_per_div_20 : lambda : self.fftsink.y_per_div == 20,
372             }
373
374
375 def next_up(v, seq):
376     """
377     Return the first item in seq that is > v.
378     """
379     for s in seq:
380         if s > v:
381             return s
382     return v
383
384 def next_down(v, seq):
385     """
386     Return the last item in seq that is < v.
387     """
388     rseq = list(seq[:])
389     rseq.reverse()
390
391     for s in rseq:
392         if s < v:
393             return s
394     return v
395
396
397 # ----------------------------------------------------------------
398 # Standalone test app
399 # ----------------------------------------------------------------
400
401 class test_top_block (stdgui2.std_top_block):
402     def __init__(self, frame, panel, vbox, argv):
403         stdgui2.std_top_block.__init__ (self, frame, panel, vbox, argv)
404
405         fft_size = 512
406
407         # build our flow graph
408         input_rate = 20.000e3
409
410         # Generate a complex sinusoid
411         self.src1 = gr.sig_source_c (input_rate, gr.GR_SIN_WAVE, 5.75e3, 1000)
412         #src1 = gr.sig_source_c (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1000)
413
414         # We add these throttle blocks so that this demo doesn't
415         # suck down all the CPU available.  Normally you wouldn't use these.
416         self.thr1 = gr.throttle(gr.sizeof_gr_complex, input_rate)
417
418         sink1 = waterfall_sink_c (panel, title="Complex Data", fft_size=fft_size,
419                                   sample_rate=input_rate, baseband_freq=100e3)
420         self.connect(self.src1, self.thr1, sink1)
421         vbox.Add (sink1.win, 1, wx.EXPAND)
422
423         # generate a real sinusoid
424         self.src2 = gr.sig_source_f (input_rate, gr.GR_SIN_WAVE, 5.75e3, 1000)
425         self.thr2 = gr.throttle(gr.sizeof_float, input_rate)
426         sink2 = waterfall_sink_f (panel, title="Real Data", fft_size=fft_size,
427                                   sample_rate=input_rate, baseband_freq=100e3)
428         self.connect(self.src2, self.thr2, sink2)
429         vbox.Add (sink2.win, 1, wx.EXPAND)
430
431
432 def main ():
433     app = stdgui2.stdapp (test_top_block, "Waterfall Sink Test App")
434     app.MainLoop ()
435
436 if __name__ == '__main__':
437     main ()