Updated license from GPL version 2 or later to GPL version 3 or later.
[debian/gnuradio] / gr-radio-astronomy / src / python / ra_waterfallsink.py
1 #!/usr/bin/env python
2 #
3 # Copyright 2003,2004,2005 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 stdgui
25 import wx
26 import gnuradio.wxgui.plot as plot
27 import Numeric
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 ra_waterfallsink_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='', ofunc=None, xydfunc=None, scaling=100):
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         self.ofunc = ofunc
48         self.xydfunc = xydfunc
49         self.scaling = scaling
50         if avg_alpha is None:
51             self.avg_alpha = 2.0 / fft_rate
52         else:
53             self.avg_alpha = avg_alpha
54         self.title = title
55         self.input_is_real = input_is_real
56         self.msgq = gr.msg_queue(2)         # queue up to 2 messages
57
58     def set_average(self, average):
59         self.average = average
60         if average:
61             self.avg.set_taps(self.avg_alpha)
62         else:
63             self.avg.set_taps(1.0)
64
65     def set_avg_alpha(self, avg_alpha):
66         self.avg_alpha = avg_alpha
67
68     def set_baseband_freq(self, baseband_freq):
69         self.baseband_freq = baseband_freq
70
71     def set_sample_rate(self, sample_rate):
72         self.sample_rate = sample_rate
73         self._set_n()
74
75     def set_scaling(self,scaling):
76         self.scaling = scaling
77         return
78
79     def _set_n(self):
80         self.one_in_n.set_n(max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
81         
82 class ra_waterfallsink_f(gr.hier_block, ra_waterfallsink_base):
83     def __init__(self, fg, parent, baseband_freq=0,
84                  y_per_div=10, ref_level=50, sample_rate=1, fft_size=512,
85                  fft_rate=default_fft_rate, average=False, avg_alpha=None,
86                  title='', scaling=100, size=default_fftsink_size, ofunc=None, xydfunc=None):
87
88         ra_waterfallsink_base.__init__(self, input_is_real=True, baseband_freq=baseband_freq,
89                                sample_rate=sample_rate, fft_size=fft_size,
90                                fft_rate=fft_rate,
91                                average=average, avg_alpha=avg_alpha, title=title, ofunc=ofunc, xydfunc=xydfunc, scaling=scaling)
92                                
93         s2p = gr.serial_to_parallel(gr.sizeof_float, self.fft_size)
94         self.one_in_n = gr.keep_one_in_n(gr.sizeof_float * self.fft_size,
95                                          max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
96         mywindow = window.blackmanharris(self.fft_size)
97         fft = gr.fft_vfc(self.fft_size, True, mywindow)
98         c2mag = gr.complex_to_mag(self.fft_size)
99         self.avg = gr.single_pole_iir_filter_ff(1.0, self.fft_size)
100         log = gr.nlog10_ff(20, self.fft_size, -20*math.log10(self.fft_size))
101         sink = gr.message_sink(gr.sizeof_float * self.fft_size, self.msgq, True)
102
103         fg.connect (s2p, self.one_in_n, fft, c2mag, self.avg, log, sink)
104         gr.hier_block.__init__(self, fg, s2p, sink)
105
106         self.win = ra_waterfall_window(self, parent, size=size)
107         self.set_average(self.average)
108
109
110 class ra_waterfallsink_c(gr.hier_block, ra_waterfallsink_base):
111     def __init__(self, fg, parent, baseband_freq=0,
112                  y_per_div=10, ref_level=50, sample_rate=1, fft_size=512,
113                  fft_rate=default_fft_rate, average=False, 
114                  avg_alpha=None, scaling=100,
115                  title='', size=default_fftsink_size, ofunc=None, xydfunc=None):
116
117         ra_waterfallsink_base.__init__(self, input_is_real=False, baseband_freq=baseband_freq,
118                                      sample_rate=sample_rate, fft_size=fft_size,
119                                      fft_rate=fft_rate,
120                                      average=average, 
121                                      avg_alpha=avg_alpha, 
122                                      title=title, ofunc=ofunc, 
123                                      xydfunc=xydfunc, scaling=scaling)
124
125         s2p = gr.serial_to_parallel(gr.sizeof_gr_complex, self.fft_size)
126         self.one_in_n = gr.keep_one_in_n(gr.sizeof_gr_complex * self.fft_size,
127                                          max(1, int(self.sample_rate/self.fft_size/self.fft_rate)))
128
129         mywindow = window.blackmanharris(self.fft_size)
130         fft = gr.fft_vcc(self.fft_size, True, mywindow)
131         c2mag = gr.complex_to_mag(self.fft_size)
132         self.avg = gr.single_pole_iir_filter_ff(1.0, self.fft_size)
133         log = gr.nlog10_ff(20, self.fft_size, -20*math.log10(self.fft_size))
134         sink = gr.message_sink(gr.sizeof_float * self.fft_size, self.msgq, True)
135
136         fg.connect(s2p, self.one_in_n, fft, c2mag, self.avg, log, sink)
137         gr.hier_block.__init__(self, fg, s2p, sink)
138
139         self.win = ra_waterfall_window(self, parent, size=size)
140         self.set_average(self.average)
141
142
143 # ------------------------------------------------------------------------
144
145 myDATA_EVENT = wx.NewEventType()
146 EVT_DATA_EVENT = wx.PyEventBinder (myDATA_EVENT, 0)
147
148
149 class DataEvent(wx.PyEvent):
150     def __init__(self, data):
151         wx.PyEvent.__init__(self)
152         self.SetEventType (myDATA_EVENT)
153         self.data = data
154
155     def Clone (self): 
156         self.__class__ (self.GetId())
157
158
159 class input_watcher (threading.Thread):
160     def __init__ (self, msgq, fft_size, event_receiver, **kwds):
161         threading.Thread.__init__ (self, **kwds)
162         self.setDaemon (1)
163         self.msgq = msgq
164         self.fft_size = fft_size
165         self.event_receiver = event_receiver
166         self.keep_running = True
167         self.start ()
168
169     def run (self):
170         while (self.keep_running):
171             msg = self.msgq.delete_head()  # blocking read of message queue
172             itemsize = int(msg.arg1())
173             nitems = int(msg.arg2())
174
175             s = msg.to_string()            # get the body of the msg as a string
176
177             # There may be more than one FFT frame in the message.
178             # If so, we take only the last one
179             if nitems > 1:
180                 start = itemsize * (nitems - 1)
181                 s = s[start:start+itemsize]
182
183             complex_data = Numeric.fromstring (s, Numeric.Float32)
184             de = DataEvent (complex_data)
185             wx.PostEvent (self.event_receiver, de)
186             del de
187     
188
189 class ra_waterfall_window (wx.Panel):
190     def __init__ (self, fftsink, parent, id = -1,
191                   pos = wx.DefaultPosition, size = wx.DefaultSize,
192                   style = wx.DEFAULT_FRAME_STYLE, name = ""):
193         wx.Panel.__init__(self, parent, id, pos, size, style, name)
194
195         self.fftsink = fftsink
196         self.bm = wx.EmptyBitmap(1024, 512, -1)
197
198         self.scale_factor = self.fftsink.scaling
199         
200         dc1 = wx.MemoryDC()
201         dc1.SelectObject(self.bm)
202         dc1.Clear()
203
204         self.pens = self.make_pens()
205
206         wx.EVT_PAINT( self, self.OnPaint )
207         wx.EVT_CLOSE (self, self.on_close_window)
208         EVT_DATA_EVENT (self, self.set_data)
209         
210         self.build_popup_menu()
211         
212         wx.EVT_CLOSE (self, self.on_close_window)
213         self.Bind(wx.EVT_RIGHT_UP, self.on_right_click)
214         self.Bind(wx.EVT_MOTION, self.on_motion)
215
216         self.input_watcher = input_watcher(fftsink.msgq, fftsink.fft_size, self)
217
218
219     def on_close_window (self, event):
220         print "ra_waterfall_window: on_close_window"
221         self.keep_running = False
222
223     def const_list(self,const,len):
224         return [const] * len
225
226     def make_colormap(self):
227         r = []
228         r.extend(self.const_list(0,96))
229         r.extend(range(0,255,4))
230         r.extend(self.const_list(255,64))
231         r.extend(range(255,128,-4))
232         
233         g = []
234         g.extend(self.const_list(0,32))
235         g.extend(range(0,255,4))
236         g.extend(self.const_list(255,64))
237         g.extend(range(255,0,-4))
238         g.extend(self.const_list(0,32))
239         
240         b = range(128,255,4)
241         b.extend(self.const_list(255,64))
242         b.extend(range(255,0,-4))
243         b.extend(self.const_list(0,96))
244         return (r,g,b)
245
246     def make_pens(self):
247         (r,g,b) = self.make_colormap()
248         pens = []
249         for i in range(0,256):
250             colour = wx.Colour(r[i], g[i], b[i])
251             pens.append( wx.Pen(colour, 2, wx.SOLID))
252         return pens
253         
254     def OnPaint(self, event):
255         dc = wx.PaintDC(self)
256         self.DoDrawing(dc)
257
258     def DoDrawing(self, dc=None):
259         if dc is None:
260             dc = wx.ClientDC(self)
261         dc.DrawBitmap(self.bm, 0, 0, False )
262     
263
264     def const_list(self,const,len):
265         a = [const]
266         for i in range(1,len):
267             a.append(const)
268         return a
269
270     def make_colormap(self):
271         r = []
272         r.extend(self.const_list(0,96))
273         r.extend(range(0,255,4))
274         r.extend(self.const_list(255,64))
275         r.extend(range(255,128,-4))
276         
277         g = []
278         g.extend(self.const_list(0,32))
279         g.extend(range(0,255,4))
280         g.extend(self.const_list(255,64))
281         g.extend(range(255,0,-4))
282         g.extend(self.const_list(0,32))
283         
284         b = range(128,255,4)
285         b.extend(self.const_list(255,64))
286         b.extend(range(255,0,-4))
287         b.extend(self.const_list(0,96))
288         return (r,g,b)
289
290     def set_data (self, evt):
291         dB = evt.data
292         L = len (dB)
293
294         if (self.fftsink.ofunc != None):
295             self.fftsink.ofunc (evt.data, L)
296
297         dc1 = wx.MemoryDC()
298         dc1.SelectObject(self.bm)
299         dc1.Blit(0,1,1024,512,dc1,0,0,wx.COPY,False,-1,-1)
300
301         x = max(abs(self.fftsink.sample_rate), abs(self.fftsink.baseband_freq))
302         if x >= 1e9:
303             sf = 1e-9
304             units = "GHz"
305         elif x >= 1e6:
306             sf = 1e-6
307             units = "MHz"
308         else:
309             sf = 1e-3
310             units = "kHz"
311
312
313         if self.fftsink.input_is_real:     # only plot 1/2 the points
314             d_max = L/2
315             p_width = 2
316         else:
317             d_max = L/2
318             p_width = 1
319
320         WATERFALL_WIDTH=1024
321         scale_factor = self.scale_factor
322         x_positions = Numeric.zeros(WATERFALL_WIDTH, Numeric.Float64)
323         y_values = Numeric.zeros(WATERFALL_WIDTH, Numeric.Float64)
324         x_scale = L / WATERFALL_WIDTH
325         x_scale = int(x_scale)
326         if self.fftsink.input_is_real:     # real fft
327            for x_pos in range(0, d_max):
328                value = int(dB[x_pos] * scale_factor)
329                value = min(255, max(0, value))
330                idx = int(x_pos/x_scale)
331                idx = min(WATERFALL_WIDTH-1,idx)
332                x_positions[idx] = idx
333                y_values[idx] = y_values[idx] + value
334         else:                               # complex fft
335            for x_pos in range(0, d_max):    # positive freqs
336                value = int(dB[x_pos] * scale_factor)
337                value = min(255, max(0, value))
338                idx = int((x_pos+d_max)/x_scale)
339                idx = min(WATERFALL_WIDTH-1,idx)
340                x_positions[idx] = idx
341                y_values[idx] = y_values[idx] + value
342            for x_pos in range(0 , d_max):   # negative freqs
343                value = int(dB[x_pos+d_max] * scale_factor)
344                value = min(255, max(0, value))
345                idx = int((x_pos)/x_scale)
346                idx = min(WATERFALL_WIDTH-1,idx)
347                x_positions[idx] = idx
348                y_values[idx] = y_values[idx] + value
349
350         for i in range(0,WATERFALL_WIDTH):
351             yv = y_values[i]/x_scale
352             yv = int(min(255,yv))
353             dc1.SetPen(self.pens[yv])
354             dc1.DrawRectangle(i*p_width, 0, p_width, 1)
355
356         self.DoDrawing (None)
357
358     def on_average(self, evt):
359         # print "on_average"
360         self.fftsink.set_average(evt.IsChecked())
361
362     def on_right_click(self, event):
363         menu = self.popup_menu
364         for id, pred in self.checkmarks.items():
365             item = menu.FindItemById(id)
366             item.Check(pred())
367         self.PopupMenu(menu, event.GetPosition())
368
369     def on_motion(self, event):
370         if not self.fftsink.xydfunc == None:
371             pos = event.GetPosition()
372             self.fftsink.xydfunc(pos)
373
374     def on_scaling(self, evt):
375         Id = evt.GetId()
376         if Id == self.id_scaling_100:
377             self.fftsink.set_scaling(100)
378         elif Id == self.id_scaling_150:
379             self.fftsink.set_scaling(150)
380         elif Id == self.id_scaling_200:
381             self.fftsink.set_scaling(200)
382         elif Id == self.id_scaling_250:
383             self.fftsink.set_scaling(250)
384         elif Id == self.id_scaling_300:
385             self.fftsink.set_scaling(300)
386
387     def build_popup_menu(self):
388         self.id_scaling_100 = wx.NewId()
389         self.id_scaling_150 = wx.NewId()
390         self.id_scaling_200 = wx.NewId()
391         self.id_scaling_250 = wx.NewId()
392         self.id_scaling_300 = wx.NewId()
393         self.id_average = wx.NewId()
394
395         self.Bind(wx.EVT_MENU, self.on_average, id=self.id_average)
396         self.Bind(wx.EVT_MENU, self.on_scaling, id=self.id_scaling_100)
397         self.Bind(wx.EVT_MENU, self.on_scaling, id=self.id_scaling_150)
398         self.Bind(wx.EVT_MENU, self.on_scaling, id=self.id_scaling_200)
399         self.Bind(wx.EVT_MENU, self.on_scaling, id=self.id_scaling_250)
400         self.Bind(wx.EVT_MENU, self.on_scaling, id=self.id_scaling_300)
401
402
403         # make a menu
404         menu = wx.Menu()
405         self.popup_menu = menu
406         menu.AppendCheckItem(self.id_average, "Average")
407         menu.AppendCheckItem(self.id_scaling_100, "100 scale factor")
408         menu.AppendCheckItem(self.id_scaling_150, "150 scale factor")
409         menu.AppendCheckItem(self.id_scaling_200, "200 scale factor")
410         menu.AppendCheckItem(self.id_scaling_250, "250 scale factor")
411         menu.AppendCheckItem(self.id_scaling_300, "300 scale factor")
412
413         self.checkmarks = {
414             self.id_average : lambda : self.fftsink.average,
415             self.id_scaling_100 : lambda : self.fftsink.scaling == 100,
416             self.id_scaling_150 : lambda : self.fftsink.scaling == 150,
417             self.id_scaling_200 : lambda : self.fftsink.scaling == 200,
418             self.id_scaling_250 : lambda : self.fftsink.scaling == 250,
419             self.id_scaling_300 : lambda : self.fftsink.scaling == 300,
420             }
421
422
423 def next_up(v, seq):
424     """
425     Return the first item in seq that is > v.
426     """
427     for s in seq:
428         if s > v:
429             return s
430     return v
431
432 def next_down(v, seq):
433     """
434     Return the last item in seq that is < v.
435     """
436     rseq = list(seq[:])
437     rseq.reverse()
438
439     for s in rseq:
440         if s < v:
441             return s
442     return v
443
444
445 # ----------------------------------------------------------------
446 #                     Deprecated interfaces
447 # ----------------------------------------------------------------
448
449 # returns (block, win).
450 #   block requires a single input stream of float
451 #   win is a subclass of wxWindow
452
453 def make_ra_waterfallsink_f(fg, parent, title, fft_size, input_rate):
454     
455     block = ra_waterfallsink_f(fg, parent, title=title, fft_size=fft_size,
456                              sample_rate=input_rate)
457     return (block, block.win)
458
459 # returns (block, win).
460 #   block requires a single input stream of gr_complex
461 #   win is a subclass of wxWindow
462
463 def make_ra_waterfallsink_c(fg, parent, title, fft_size, input_rate):
464     block = ra_waterfallsink_c(fg, parent, title=title, fft_size=fft_size,
465                              sample_rate=input_rate)
466     return (block, block.win)
467
468
469 # ----------------------------------------------------------------
470 # Standalone test app
471 # ----------------------------------------------------------------
472
473 class test_app_flow_graph (stdgui.gui_flow_graph):
474     def __init__(self, frame, panel, vbox, argv):
475         stdgui.gui_flow_graph.__init__ (self, frame, panel, vbox, argv)
476
477         fft_size = 512
478
479         # build our flow graph
480         input_rate = 20.000e3
481
482         # Generate a complex sinusoid
483         src1 = gr.sig_source_c (input_rate, gr.GR_SIN_WAVE, 5.75e3, 1000)
484         #src1 = gr.sig_source_c (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1000)
485
486         # We add these throttle blocks so that this demo doesn't
487         # suck down all the CPU available.  Normally you wouldn't use these.
488         thr1 = gr.throttle(gr.sizeof_gr_complex, input_rate)
489
490         sink1 = ra_waterfallsink_c (self, panel, title="Complex Data", fft_size=fft_size,
491                                   sample_rate=input_rate, baseband_freq=100e3)
492         vbox.Add (sink1.win, 1, wx.EXPAND)
493         self.connect (src1, thr1, sink1)
494
495         # generate a real sinusoid
496         src2 = gr.sig_source_f (input_rate, gr.GR_SIN_WAVE, 5.75e3, 1000)
497         #src2 = gr.sig_source_f (input_rate, gr.GR_CONST_WAVE, 5.75e3, 1000)
498         thr2 = gr.throttle(gr.sizeof_float, input_rate)
499         sink2 = ra_waterfallsink_f (self, panel, title="Real Data", fft_size=fft_size,
500                                   sample_rate=input_rate, baseband_freq=100e3)
501         vbox.Add (sink2.win, 1, wx.EXPAND)
502         self.connect (src2, thr2, sink2)
503
504 def main ():
505     app = stdgui.stdapp (test_app_flow_graph,
506                          "Waterfall Sink Test App")
507     app.MainLoop ()
508
509 if __name__ == '__main__':
510     main ()