[Matplotlib-users] PSD amplitudes

##----------------------------------------------------------------------------
## Name: psd_scale.py
## 
## Purpose: Test Power Spectral Density of 1Vrms data
## Depends on Python SciPy and NumPy
## 
## Author: J Park
##
## Created: 10/17/07
##
## Modified: 
##----------------------------------------------------------------------------
try:
 from numpy import * # www.numpy.org numpy.scipy.org
except ImportError: 
 print "Failed to import numpy."
 
try:
 import pylab as mp # matplotlib.sourceforge.net
 from matplotlib.font_manager import fontManager, FontProperties 
except ImportError: 
 print "Failed to import pylab."
 
# Default Parameters
nFFT = 1024 
overlap = 512 
freqSample = 100. 
PlotAll = False
WriteOutput = False
##----------------------------------------------------------------------------
## Main module
def main():
 deltaF = freqSample/nFFT # Frequency resolution in Hz
 deltaT = 1./freqSample # Sample interval
 print 'Sample interval %e (s)' % (deltaT)
 print 'Frequency resolution %e (Hz)' % (deltaF)
 # Setup Plots
 # ----------------------------------------------------------------------
 mp.figure(1)
 mp.title ( "PSD" )
 mp.ylabel( "(dB)" )
 mp.xlabel( "Frequency (Hz)" )
 legendFont = FontProperties(size='small')
 ymin = 0
 ymax = 30
 xmin = 0
 xmax = 50
 xticks = 5
 yticks = 5
 if PlotAll:
 mp.figure(2)
 mp.title ( "Input Timeseries" )
 mp.ylabel( "Amplitude" )
 mp.xlabel( "time (s)" )
 # Create some synthetic data with unity RMS amplitude = 0 dB
 # ----------------------------------------------------------------------
 t = mp.arange(0., 60., deltaT) # 60 seconds at deltaT interval
 A = 1.414
 
 y0 = A * sin( 2. * math.pi * 5 * t )
 y1 = A * sin( 2. * math.pi * 10 * t )
 y2 = A * sin( 2. * math.pi * 20 * t )
 y3 = A * sin( 2. * math.pi * 30 * t )
 y4 = A * sin( 2. * math.pi * 40 * t )
 y5 = A * sin( 2. * math.pi * 45 * t )
 dataList = [ y0, y1, y2, y3, y4, y5 ]
 
 for data in dataList:
 inputDataLen = len( data )
 numAverages = math.floor( inputDataLen / (overlap) ) - 1
 normalizedRandomError = 1./math.sqrt( numAverages )
 print "%d points" % ( inputDataLen ),
 print "%d averages" % (numAverages),
 print "normalized random error %.3f" % ( normalizedRandomError )
 mp.figure(1)
 (Pxx, freqs) = mp.psd( data,
 NFFT = nFFT,
 Fs = freqSample,
 noverlap = overlap,
 lw = 2,
 label = '' )
 Pxx_dB = 10.*log10(Pxx)
 
 if PlotAll:
 mp.figure(2)
 mp.plot(t, data, label='' )
 # Write Output data
 # ----------------------------------------------------------------------
 if WriteOutput:
 PxxLen = len(Pxx)
 OutputFile = "PSD.dat"
 fdOutFile = open( OutputFile, 'a' )
 fdOutFile.write( "Freq\t\tPower(dB)\n" )
 for i in range(PxxLen):
 fdOutFile.write( "%.4e\t%.3f\n" % ( freqs[i], Pxx_dB[i] ) )
 fdOutFile.close()
 print "Wrote ", PxxLen, " points to ", OutputFile
 
 # Show the Plot
 # ----------------------------------------------------------------------
 mp.figure(1)
 mp.axis([xmin, xmax, ymin, ymax])
 mp.xticks( arange(xmin, xmax+1, xticks) )
 mp.yticks( arange(ymin, ymax , yticks) )
 mp.title('')
 mp.xlabel('Frequency (Hz)')
 mp.ylabel(r'$\tt{dB re V^2/Hz}$')
 #mp.legend( loc='upper right', prop=legendFont )
 if WriteOutput:
 plotFileName = "PSD.png"
 mp.savefig( plotFileName )
 print "Wrote png image to ", plotFileName
 if PlotAll:
 mp.figure(2)
 #mp.legend( loc='lower left', prop=legendFont )
 mp.show()
 print "Normal Exit"
## Main module
##----------------------------------------------------------------------------
##----------------------------------------------------------------------------
## Provide for cmd line invocation
if __name__ == "__main__":
 main()