git.gag.com Git - debian/gnuradio/blob - gr-utils/src/python/gr_plot_qt.py

   1 #!/usr/bin/env python
   2
   3 try:
   4     import scipy
   5     from scipy import fftpack
   6 except ImportError:
   7     print "Please install SciPy to run this script (http://www.scipy.org/)"
   8     raise SystemExit, 1
   9
  10 import sys, os
  11 from PyQt4 import Qt, QtCore, QtGui
  12 import PyQt4.Qwt5 as Qwt
  13 from matplotlib import mlab
  14 from optparse import OptionParser
  15 from gnuradio import eng_notation
  16
  17 from pyqt_plot import Ui_MainWindow
  18
  19 class gr_plot_qt(QtGui.QMainWindow):
  20     def __init__(self, qapp, filename, options, parent=None):
  21         QtGui.QWidget.__init__(self, parent)
  22         self.gui = Ui_MainWindow()
  23         self.gui.setupUi(self)
  24
  25         self.block_length = options.block_length
  26         self.start = options.start
  27         self.sample_rate = options.sample_rate
  28         self.psdfftsize = options.psd_size
  29         self.specfftsize = options.spec_size
  30         self.winfunc = scipy.blackman
  31         self.sizeof_data = 8
  32         self.datatype = scipy.complex64
  33         self.iq = list()
  34         self.time = list()
  35
  36         # Set up basic plot attributes
  37         self.gui.timePlot.setAxisTitle(self.gui.timePlot.xBottom, "Time (sec)")
  38         self.gui.timePlot.setAxisTitle(self.gui.timePlot.yLeft, "Amplitude (V)")
  39         self.gui.freqPlot.setAxisTitle(self.gui.freqPlot.xBottom, "Frequency (Hz)")
  40         self.gui.freqPlot.setAxisTitle(self.gui.freqPlot.yLeft, "Magnitude (dB)")
  41
  42         # Set up FFT size combo box
  43         self.gui.fftComboBox.addItems(["128", "256", "512", "1024", "2048",
  44                                        "4096", "8192", "16384", "32768"])
  45         pos = self.gui.fftComboBox.findText(Qt.QString("%1").arg(self.psdfftsize))
  46         self.gui.fftComboBox.setCurrentIndex(pos)
  47         self.connect(self.gui.fftComboBox,
  48                      Qt.SIGNAL("activated (const QString&)"),
  49                      self.fftComboBoxEdit)
  50
  51         # Set up color scheme box
  52         self.color_modes = {"Black on White" : self.color_black_on_white,
  53                             "White on Black" : self.color_white_on_black,
  54                             "Blue on Black"  : self.color_blue_on_black,
  55                             "Green on Black" : self.color_green_on_black}
  56         self.gui.colorComboBox.addItems(self.color_modes.keys())
  57         pos = self.gui.colorComboBox.findText("Blue on Black")
  58         self.gui.colorComboBox.setCurrentIndex(pos)
  59         self.connect(self.gui.colorComboBox,
  60                      Qt.SIGNAL("activated (const QString&)"),
  61                      self.colorComboBoxEdit)
  62
  63
  64         # Create zoom functionality for the plots
  65         self.timeZoomer = Qwt.QwtPlotZoomer(self.gui.timePlot.xBottom,
  66                                             self.gui.timePlot.yLeft,
  67                                             Qwt.QwtPicker.PointSelection,
  68                                             Qwt.QwtPicker.AlwaysOn,
  69                                             self.gui.timePlot.canvas())
  70
  71         self.freqZoomer = Qwt.QwtPlotZoomer(self.gui.freqPlot.xBottom,
  72                                             self.gui.freqPlot.yLeft,
  73                                             Qwt.QwtPicker.PointSelection,
  74                                             Qwt.QwtPicker.AlwaysOn,
  75                                             self.gui.freqPlot.canvas())
  76
  77         self.picker = Qwt.QwtPlotPicker(self.gui.timePlot.xBottom,
  78                                         self.gui.timePlot.yLeft,
  79                                         Qwt.QwtPicker.PointSelection,
  80                                         Qwt.QwtPlotPicker.CrossRubberBand,
  81                                         Qwt.QwtPicker.AlwaysOn,
  82                                         self.gui.timePlot.canvas())
  83         self.picker.connect(self.picker,
  84                             Qt.SIGNAL('selected(const QwtDoublePoint&)'),
  85                             self.clickMe)
  86
  87         # Set up action when tab is changed
  88         self.connect(self.gui.tabGroup,
  89                      Qt.SIGNAL("currentChanged (int)"),
  90                      self.tabChanged)
  91
  92         # Add a legend to the Time plot
  93         legend_real = Qwt.QwtLegend()
  94         self.gui.timePlot.insertLegend(legend_real)
  95
  96         # Set up slider
  97         self.gui.plotHBar.setSingleStep(1)
  98         self.gui.plotHBar.setPageStep(self.block_length)
  99         self.gui.plotHBar.setMinimum(0)
 100         self.gui.plotHBar.setMaximum(self.block_length)
 101         self.connect(self.gui.plotHBar,
 102                      Qt.SIGNAL("valueChanged(int)"),
 103                      self.sliderMoved)
 104
 105         # Connect Open action to Open Dialog box
 106         self.connect(self.gui.action_open,
 107                      Qt.SIGNAL("activated()"),
 108                      self.open_file)
 109
 110         # Set up file position boxes to update current figure
 111         self.connect(self.gui.filePosStartLineEdit,
 112                      Qt.SIGNAL("editingFinished()"),
 113                      self.file_position_changed)
 114         self.connect(self.gui.filePosStopLineEdit,
 115                      Qt.SIGNAL("editingFinished()"),
 116                      self.file_position_changed)
 117         self.connect(self.gui.filePosLengthLineEdit,
 118                      Qt.SIGNAL("editingFinished()"),
 119                      self.file_length_changed)
 120
 121         self.connect(self.gui.fileTimeStartLineEdit,
 122                      Qt.SIGNAL("editingFinished()"),
 123                      self.file_time_changed)
 124         self.connect(self.gui.fileTimeStopLineEdit,
 125                      Qt.SIGNAL("editingFinished()"),
 126                      self.file_time_changed)
 127         self.connect(self.gui.fileTimeLengthLineEdit,
 128                      Qt.SIGNAL("editingFinished()"),
 129                      self.file_time_length_changed)
 130
 131         self.rcurve = Qwt.QwtPlotCurve("Real")
 132         self.icurve = Qwt.QwtPlotCurve("Imaginary")
 133
 134         self.icurve.attach(self.gui.timePlot)
 135         self.rcurve.attach(self.gui.timePlot)
 136
 137         self.psdcurve = Qwt.QwtPlotCurve("PSD")
 138         self.psdcurve.attach(self.gui.freqPlot)
 139
 140         # Set up initial color scheme
 141         self.color_modes["Blue on Black"]()
 142
 143         self.set_sample_rate(self.sample_rate)
 144         self.connect(self.gui.sampleRateLineEdit,
 145                      Qt.SIGNAL("editingFinished()"),
 146                      self.sample_rate_changed)
 147
 148         if(filename is not None):
 149             self.initialize(filename)
 150
 151         self.show()
 152
 153     def open_file(self):
 154         filename = Qt.QFileDialog.getOpenFileName(self, "Open", ".")
 155         if(filename != ""):
 156             print filename
 157             self.initialize(filename)
 158
 159     def initialize(self, filename):
 160         self.hfile = open(filename, "r")
 161
 162         self.setWindowTitle(("GNU Radio File Plot Utility: %s" % filename))
 163
 164         self.gui.filePosStartLineEdit.setText("0")
 165         self.gui.filePosStopLineEdit.setText("0")
 166         self.gui.fileTimeStartLineEdit.setText("0")
 167         self.gui.fileTimeStopLineEdit.setText("0")
 168
 169         self.cur_start = 0
 170         self.cur_stop = self.block_length
 171
 172         self.init_data_input()
 173         self.get_data(self.cur_start, self.cur_stop)
 174         self.get_psd()
 175         self.gui.plotHBar.setSliderPosition(0)
 176         self.gui.plotHBar.setMaximum(self.signal_size)
 177
 178
 179         self.update_time_curves()
 180         self.update_psd_curves()
 181
 182     def init_data_input(self):
 183         self.hfile.seek(0, os.SEEK_END)
 184         self.signal_size = self.hfile.tell()/self.sizeof_data
 185         print "Sizeof File: ", self.signal_size
 186         self.hfile.seek(0, os.SEEK_SET)
 187
 188     def get_data(self, start, end):
 189         if(end > start):
 190             self.hfile.seek(start*self.sizeof_data, os.SEEK_SET)
 191             self.position = start
 192             try:
 193                 iq = scipy.fromfile(self.hfile, dtype=self.datatype,
 194                                     count=end-start)
 195
 196                 if(len(iq) < (end-start)):
 197                     end = len(iq)
 198                     self.gui.filePosLengthLineEdit.setText(Qt.QString("%1").arg(end))
 199                     self.gui.plotHBar.setMaximum(end)
 200                     self.gui.plotHBar.setSingleStep(end)
 201                     self.file_length_changed()
 202
 203                 tstep = 1.0 / self.sample_rate
 204                 self.iq = iq
 205                 self.time = [tstep*(self.position + i) for i in xrange(len(self.iq))]
 206
 207                 self.set_file_pos_box(start, end)
 208             except MemoryError:
 209                 pass
 210         else:
 211             # Do we want to do anything about this?
 212             pass
 213
 214     def get_psd(self):
 215         winpoints = self.winfunc(self.psdfftsize)
 216         iq_psd, freq = mlab.psd(self.iq, Fs=self.sample_rate,
 217                                 NFFT=self.psdfftsize,
 218                                 noverlap=self.psdfftsize/4.0,
 219                                 window=winpoints,
 220                                 scale_by_freq=False)
 221         self.iq_psd = 10.0*scipy.log10(abs(fftpack.fftshift(iq_psd)))
 222         self.freq = freq - self.sample_rate/2.0
 223
 224
 225     def clickMe(self, qPoint):
 226         print qPoint.x()
 227
 228     def fftComboBoxEdit(self, fftSize):
 229         self.psdfftsize = fftSize.toInt()[0]
 230         self.get_psd()
 231         self.update_psd_curves()
 232
 233     def colorComboBoxEdit(self, colorSelection):
 234         colorstr = str(colorSelection.toAscii())
 235         color_func = self.color_modes[colorstr]
 236         color_func()
 237
 238     def sliderMoved(self, value):
 239         pos_start = value
 240         pos_end = value + self.gui.plotHBar.pageStep()
 241
 242         self.get_data(pos_start, pos_end)
 243         self.get_psd()
 244         self.update_time_curves()
 245         self.update_psd_curves()
 246
 247     def set_sample_rate(self, sr):
 248         self.sample_rate = sr
 249         srstr = eng_notation.num_to_str(self.sample_rate)
 250         self.gui.sampleRateLineEdit.setText(Qt.QString("%1").arg(srstr))
 251
 252     def sample_rate_changed(self):
 253         srstr = self.gui.sampleRateLineEdit.text().toAscii()
 254         self.sample_rate = eng_notation.str_to_num(srstr)
 255         self.set_file_pos_box(self.cur_start, self.cur_stop)
 256         self.get_data(self.cur_start, self.cur_stop)
 257         self.get_psd()
 258         self.update_time_curves()
 259         self.update_psd_curves()
 260
 261     def set_file_pos_box(self, start, end):
 262         tstart = start / self.sample_rate
 263         tend = end / self.sample_rate
 264
 265         self.gui.filePosStartLineEdit.setText(Qt.QString("%1").arg(start))
 266         self.gui.filePosStopLineEdit.setText(Qt.QString("%1").arg(end))
 267         self.gui.filePosLengthLineEdit.setText(Qt.QString("%1").arg(end-start))
 268
 269         self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
 270         self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
 271         self.gui.fileTimeLengthLineEdit.setText(Qt.QString("%1").arg(tend-tstart))
 272
 273     def file_position_changed(self):
 274         start  = self.gui.filePosStartLineEdit.text().toInt()
 275         end    = self.gui.filePosStopLineEdit.text().toInt()
 276         if((start[1] == True) and (end[1] == True)):
 277             self.cur_start = start[0]
 278             self.cur_stop = end[0]
 279
 280             tstart = self.cur_start / self.sample_rate
 281             tend = self.cur_stop / self.sample_rate
 282             self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
 283             self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
 284
 285             self.get_data(self.cur_start, self.cur_stop)
 286
 287             self.update_time_curves()
 288             self.update_psd_curves()
 289
 290         # If there's a non-digit character, reset box
 291         else:
 292             self.set_file_pos_box(self.cur_start, self.cur_stop)
 293
 294     def file_time_changed(self):
 295         tstart = self.gui.fileTimeStartLineEdit.text().toDouble()
 296         tstop  = self.gui.fileTimeStopLineEdit.text().toDouble()
 297         if((tstart[1] == True) and (tstop[1] == True)):
 298             self.cur_start = int(tstart[0] * self.sample_rate)
 299             self.cur_stop = int(tstop[0] * self.sample_rate)
 300             self.get_data(self.cur_start, self.cur_stop)
 301
 302             self.gui.filePosStartLineEdit.setText(Qt.QString("%1").arg(self.cur_start))
 303             self.gui.filePosStopLineEdit.setText(Qt.QString("%1").arg(self.cur_stop))
 304
 305             self.update_time_curves()
 306             self.update_psd_curves()
 307         # If there's a non-digit character, reset box
 308         else:
 309             self.set_file_pos_box(self.cur_start, self.cur_stop)
 310
 311     def file_length_changed(self):
 312         start = self.gui.filePosStartLineEdit.text().toInt()
 313         length = self.gui.filePosLengthLineEdit.text().toInt()
 314
 315         if((start[1] == True) and (length[1] == True)):
 316             self.cur_start = start[0]
 317             self.block_length = length[0]
 318             self.cur_stop = self.cur_start + self.block_length
 319
 320             tstart = self.cur_start / self.sample_rate
 321             tend = self.cur_stop / self.sample_rate
 322             tlen = self.block_length / self.sample_rate
 323             self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
 324             self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
 325             self.gui.fileTimeLengthLineEdit.setText(Qt.QString("%1").arg(tlen))
 326
 327             self.gui.plotHBar.setPageStep(self.block_length)
 328
 329             self.get_data(self.cur_start, self.cur_stop)
 330             self.get_psd()
 331
 332             self.update_time_curves()
 333             self.update_psd_curves()
 334         # If there's a non-digit character, reset box
 335         else:
 336             self.set_file_pos_box(self.cur_start, self.cur_stop)
 337
 338     def file_time_length_changed(self):
 339         tstart = self.gui.fileTimeStartLineEdit.text().toDouble()
 340         tlength = self.gui.fileTimeLengthLineEdit.text().toDouble()
 341         if((tstart[1] == True) and (tlength[1] == True)):
 342             self.cur_start = int(tstart[0] * self.sample_rate)
 343             self.block_length = int(tlength[0] * self.sample_rate)
 344             self.cur_stop = self.cur_start + self.block_length
 345
 346             tstart = self.cur_start / self.sample_rate
 347             tend = self.cur_stop / self.sample_rate
 348             tlen = self.block_length / self.sample_rate
 349             self.gui.fileTimeStartLineEdit.setText(Qt.QString("%1").arg(tstart))
 350             self.gui.fileTimeStopLineEdit.setText(Qt.QString("%1").arg(tend))
 351             self.gui.fileTimeLengthLineEdit.setText(Qt.QString("%1").arg(tlen))
 352
 353             self.get_data(self.cur_start, self.cur_stop)
 354             self.get_psd()
 355
 356             self.update_time_curves()
 357             self.update_psd_curves()
 358         # If there's a non-digit character, reset box
 359         else:
 360             self.set_file_pos_box(self.cur_start, self.cur_stop)
 361
 362
 363     def update_time_curves(self):
 364         self.icurve.setData(self.time, self.iq.imag)
 365         self.rcurve.setData(self.time, self.iq.real)
 366
 367         # Reset the x-axis to the new time scale
 368         iqmax = 1.5 * max(max(self.iq.real), max(self.iq.imag))
 369         iqmin = 1.5 * min(min(self.iq.real), min(self.iq.imag))
 370         self.gui.timePlot.setAxisScale(self.gui.timePlot.xBottom,
 371                                        min(self.time),
 372                                        max(self.time))
 373         self.gui.timePlot.setAxisScale(self.gui.timePlot.yLeft,
 374                                        iqmin,
 375                                        iqmax)
 376
 377         # Set the zoomer base to unzoom to the new axis
 378         self.timeZoomer.setZoomBase()
 379
 380         self.gui.timePlot.replot()
 381
 382     def update_psd_curves(self):
 383         self.psdcurve.setData(self.freq, self.iq_psd)
 384
 385         self.gui.freqPlot.setAxisScale(self.gui.freqPlot.xBottom,
 386                                        min(self.freq),
 387                                        max(self.freq))
 388
 389         # Set the zoomer base to unzoom to the new axis
 390         self.freqZoomer.setZoomBase()
 391
 392         self.gui.freqPlot.replot()
 393
 394     def tabChanged(self, index):
 395         self.gui.timePlot.replot()
 396         self.gui.freqPlot.replot()
 397
 398     def color_black_on_white(self):
 399         blue = QtGui.qRgb(0x00, 0x00, 0xFF)
 400         red = QtGui.qRgb(0xFF, 0x00, 0x00)
 401
 402         blackBrush = Qt.QBrush(Qt.QColor("black"))
 403         blueBrush = Qt.QBrush(Qt.QColor(blue))
 404         redBrush = Qt.QBrush(Qt.QColor(red))
 405
 406         self.gui.timePlot.setCanvasBackground(Qt.QColor("white"))
 407         self.gui.freqPlot.setCanvasBackground(Qt.QColor("white"))
 408         self.picker.setTrackerPen(Qt.QPen(blackBrush, 2))
 409         self.timeZoomer.setTrackerPen(Qt.QPen(blackBrush, 2))
 410         self.timeZoomer.setRubberBandPen(Qt.QPen(blackBrush, 2))
 411         self.freqZoomer.setTrackerPen(Qt.QPen(blackBrush, 2))
 412         self.freqZoomer.setRubberBandPen(Qt.QPen(blackBrush, 2))
 413         self.psdcurve.setPen(Qt.QPen(blueBrush, 1))
 414         self.rcurve.setPen(Qt.QPen(blueBrush, 2))
 415         self.icurve.setPen(Qt.QPen(redBrush, 2))
 416
 417         self.gui.timePlot.replot()
 418         self.gui.freqPlot.replot()
 419
 420     def color_white_on_black(self):
 421         white = QtGui.qRgb(0xFF, 0xFF, 0xFF)
 422         red = QtGui.qRgb(0xFF, 0x00, 0x00)
 423
 424         whiteBrush = Qt.QBrush(Qt.QColor("white"))
 425         whiteBrush = Qt.QBrush(Qt.QColor(white))
 426         redBrush = Qt.QBrush(Qt.QColor(red))
 427
 428         self.gui.timePlot.setCanvasBackground(QtGui.QColor("black"))
 429         self.gui.freqPlot.setCanvasBackground(QtGui.QColor("black"))
 430         self.picker.setTrackerPen(Qt.QPen(whiteBrush, 2))
 431         self.timeZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
 432         self.timeZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
 433         self.freqZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
 434         self.freqZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
 435         self.psdcurve.setPen(Qt.QPen(whiteBrush, 1))
 436         self.rcurve.setPen(Qt.QPen(whiteBrush, 2))
 437         self.icurve.setPen(Qt.QPen(redBrush, 2))
 438
 439         self.gui.timePlot.replot()
 440         self.gui.freqPlot.replot()
 441
 442
 443     def color_green_on_black(self):
 444         green = QtGui.qRgb(0x00, 0xFF, 0x00)
 445         red = QtGui.qRgb(0xFF, 0x00, 0x50)
 446
 447         whiteBrush = Qt.QBrush(Qt.QColor("white"))
 448         greenBrush = Qt.QBrush(Qt.QColor(green))
 449         redBrush = Qt.QBrush(Qt.QColor(red))
 450
 451         self.gui.timePlot.setCanvasBackground(QtGui.QColor("black"))
 452         self.gui.freqPlot.setCanvasBackground(QtGui.QColor("black"))
 453         self.picker.setTrackerPen(Qt.QPen(whiteBrush, 2))
 454         self.timeZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
 455         self.timeZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
 456         self.freqZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
 457         self.freqZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
 458         self.psdcurve.setPen(Qt.QPen(greenBrush, 1))
 459         self.rcurve.setPen(Qt.QPen(greenBrush, 2))
 460         self.icurve.setPen(Qt.QPen(redBrush, 2))
 461
 462         self.gui.timePlot.replot()
 463         self.gui.freqPlot.replot()
 464
 465     def color_blue_on_black(self):
 466         blue = QtGui.qRgb(0x00, 0x00, 0xFF)
 467         red = QtGui.qRgb(0xFF, 0x00, 0x00)
 468
 469         whiteBrush = Qt.QBrush(Qt.QColor("white"))
 470         blueBrush = Qt.QBrush(Qt.QColor(blue))
 471         redBrush = Qt.QBrush(Qt.QColor(red))
 472
 473         self.gui.timePlot.setCanvasBackground(QtGui.QColor("black"))
 474         self.gui.freqPlot.setCanvasBackground(QtGui.QColor("black"))
 475         self.picker.setTrackerPen(Qt.QPen(whiteBrush, 2))
 476         self.timeZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
 477         self.timeZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
 478         self.freqZoomer.setTrackerPen(Qt.QPen(whiteBrush, 2))
 479         self.freqZoomer.setRubberBandPen(Qt.QPen(whiteBrush, 2))
 480         self.psdcurve.setPen(Qt.QPen(blueBrush, 1))
 481         self.rcurve.setPen(Qt.QPen(blueBrush, 2))
 482         self.icurve.setPen(Qt.QPen(redBrush, 2))
 483
 484         self.gui.timePlot.replot()
 485         self.gui.freqPlot.replot()
 486
 487 def setup_options():
 488     usage="%prog: [options] (input_filename)"
 489     description = ""
 490
 491     parser = OptionParser(conflict_handler="resolve", usage=usage, description=description)
 492     parser.add_option("-B", "--block-length", type="int", default=8192,
 493                       help="Specify the block size [default=%default]")
 494     parser.add_option("-s", "--start", type="int", default=0,
 495                       help="Specify where to start in the file [default=%default]")
 496     parser.add_option("-R", "--sample-rate", type="float", default=1.0,
 497                       help="Set the sampler rate of the data [default=%default]")
 498     parser.add_option("", "--psd-size", type="int", default=2048,
 499                       help="Set the size of the PSD FFT [default=%default]")
 500     parser.add_option("", "--spec-size", type="int", default=256,
 501                       help="Set the size of the spectrogram FFT [default=%default]")
 502
 503     return parser
 504
 505 def main(args):
 506     parser = setup_options()
 507     (options, args) = parser.parse_args ()
 508
 509     if(len(args) == 1):
 510         filename = args[0]
 511     else:
 512         filename = None
 513
 514     app = Qt.QApplication(args)
 515     gplt = gr_plot_qt(app, filename, options)
 516     app.exec_()
 517
 518 if __name__ == '__main__':
 519     main(sys.argv)
 520