Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Newer
Older
100644 162 lines (125 sloc) 5.063 kb
7218059b » Roger
2012-05-05 add waterfall demo
1 from __future__ import division
2 import matplotlib.animation as animation
3 from matplotlib.mlab import psd
4 import pylab as pyl
5 import numpy as np
e11ad522 »
2012-05-21 fix demo for mac osx
6 import sys
7218059b » Roger
2012-05-05 add waterfall demo
7 from rtlsdr import RtlSdr
8
9 # A simple waterfall, spectrum plotter
10 #
11 # Controls:
12 #
13 # * Scroll mouse-wheel up and down to change the center frequency (hold shift for
14 # finer control)
15 # * Press "+" and "-" to control gain, and space to enable AGC.
16 # * Type a frequency (in MHz) and press enter to directly change the center frequency
17
18 NFFT = 1024*4
19 NUM_SAMPLES_PER_SCAN = NFFT*16
20 NUM_BUFFERED_SWEEPS = 100
21
22 # change this to control the number of scans that are combined in a single sweep
23 # (e.g. 2, 3, 4, etc.) Note that it can slow things down
24 NUM_SCANS_PER_SWEEP = 1
25
26 # these are the increments when scrolling the mouse wheel or pressing '+' or '-'
27 FREQ_INC_COARSE = 1e6
28 FREQ_INC_FINE = 0.1e6
29 GAIN_INC = 5
30
31 class Waterfall(object):
32 keyboard_buffer = []
33 shift_key_down = False
34 image_buffer = -100*np.ones((NUM_BUFFERED_SWEEPS,\
35 NUM_SCANS_PER_SWEEP*NFFT))
36
37 def __init__(self, sdr=None, fig=None):
38 self.fig = fig if fig else pyl.figure()
39 self.sdr = sdr if sdr else RtlSdr()
40
41 self.init_plot()
42
43 def init_plot(self):
44 self.ax = self.fig.add_subplot(1,1,1)
45 self.image = self.ax.imshow(self.image_buffer, aspect='auto',\
46 interpolation='nearest', vmin=-50, vmax=10)
47 self.ax.set_xlabel('Current frequency (MHz)')
48 self.ax.get_yaxis().set_visible(False)
49
50 self.fig.canvas.mpl_connect('scroll_event', self.on_click)
51 self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)
52 self.fig.canvas.mpl_connect('key_release_event', self.on_key_release)
53
54 def update_plot_labels(self):
55 fc = self.sdr.fc
56 rs = self.sdr.rs
57 freq_range = (fc - rs/2)/1e6, (fc + rs*(NUM_SCANS_PER_SWEEP - 0.5))/1e6
58
59 self.image.set_extent(freq_range + (0, 1))
60 self.fig.canvas.draw_idle()
61
62 def on_click(self, event):
63 if event.button == 'up':
64 self.sdr.fc += FREQ_INC_FINE if self.shift_key_down else FREQ_INC_COARSE
65
66 self.update_plot_labels()
67 elif event.button == 'down':
68 self.sdr.fc -= FREQ_INC_FINE if self.shift_key_down else FREQ_INC_COARSE
69
70 self.update_plot_labels()
71
72 def on_key_press(self, event):
73 if event.key == '+':
74 self.sdr.gain += GAIN_INC
75 elif event.key == '-':
76 self.sdr.gain -= GAIN_INC
77 elif event.key == ' ':
78 self.sdr.gain = 'auto'
79 elif event.key == 'shift':
80 self.shift_key_down = True
81 elif event.key == 'enter':
82 # see if valid frequency was entered, then change center frequency
83 try:
84 # join individual key presses into a string
85 input = ''.join(self.keyboard_buffer)
86
87 # if we're doing multiple adjacent scans, we need to figure out
88 # the appropriate center freq for the leftmost scan
89 center_freq = float(input)*1e6 + (self.sdr.rs/2)*(1 - NUM_SCANS_PER_SWEEP)
90 self.sdr.fc = center_freq
91
92 self.update_plot_labels()
93 except ValueError:
94 pass
95
96 self.keyboard_buffer = []
97 else:
98 self.keyboard_buffer.append(event.key)
99
100 def on_key_release(self, event):
101 if event.key == 'shift':
102 self.shift_key_down = False
103
104 def update(self, *args):
105 # save center freq. since we're gonna be changing it
106 start_fc = self.sdr.fc
107
108 # prepare space in buffer
109 # TODO: use indexing to avoid recreating buffer each time
110 self.image_buffer = np.roll(self.image_buffer, 1, axis=0)
111
112 for scan_num, start_ind in enumerate(range(0, NUM_SCANS_PER_SWEEP*NFFT, NFFT)):
113 self.sdr.fc += self.sdr.rs*scan_num
114
115 # estimate PSD for one scan
116 samples = self.sdr.read_samples(NUM_SAMPLES_PER_SCAN)
117 psd_scan, f = psd(samples, NFFT=NFFT)
118
119 self.image_buffer[0, start_ind: start_ind+NFFT] = 10*np.log10(psd_scan)
120
121 # plot entire sweep
122 self.image.set_array(self.image_buffer)
123
124 # restore original center freq.
125 self.sdr.fc = start_fc
126
127 return self.image,
128
129 def start(self):
130 self.update_plot_labels()
e11ad522 »
2012-05-21 fix demo for mac osx
131 if sys.platform == 'darwin':
132 # Disable blitting. The matplotlib.animation's restore_region()
133 # method is only implemented for the Agg-based backends,
134 # which the macosx backend is not.
135 blit = False
136 else:
137 blit = True
138 ani = animation.FuncAnimation(self.fig, self.update, interval=50,
139 blit=blit)
7218059b » Roger
2012-05-05 add waterfall demo
140
141 pyl.show()
142
143 return
144
145
146 def main():
147 sdr = RtlSdr()
148 wf = Waterfall(sdr)
149
150 # some defaults
151 sdr.rs = 3e6
152 sdr.fc = 100e6
153 sdr.gain = 10
154
155 wf.start()
156
157 # cleanup
158 sdr.close()
159
160
161 if __name__ == '__main__':
e11ad522 »
2012-05-21 fix demo for mac osx
162 main()
Something went wrong with that request. Please try again.