rtspectrum.py

#!/usr/bin/env python3

# Provides "fast" waveform and spectrum visualization for SDR signals using pyqtgraph.
# The idea is that you change the get_samples() function for whatever you need to 
# obtain you own data for analisys.
# Hopefully this can be useful, of at least, fun for someone else ;) 

# TODO: 
#   * Command line args / UI Selection for Plot parameters
#   * Grid and legend on the "Digital fosphor display"
#   * Possibility to use numba/cuSignal for faster processing
#   * Create Watefall display
#   * Proper widowing for the fosfor plot
#  

# Inspired by:
# https://teaandtechtime.com/python-intensity-graded-fft-plots/
# https://stackoverflow.com/questions/40126176/fast-live-plotting-in-matplotlib-pyplot
# https://dl.cdn-anritsu.com/en-us/test-measurement/files/Technical-Notes/White-Paper/11410-01138B.pdf
# http://www.sjsu.edu/people/burford.furman/docs/me120/FFT_tutorial_NI.pdf

import sys, time

import numpy as np
from numpy.fft import fftshift, fft, ifft
from pyqtgraph.colormap import ColorMap
import scipy as sp
import scipy.signal as signal

from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph as pg

# fft = np.fft.fft
# fftshift = np.fft.fftshift

# Performance options (not tested)
#pg.setConfigOption('useNumba', True)              # Use numba   
#pg.setConfigOption('imageAxisOrder', 'row-major') # best performance
#pg.setConfigOption('useCupy', True)               # Use cupy if available

class App(QtGui.QMainWindow):
    def __init__(self, parent=None):
        super(App, self).__init__(parent)

        # Plot Parameters
        self.imgXBins  = 1024    # FFT Size of the Fosfor plot 
        self.imgYMaxdB = 10      # Max value for the fosphor plot in dB
        self.imgYMindB = -140    # Min value for the fosphor plot in dB
        self.displayPesistance = 0.9  # Fosphor display persistance
        
        #### Create Gui Elements ###########
        self.mainbox = QtGui.QWidget()
        self.setCentralWidget(self.mainbox)
        self.mainbox.setLayout(QtGui.QVBoxLayout())

        self.canvas = pg.GraphicsLayoutWidget()
        self.mainbox.layout().addWidget(self.canvas)

        self.label = QtGui.QLabel()
        self.mainbox.layout().addWidget(self.label)
      
        # Digital fosphor image plot
        self.imv = pg.ImageView()
        self.mainbox.layout().addWidget(self.imv)
        self.imv.setPredefinedGradient('flame')  # thermal, flame, inferno, yellowy, magma
        # self.imv.setLevels(0, 16)       
        self.imv.view.invertY(False)

        # Time domain I and Q plots
        self.timePlot = self.canvas.addPlot()
        self.realTrace = self.timePlot.plot(pen='y')
        self.imagTrace = self.timePlot.plot(pen='c')

        # Frequency domain plot
        self.freqPlot  = self.canvas.addPlot()
        self.freqTrace = self.freqPlot.plot(pen='m')

        # Initialization
        self.counter = 0
        self.fps = 0.
        self.lastupdate = time.time()
        self._update()
    

    def _update(self) :
        """ Function to accquire data and update the graph display """
        ## Replace with the function to get your samples ===================
        x = signal.resample(generate_samples(256),1024)
        x = x* ( 1 + 0.0001*np.abs(x)**2)    # Some non-linear distortion
        x.real = np.clip(x.real, -2.9, 2.9)  # Some clipping
        x.imag = np.clip(x.imag, -2.9, 2.9)  # Some clipping
        ## =================================================================

        s = 20*np.log10(np.abs(fftshift(fft(x, norm='ortho'))))
        mag = 20*np.log10(np.abs(fftshift(fft(x[:self.imgXBins], norm='ortho'))))
        mag  = np.rint(mag.flat)
        idx  = np.arange(self.imgXBins)

        # Remove points out of imgScale
        xpos = idx[(mag>=self.imgYMindB)&(mag<=self.imgYMaxdB)]
        ypos = np.take(mag, xpos) - self.imgYMindB

        self.imgYBins  = self.imgYMaxdB - self.imgYMindB + 1
        # Hit histogram in a sparse matrix for cheaper updates
        hitmap = sp.sparse.csc_matrix(
            (np.ones(len(xpos)),       # Value to the Hit postions
            (xpos, ypos)),             # x and  non zero postions
            shape=(self.imgXBins,self.imgYBins), 
            dtype=np.float32
        )

        r = self.displayPesistance
        try :
            hitmap = (1.0-r) * hitmap + r * self.oldmap
        except:
            self.oldmap = hitmap.tocoo()    

        hitmap /= hitmap.max()
        self.imv.setImage(hitmap.toarray(), autoLevels=True)
        self.oldmap = hitmap.tocoo()

        # Update data on line plots        
        self.realTrace.setData(x.real.flat)
        self.imagTrace.setData(x.imag.flat)
        self.freqTrace.setData(s.flat)

        # Calculate FPS rate
        now = time.time()
        dt = (now-self.lastupdate)
        if dt <= 0:
            dt = 1e-6
        fps2 = 1.0 / dt
        self.lastupdate = now
        self.fps = self.fps * 0.9 + fps2 * 0.1
        tx = 'Mean Frame Rate:  {fps:.3f} FPS'.format(fps=self.fps )
        self.label.setText(tx)
        QtCore.QTimer.singleShot(1, self._update)
        self.counter += 1

def generate_samples(numSamples: int=4096) :
    """Generate random complex signal """
    x = np.zeros(numSamples,dtype=np.complex64)
    x.real = np.sqrt(2)**-1*np.random.randn(numSamples)
    x.imag = np.sqrt(2)**-1*np.random.randn(numSamples)
    return x

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    thisapp = App()
    thisapp.show()
    sys.exit(app.exec_())