Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Newer
Older
100644 163 lines (125 sloc) 5.063 kb
7218059 add waterfall demo
Roger authored
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
e11ad52 @xyb fix demo for mac osx
authored
6 import sys
7218059 add waterfall demo
Roger authored
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()
e11ad52 @xyb fix demo for mac osx
authored
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)
7218059 add waterfall demo
Roger authored
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__':
e11ad52 @xyb fix demo for mac osx
authored
162 main()
Something went wrong with that request. Please try again.