KaptureControlGui/KCG/widgets/FrequencyExtractWidget.py

from PyQt4 import QtCore, QtGui


from ..base import kcgwidget as kcgw
from ..base.globals import glob as global_objects
from ..base.backend import board
from ..base.backend.board import available_boards
from .. import config

import pyqtgraph as pg
from math import ceil
import numpy as np



from .PlotWidget import SubPlotWidget

__widget_id__ = None


class FrequencyExtractWidget(kcgw.KCGWidgets):
    def __init__(self, unique_id, parent):
        super(FrequencyExtractWidget, self).__init__()

        self.id = unique_id
        self.par = parent
        
        #ToDo: Maybe find valid data from the get-go?
        self.data = None


        self.board_config = board.get_board_config(available_boards[0]) 
        self.board_config.observe(self, self.dataSetChanged, 'lastDataSet')
        self.numADCs = self.board_config.get('adc_number')


        #Use a VBox. Top half should be the SubPlotWidget, bottom
        #half should be the controls
        self.layout = QtGui.QVBoxLayout()

        
        self.plotBox = QtGui.QVBoxLayout()
        self.plot = SubPlotWidget()

        self.region = pg.LinearRegionItem(values=[10,20])
        self.plot.plotItem.addItem(self.region)
        self.region.sigRegionChangeFinished.connect(self.doPlot)

        self.line = pg.InfiniteLine(angle=0, label='Y={value:0.2f}',
                                    labelOpts={'movable': True})
        self.plot.plotItem.addItem(self.line)
        self.line.hide()

        self.plotBox.addWidget(self.plot)


        #Controls for the plot
        self.controlsBox = QtGui.QHBoxLayout()

        self.adcSelect = QtGui.QComboBox(self)
        for i in range(self.numADCs):
            self.adcSelect.addItem("ADC {}".format(i+1))
        self.adcSelect.currentIndexChanged.connect(self.doPlot)


        self.bucketSelect = QtGui.QComboBox(self)
        self.bucketSelect.addItem("Mean")
        for i in range(184):
            self.bucketSelect.addItem("Bucket {}".format(i+1))
        self.bucketSelect.currentIndexChanged.connect(self.doPlot)


        self.fromBox = self.createSpinbox(0, 100000000, interval=100, connect=self.doPlot)
        self.toBox = self.createSpinbox(0, 100000000, start_value=1000, interval=100, connect=self.doPlot)

        self.controlsBox.addWidget(self.adcSelect)
        self.controlsBox.addWidget(self.bucketSelect)
        self.controlsBox.addStretch()
        self.controlsBox.addWidget(self.createLabel(text="From:"))
        self.controlsBox.addWidget(self.fromBox)
        self.controlsBox.addWidget(self.createLabel(text="To:"))
        self.controlsBox.addWidget(self.toBox)


        self.frequencyTools = QtGui.QHBoxLayout()
        
        self.freqText = QtGui.QLineEdit(self)
        self.isFreqValid = False

        self.frequencyTools.addWidget(self.createLabel(text="Frequency:"))
        self.frequencyTools.addWidget(self.freqText)


        self.layout.addLayout(self.plotBox)
        self.layout.addLayout(self.controlsBox)
        self.layout.addLayout(self.frequencyTools)
        self.setLayout(self.layout)
        self.setWindowTitle("Frequency Extract")


    def dataSetChanged(self, data=None):
        self.data = data
        self.doPlot()


    def doPlot(self):

        if not self.data:
            return

        #PlotWidget's "plot" function expects the fft_mode to have two entries
        #which we don't need here. So we need to shift our index by +2 in order
        #to accomodate for the two missing options in the bucket_select ComboBox
        fft_mode = self.bucketSelect.currentIndex() + 2
        adc = self.adcSelect.currentIndex()


        self.fftData = self.data.fft(adc=self.adcSelect.currentIndex(),
                                     frm=self.fromBox.value(),
                                     to=self.toBox.value(),
                                     drop_first_bin=True,
                                     nobunching=False)
        self.plot.plot(
                self.fftData,
                autorange=False,
                xvalueborders=[self.data.fftFreqDist(), self.data.fftMaxFreq()],
                fft_mode=fft_mode,
                log=False)

        self.plot.plotItem.setLabel('left', 'Spectral Intensity')

        #I am cheating a little bit and getting the readily processed data
        #from the plotItems, so I don't have to process, convert, format, etc.
        #the fftData again. We already did that in the SubPlotWidget.plot
        #function, so no need to do it again. Just steal the data from the
        #plotItems!
        self.fftYValues = self.plot.plotItemPlot[0].yData
        self.fftXValues = self.plot.plotItemPlot[0].xData

        region = self.region.getRegion()
   
        #Clip to valid data ranges
        rLeft = ceil(region[0]) if region[0] >=0 else 0
        rRight = ceil(region[1]) if region[1] < len(self.fftYValues) else len(self.fftYValues)

        #Make sure the region is within the boundaries of our data array
        if rLeft >= len(self.fftYValues) or rRight < 0:
            self.freqText.setText("Region boundaries out of data range")
            self.line.hide()
            self.isFreqValid = False
            self.freqText.setText("NaN")
            return

        #If the region is not really a region, ignore it
        if rLeft == rRight:
            self.freqText.setText("Region too small")
            self.line.hide()
            self.isFreqValid = False
            self.freqText.setText("NaN")
            return
    
        chunk = self.fftYValues[rLeft:rRight]
        maxVal = np.max(chunk)
        maxIndexInChunk = np.where(chunk == maxVal)

        self.line.setValue(maxVal)
        self.line.show()

        
        #np.where returns a tuple of indices (because it supports
        #multidimensional arrays). But since we know that our array is only
        #1-Dimensional, we can simply pick only the first dimension
        maxIndexInChunk = maxIndexInChunk[0]
        indexInData = rLeft + maxIndexInChunk

        #We are basically just repeating what the plot item also did to
        #determine the X-Axis labeling, and use it calculate the X-Value
        #at the index we have identified
        xAxisScale = (self.data.fftMaxFreq() - self.data.fftFreqDist()) / float(self.fftYValues.shape[0])

        self.freqText.setText(str(self.fftXValues[indexInData][0]*xAxisScale))
        self.isFreqValid = True



    def closeEvent(self, event):
        global __widget_id__
        __widget_id__ = None
        del self.par.widgets[self.id]

        self.board_config.unobserve(self, 'lastDataSet')


def addFrequencyExtractWidget():
    global __widget_id__
    if __widget_id__:
        global_objects.get_global('area').widgets[__widget_id__].setFocus()
    else:
        nid = kcgw.idg.genid()
        __widget_id__ = nid
        w = FrequencyExtractWidget(nid, global_objects.get_global('area'))
        global_objects.get_global('area').newWidget(w, "Frequency Extract", nid, widget_type=4)

kcgw.register_widget(QtGui.QIcon(config.icon_path("sproject.svg")), "Frequency Extract", addFrequencyExtractWidget, "Ctrl+e")