1import collections, math
2
3from PyQt5 import QtCore
4import pyqtgraph as pg
5
6# Basic PyQtGraph settings
7pg.setConfigOptions(antialias=True)
8
9
10class SpectrumPlotWidget:
11    """Main spectrum plot"""
12    def __init__(self, layout):
13        if not isinstance(layout, pg.GraphicsLayoutWidget):
14            raise ValueError("layout must be instance of pyqtgraph.GraphicsLayoutWidget")
15
16        self.layout = layout
17
18        self.main_curve = True
19        self.main_color = pg.mkColor("y")
20        self.persistence = False
21        self.persistence_length = 5
22        self.persistence_decay = "exponential"
23        self.persistence_color = pg.mkColor("g")
24        self.persistence_data = None
25        self.persistence_curves = None
26        self.peak_hold_max = False
27        self.peak_hold_max_color = pg.mkColor("r")
28        self.peak_hold_min = False
29        self.peak_hold_min_color = pg.mkColor("b")
30        self.average = False
31        self.average_color = pg.mkColor("c")
32
33        self.create_plot()
34
35    def create_plot(self):
36        """Create main spectrum plot"""
37        self.posLabel = self.layout.addLabel(row=0, col=0, justify="right")
38        self.plot = self.layout.addPlot(row=1, col=0)
39        self.plot.showGrid(x=True, y=True)
40        self.plot.setLabel("left", "Power", units="dB")
41        self.plot.setLabel("bottom", "Frequency", units="Hz")
42        self.plot.setLimits(xMin=0)
43        self.plot.showButtons()
44
45        #self.plot.setDownsampling(mode="peak")
46        #self.plot.setClipToView(True)
47
48        self.create_persistence_curves()
49        self.create_average_curve()
50        self.create_peak_hold_min_curve()
51        self.create_peak_hold_max_curve()
52        self.create_main_curve()
53
54        # Create crosshair
55        self.vLine = pg.InfiniteLine(angle=90, movable=False)
56        self.vLine.setZValue(1000)
57        self.hLine = pg.InfiniteLine(angle=0, movable=False)
58        self.vLine.setZValue(1000)
59        self.plot.addItem(self.vLine, ignoreBounds=True)
60        self.plot.addItem(self.hLine, ignoreBounds=True)
61        self.mouseProxy = pg.SignalProxy(self.plot.scene().sigMouseMoved,
62                                         rateLimit=60, slot=self.mouse_moved)
63
64    def create_main_curve(self):
65        """Create main spectrum curve"""
66        self.curve = self.plot.plot(pen=self.main_color)
67        self.curve.setZValue(900)
68
69    def create_peak_hold_max_curve(self):
70        """Create max. peak hold curve"""
71        self.curve_peak_hold_max = self.plot.plot(pen=self.peak_hold_max_color)
72        self.curve_peak_hold_max.setZValue(800)
73
74    def create_peak_hold_min_curve(self):
75        """Create min. peak hold curve"""
76        self.curve_peak_hold_min = self.plot.plot(pen=self.peak_hold_min_color)
77        self.curve_peak_hold_min.setZValue(800)
78
79    def create_average_curve(self):
80        """Create average curve"""
81        self.curve_average = self.plot.plot(pen=self.average_color)
82        self.curve_average.setZValue(700)
83
84    def create_persistence_curves(self):
85        """Create spectrum persistence curves"""
86        z_index_base = 600
87        decay = self.get_decay()
88        self.persistence_curves = []
89        for i in range(self.persistence_length):
90            alpha = 255 * decay(i + 1, self.persistence_length + 1)
91            color = self.persistence_color
92            curve = self.plot.plot(pen=(color.red(), color.green(), color.blue(), alpha))
93            curve.setZValue(z_index_base - i)
94            self.persistence_curves.append(curve)
95
96    def set_colors(self):
97        """Set colors of all curves"""
98        self.curve.setPen(self.main_color)
99        self.curve_peak_hold_max.setPen(self.peak_hold_max_color)
100        self.curve_peak_hold_min.setPen(self.peak_hold_min_color)
101        self.curve_average.setPen(self.average_color)
102
103        decay = self.get_decay()
104        for i, curve in enumerate(self.persistence_curves):
105            alpha = 255 * decay(i + 1, self.persistence_length + 1)
106            color = self.persistence_color
107            curve.setPen((color.red(), color.green(), color.blue(), alpha))
108
109    def decay_linear(self, x, length):
110        """Get alpha value for persistence curve (linear decay)"""
111        return (-x / length) + 1
112
113    def decay_exponential(self, x, length, const=1 / 3):
114        """Get alpha value for persistence curve (exponential decay)"""
115        return math.e**(-x / (length * const))
116
117    def get_decay(self):
118        """Get decay function"""
119        if self.persistence_decay == 'exponential':
120            return self.decay_exponential
121        else:
122            return self.decay_linear
123
124    def update_plot(self, data_storage, force=False):
125        """Update main spectrum curve"""
126        if data_storage.x is None:
127            return
128
129        if self.main_curve or force:
130            self.curve.setData(data_storage.x, data_storage.y)
131            if force:
132                self.curve.setVisible(self.main_curve)
133
134    def update_peak_hold_max(self, data_storage, force=False):
135        """Update max. peak hold curve"""
136        if data_storage.x is None:
137            return
138
139        if self.peak_hold_max or force:
140            self.curve_peak_hold_max.setData(data_storage.x, data_storage.peak_hold_max)
141            if force:
142                self.curve_peak_hold_max.setVisible(self.peak_hold_max)
143
144    def update_peak_hold_min(self, data_storage, force=False):
145        """Update min. peak hold curve"""
146        if data_storage.x is None:
147            return
148
149        if self.peak_hold_min or force:
150            self.curve_peak_hold_min.setData(data_storage.x, data_storage.peak_hold_min)
151            if force:
152                self.curve_peak_hold_min.setVisible(self.peak_hold_min)
153
154    def update_average(self, data_storage, force=False):
155        """Update average curve"""
156        if data_storage.x is None:
157            return
158
159        if self.average or force:
160            self.curve_average.setData(data_storage.x, data_storage.average)
161            if force:
162                self.curve_average.setVisible(self.average)
163
164    def update_persistence(self, data_storage, force=False):
165        """Update persistence curves"""
166        if data_storage.x is None:
167            return
168
169        if self.persistence or force:
170            if self.persistence_data is None:
171                self.persistence_data = collections.deque(maxlen=self.persistence_length)
172            else:
173                for i, y in enumerate(self.persistence_data):
174                    curve = self.persistence_curves[i]
175                    curve.setData(data_storage.x, y)
176                    if force:
177                        curve.setVisible(self.persistence)
178            self.persistence_data.appendleft(data_storage.y)
179
180    def recalculate_plot(self, data_storage):
181        """Recalculate plot from history"""
182        if data_storage.x is None:
183            return
184
185        QtCore.QTimer.singleShot(0, lambda: self.update_plot(data_storage, force=True))
186        QtCore.QTimer.singleShot(0, lambda: self.update_average(data_storage, force=True))
187        QtCore.QTimer.singleShot(0, lambda: self.update_peak_hold_max(data_storage, force=True))
188        QtCore.QTimer.singleShot(0, lambda: self.update_peak_hold_min(data_storage, force=True))
189
190    def recalculate_persistence(self, data_storage):
191        """Recalculate persistence data and update persistence curves"""
192        if data_storage.x is None:
193            return
194
195        self.clear_persistence()
196        self.persistence_data = collections.deque(maxlen=self.persistence_length)
197        for i in range(min(self.persistence_length, data_storage.history.history_size - 1)):
198            data = data_storage.history[-i - 2]
199            if data_storage.smooth:
200                data = data_storage.smooth_data(data)
201            self.persistence_data.append(data)
202        QtCore.QTimer.singleShot(0, lambda: self.update_persistence(data_storage, force=True))
203
204    def mouse_moved(self, evt):
205        """Update crosshair when mouse is moved"""
206        pos = evt[0]
207        if self.plot.sceneBoundingRect().contains(pos):
208            mousePoint = self.plot.vb.mapSceneToView(pos)
209            self.posLabel.setText(
210                "<span style='font-size: 12pt'>f={:0.3f} MHz, P={:0.3f} dB</span>".format(
211                    mousePoint.x() / 1e6,
212                    mousePoint.y()
213                )
214            )
215            self.vLine.setPos(mousePoint.x())
216            self.hLine.setPos(mousePoint.y())
217
218    def clear_plot(self):
219        """Clear main spectrum curve"""
220        self.curve.clear()
221
222    def clear_peak_hold_max(self):
223        """Clear max. peak hold curve"""
224        self.curve_peak_hold_max.clear()
225
226    def clear_peak_hold_min(self):
227        """Clear min. peak hold curve"""
228        self.curve_peak_hold_min.clear()
229
230    def clear_average(self):
231        """Clear average curve"""
232        self.curve_average.clear()
233
234    def clear_persistence(self):
235        """Clear spectrum persistence curves"""
236        self.persistence_data = None
237        for curve in self.persistence_curves:
238            curve.clear()
239            self.plot.removeItem(curve)
240        self.create_persistence_curves()
241
242
243class WaterfallPlotWidget:
244    """Waterfall plot"""
245    def __init__(self, layout, histogram_layout=None):
246        if not isinstance(layout, pg.GraphicsLayoutWidget):
247            raise ValueError("layout must be instance of pyqtgraph.GraphicsLayoutWidget")
248
249        if histogram_layout and not isinstance(histogram_layout, pg.GraphicsLayoutWidget):
250            raise ValueError("histogram_layout must be instance of pyqtgraph.GraphicsLayoutWidget")
251
252        self.layout = layout
253        self.histogram_layout = histogram_layout
254
255        self.history_size = 100
256        self.counter = 0
257
258        self.create_plot()
259
260    def create_plot(self):
261        """Create waterfall plot"""
262        self.plot = self.layout.addPlot()
263        self.plot.setLabel("bottom", "Frequency", units="Hz")
264        self.plot.setLabel("left", "Time")
265
266        self.plot.setYRange(-self.history_size, 0)
267        self.plot.setLimits(xMin=0, yMax=0)
268        self.plot.showButtons()
269        #self.plot.setAspectLocked(True)
270
271        #self.plot.setDownsampling(mode="peak")
272        #self.plot.setClipToView(True)
273
274        # Setup histogram widget (for controlling waterfall plot levels and gradients)
275        if self.histogram_layout:
276            self.histogram = pg.HistogramLUTItem()
277            self.histogram_layout.addItem(self.histogram)
278            self.histogram.gradient.loadPreset("flame")
279            #self.histogram.setHistogramRange(-50, 0)
280            #self.histogram.setLevels(-50, 0)
281
282    def update_plot(self, data_storage):
283        """Update waterfall plot"""
284        self.counter += 1
285
286        # Create waterfall image on first run
287        if self.counter == 1:
288            self.waterfallImg = pg.ImageItem()
289            self.waterfallImg.scale((data_storage.x[-1] - data_storage.x[0]) / len(data_storage.x), 1)
290            self.plot.clear()
291            self.plot.addItem(self.waterfallImg)
292
293        # Roll down one and replace leading edge with new data
294        self.waterfallImg.setImage(data_storage.history.buffer[-self.counter:].T,
295                                   autoLevels=False, autoRange=False)
296
297        # Move waterfall image to always start at 0
298        self.waterfallImg.setPos(
299            data_storage.x[0],
300            -self.counter if self.counter < self.history_size else -self.history_size
301        )
302
303        # Link histogram widget to waterfall image on first run
304        # (must be done after first data is received or else levels would be wrong)
305        if self.counter == 1 and self.histogram_layout:
306            self.histogram.setImageItem(self.waterfallImg)
307
308    def clear_plot(self):
309        """Clear waterfall plot"""
310        self.counter = 0
311