-
Notifications
You must be signed in to change notification settings - Fork 1
/
plot_fft_3d.py
121 lines (93 loc) · 2.95 KB
/
plot_fft_3d.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
#!/usr/bin/python
# import all the necessary Libraries
import sys
import numpy as np
from pyrf.devices.thinkrf import WSA4000
from pyrf.util import read_data_and_context
from pyrf.config import TriggerSettings
from pyrf.numpy_util import compute_fft
from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph.opengl as gl
from pyqtgraph_util import represent_fft_to_plot
from pyqtgraph_util import create_color_heatmap
# hold a constant reflevel
STATIC_REFLEVEL = -35
xSize = 128
ySize = 150
dut = WSA4000()
dut.connect(sys.argv[1])
powerSize = dut.spp()
if len(sys.argv) > 2:
center_freq = int(sys.argv[2]) * 1e6
else:
center_freq = 2450e6
# setup WSA settings
dut.reset()
dut.request_read_perm()
dut.ifgain(0)
dut.freq(center_freq)
dut.gain('high')
dut.fshift(0)
dut.decimation(0)
# Create a GL View widget to display data
app = QtGui.QApplication([])
w = gl.GLViewWidget()
w.show()
w.setWindowTitle('PYRF pyqtgraph example: 3D FFT Plot')
w.setCameraPosition(distance=100)
# add grid
gridSize = QtGui.QVector3D(4,5,6)
g = gl.GLGridItem(size=gridSize)
g.scale(5.85,5,0)
w.addItem(g)
# the initial plot data before data is received
z = np.random.normal(size=(xSize,ySize))
# initialize the x-y boundries of the plot
x = np.linspace(-73.5, 144.3, xSize)
y = np.linspace(0.1, 5, ySize)
# initialize the colors of the plots (light blue)
colors = np.ones((xSize,ySize,4), dtype=float)
colors[:,:,0] = 0
colors[:,:,1] = 0.4
colors[:,:,2] = 1
# plot the data
p3 = gl.GLSurfacePlotItem(x = x, y = y, z = z, shader = 'shaded',
colors = colors.reshape(xSize*ySize,4), smooth=False)
# determine the size
p3.scale(16./49., 20, 0.1)
# determine the location on the gride
p3.translate(-12, -50, 0)
w.addItem(p3)
w.pan(-10,-10,-10)
def update():
# update the plot to show new data
global p3, z, colors
# read data
data, context = read_data_and_context(dut, powerSize)
# ignore the reference level so plot doesn't change positions
context['reflevel'] = STATIC_REFLEVEL
# compute the fft of the complex data
powData = compute_fft(dut, data, context)
# compress the FFT into a 128 array
zData = represent_fft_to_plot(powData)
# move the data stream as well as colours back to show new data
for i in range (ySize):
if i == 0:
continue
else:
colors[:,ySize - i,:] = colors[:, ySize - i - 1,:]
z[:,ySize - i] = z[:, ySize - i - 1]
z[:,0] = zData[:,0]
# grab new color
colors[:,0,:] = create_color_heatmap(powData)
colors[:,1,:] = create_color_heatmap(powData)
# update the plot
p3.setData(z = z, colors = colors.reshape(xSize*ySize,4))
timer = QtCore.QTimer()
timer.timeout.connect(update)
timer.start(5)
## Start Qt event loop unless running in interactive mode.
if __name__ == '__main__':
import sys
if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
QtGui.QApplication.instance().exec_()