# HG changeset patch
# User David Douard <david.douard@logilab.fr>
# Date 1202842502 -3600
# Node ID 235686915f923bf1b7f5ba9e4173f51cf71478fa
# Parent  cafc62915d390404197c320382024474a1943315
add THD computation for simple traces of power spectrum

diff -r cafc62915d39 -r 235686915f92 pygpibtoolkit/HP3562A/datablockwidget.py
--- a/pygpibtoolkit/HP3562A/datablockwidget.py	Mon Feb 11 23:39:21 2008 +0100
+++ b/pygpibtoolkit/HP3562A/datablockwidget.py	Tue Feb 12 19:55:02 2008 +0100
@@ -9,6 +9,7 @@
 import  trace_decoder
 import  coord_decoder
 from pygpibtoolkit.qt4.mpl import QMplCanvas
+from pygpibtoolkit.HP3562A import mathtools
 
 children = [ ]
 
@@ -206,6 +207,7 @@
         l.setMargin(0)
         self.canvas = QMplCanvas(self)
         l.addWidget(self.canvas, 1)
+        
         self.setCentralWidget(mainw)
         
     def updateTraceData(self):
@@ -234,7 +236,29 @@
         self.canvas.axes.grid(True)
         if self.header[0]['Display function']:
             self.canvas.axes.set_title(self.header[0]['Display function'])
-        
+
+        y = self.trace.copy()
+        if f0 > 0:
+            # must add some initial zeros
+            yy = numpy.zeros(f0/dx+len(y))
+            yy[-len(y):] = y
+            y = yy
+        msg = ""
+        try:
+            f0, thd = mathtools.thd(y, db=True)
+            f0 = f0*dx
+            assert thd
+        except:
+            pass
+        else:
+            msg += 'THD:%.2g db  Freq:%.2f Hz  '%(thd, f0)
+        try:
+            thdn = mathtools.thd_n(y, db=True)
+        except:
+            pass
+        else:
+            msg += 'THD+N:%.2g db  '%thdn
+        self.statusBar().showMessage(msg)
         self.canvas.repaint() 
 
 class CoordBinaryDatablockMDIChild(TraceBinaryDatablockMDIChild):
diff -r cafc62915d39 -r 235686915f92 pygpibtoolkit/HP3562A/mathtools.py
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/pygpibtoolkit/HP3562A/mathtools.py	Tue Feb 12 19:55:02 2008 +0100
@@ -0,0 +1,61 @@
+import numpy
+
+def thd(spectrum, squared=True, db=True):
+    """
+    compute THD from a spectrum data, given by a simple 1D data
+    vector.
+
+    The result is the couple (fmax, thd), this latter being expressed
+    in 'db' if db is True; if not, in percentage.
+    
+    """
+    n = len(spectrum)
+    # first, find the fundamental frequency
+    fmax = spectrum.argmax()
+    vmax = spectrum[fmax]
+
+    # then, the compute the several max values for harmonics
+    h_freqs = numpy.arange(0, n, fmax)
+    h_vals = []
+    h_freqs2 = [fmax]
+    for i, f in enumerate(h_freqs[2:]):
+        fm = spectrum[f-5:f+5].argmax()
+        h_vals.append(spectrum[f-5+fm])
+        h_freqs2.append((f-5.+fm)/(i+2))
+    h_vals = numpy.array(h_vals)
+    h_freqs2 = numpy.array(h_freqs2)
+    if not squared:
+        h_vals = h_vals**2
+        vmax = vmax**2
+    thd = h_vals.sum()/vmax
+    if db:
+        thd = 10*numpy.log10(thd)
+    else:
+        thd = 100.0*numpy.sqrt(thd)
+    return h_freqs2.mean(), thd
+
+def thd_n(spectrum, squared=True, db=True):
+    """
+    compute THD+N from a spectrum data, given by a simple 1D data
+    vector.
+
+    The result is the thdN, expressed in 'db' if db is True; if not,
+    in percentage.
+    
+    """
+    if not squared:
+        spectrum = spectrum**2
+    fmax = spectrum.argmax()
+    w = int(fmax/10) # argh
+    vmax = spectrum[fmax-w:fmax+w].sum()
+    print  spectrum[fmax-w],  spectrum[fmax], spectrum[fmax+w]
+    thd = (spectrum.sum()-vmax)/(vmax)
+    if db:
+        thd = 10*numpy.log10(thd)
+    else:
+        thd = 100.0*numpy.sqrt(thd)
+    return thd
+
+
+    
+    
diff -r cafc62915d39 -r 235686915f92 pygpibtoolkit/HP3562A/trace_decoder.py
--- a/pygpibtoolkit/HP3562A/trace_decoder.py	Mon Feb 11 23:39:21 2008 +0100
+++ b/pygpibtoolkit/HP3562A/trace_decoder.py	Tue Feb 12 19:55:02 2008 +0100
@@ -114,13 +114,12 @@
 
     if options.displayheader:
         print format_header(header, HEADER, 100)
+    f0 = header['Start freq value']
+    dx = header['Delta X-axis']
+    n = header['Number of elements']
+    x = numpy.linspace(f0, f0+dx*n, len(data)) 
+    y = data.copy()
     if options.plot:
-        f0 = header['Start freq value']
-        dx = header['Delta X-axis']
-        n = header['Number of elements']
-        x = numpy.linspace(f0, f0+dx*n, len(data)) 
-        y = data.copy()
-
         import pylab
         if options.ymode != "lin":
             minv = min(y[y>0])
@@ -132,8 +131,8 @@
         pylab.grid()
         pylab.plot(x, y)
         pylab.xlabel('frequency')
-        pylab.show()
-    
-          
+        pylab.show()    
+    return header, x, y
+
 if __name__ == "__main__":
-    main()
+    h, x, y = main()