Getting Frequency and FFT plots with THD+N figure from Bitscope DDR file

bitscope.py

#!/usr/bin/python3
#
# bitscope.py: 
# (c) 2016 Derrik Walker v2.0
#
# Python Library for bitscope functions
#
# This is the Library goes with the blog post:
#
# http://www.doomd.net/2016/10/the-bitscope-linux-and-python.html 
#
# This is licensed for use under the GNU General Pulbic License v2
#
# 2016-08-10	DW2	Initial Version
# 2016-09-08	DW2	Fixed a bug in the file parsing
# 2016-10-12	DW2	Initial release
#

import sys, csv
import numpy as np

class dso_data:

	channel0 = 0
	channel1 = 1

	sig0 = []
	sig1 = []

	rate0 = None
	rate1 = None

	count0 = None
	count1 = None

	data0 = []
	data1 = []

	def __init__( self, fname ):

		try:
    			f = open(fname, 'r')

		except IOError:
    			print ( "Error reading file:", fname )
    			sys.exit()

		with f:
			reader = csv.reader( f )
			next ( reader )	# skip the header

			for row in reader:
				if int( row[2] ) == dso_data.channel0:
					self.rate0 = float( row[7] )
					self.count0 = float( row[8] )
					self.data0.extend( [ float( x ) for x in row[9:len(row)] ])
					
				if int( row[2] ) == dso_data.channel1:
					self.rate1 = float( row[7] )
					self.count1 = float( row[8] )
					self.data1.extend( [ float( x ) for x in row[9:len(row)] ])

		self.sig0 = np.array( self.data0 )
		self.sig1 = np.array( self.data1 )

	def data( self, channel ):
		# returns the raw data
		
		if( channel == 0 ):
			return self.data0

		elif ( channel == 1 ):
			return self.data0

		else:
			return None

	def sig( self, channel ):
		# returns the numpy array of data
		
		if( channel == 0 ):
			return self.sig0

		elif ( channel == 1 ):
			return self.sig1

		else:
			return None

	def rate( self, channel ):
		# returns the sampling rate

		if( channel == 0 ):
			return self.rate0

		elif( channel == 1 ):
			return self.rate1

		else:
			return None

dso_THD_N.py

#!/usr/bin/python3
#
# dso_THD_N.py: 
# (c) 2016 Derrik Walker v2.0
# adapted by Marco Duiker
#
# Generate an fft plot data 'recorded' from the Bitscope DSO program and the dso_data class from bitscope.py. 
# Add THD + N to the title
#
# usage:
#
#	dso_fft.py [-c <channel>] -f <file>.csv
#	
#	Where <file.csv> is a "recorded" output from the Bitscope DSO program
#
# This is the scipt goes with the blog post:
#
# http://www.doomd.net/2016/10/the-bitscope-linux-and-python.html 
#
#
# And is enhanced by Marco Duiker with https://gist.github.com/endolith/246092
#
#
# This is licensed for use under the GNU General Pulbic License v2
#
# 2016-06-15	DW2	Initial Version
# 2016-07-04	DW2	Added cmd line support for specifying the channel
# 2016-08-10	DW2	Moved bitscope spacific funtions to bitscope.py
# 2016-10-12	DW2	Initial release
#

from __future__ import division
import sys
from scipy.signal import blackmanharris
from numpy.fft import rfft, irfft, rfftfreq
from numpy import argmax, sqrt, mean, absolute, arange, log10
import numpy as np
import getopt
import matplotlib.pyplot as plt

import bitscope

def rms_flat(a):
    """
    Return the root mean square of all the elements of *a*, flattened out.
    """
    return sqrt(mean(absolute(a)**2))


def find_range(f, x):
    """
    Find range between nearest local minima from peak at index x
    """
    for i in arange(x+1, len(f)):
        if f[i+1] >= f[i]:
            uppermin = i
            break
    for i in arange(x-1, 0, -1):
        if f[i] <= f[i-1]:
            lowermin = i + 1
            break
    return (lowermin, uppermin)


def THDN(signal, sample_rate):
    """
    Measure the THD+N for a signal and print the results

    Prints the estimated fundamental frequency and the measured THD+N.  This is
    calculated from the ratio of the entire signal before and after
    notch-filtering.

    Currently this tries to find the "skirt" around the fundamental and notch
    out the entire thing.  A fixed-width filter would probably be just as good,
    if not better.
    """
    # Get rid of DC and window the signal
    signal -= mean(signal) # TODO: Do this in the frequency domain, and take any skirts with it?
    windowed = signal * blackmanharris(len(signal))  # TODO Kaiser?

    # Measure the total signal before filtering but after windowing
    total_rms = rms_flat(windowed)

    # Find the peak of the frequency spectrum (fundamental frequency), and
    # filter the signal by throwing away values between the nearest local
    # minima
    f = rfft(windowed)
    i = argmax(abs(f))

    frequency = sample_rate * (i / len(windowed))
    #print 'Frequency: %f Hz' % (frequency)  # Not exact

    lowermin, uppermin = find_range(abs(f), i)
    f[lowermin: uppermin] = 0

    # Transform noise back into the signal domain and measure it
    # TODO: Could probably calculate the RMS directly in the frequency domain instead
    noise = irfft(f)
    THDN = rms_flat(noise) / total_rms
    #print "THD+N:     %.4f%% or %.1f dB" % (THDN * 100, 20 * log10(THDN))

    return THDN, frequency

def usage():
	print( "Usage: dso_fft [-c <channel> ] -f <file.csv>, where <file.csv> is the recorded output from the Bitscope's DSO" )
	exit(1)

file = None
channel = 0 

try:
	opts, args = getopt.getopt( sys.argv[1:], "hc:f:", ["help", "channel=", "file=" ] )

except getopt.GetoptError as err:
	print(err) 
	usage()

for o, a in opts:

	if  o in ( "-c", "--channel" ):
		channel = int( a ) 

	elif o in ("-h", "--help"):
		usage()

	elif o in ("-f", "--file"):
		file =  a
	else:
		assert False, "ERRRR"	

if not file:
	usage()

if channel > 1 or channel < 0:
	print( "channel out of range!" )
	usage()

dso = bitscope.dso_data( file )

data = dso.data( channel )	# raw data as Python array
rate = dso.rate( channel )	# just the rate
sig = dso.sig( channel )	# a numpy array of the data

if rate == None:
	print( "It appears that channel", channel, "has no data!" )
	exit()

time = [ x/rate for x in range( 0, len(data) ) ]

fft = abs( rfft( sig ) ) / (  len(sig)/2 ) 
freq = abs( rfftfreq( sig.size, d=1/rate ) )

thdn, frequency = THDN(sig, rate)

plt.subplot( 211 )
plt.xlabel( "Time (sec)" )
plt.ylabel( "Voltage" )
plt.title( "Original Signal" )
plt.plot( time, dso.data0 )
 
plt.subplot( 212 )
plt.xlabel( "Frequency (Hz) " + "(Fundamental frequency: %s Hz)" % round(frequency) )
plt.ylabel( "Magnitude" )
plt.title( "FFT of Signal" + " (THD+N:  %.4f%% or %.1f dB)" % (thdn * 100, 20 * log10(thdn)) )
plt.plot( freq, fft ) 
 
plt.tight_layout()
plt.show()