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