Moved fitting from GnuplotBiDir to scipy.optimize.leastsq in vib_analyze.

[calibcant.git] / calibcant / vib_analyze.py

diff --git a/calibcant/vib_analyze.py b/calibcant/vib_analyze.py
index 283f1ae49d8b389cf549c497162c1c32e6ae7f0a..b8e8c7c3397e50ebab2af581a660c64e14a0efb9 100755 (executable)
--- a/calibcant/vib_analyze.py
+++ b/calibcant/vib_analyze.py
@@ -33,7 +33,8 @@ The relevent physical quantities are :
 """
 
 import os, time, numpy
-import GnuplotBiDir  # used for fitting lorentzian
+import GnuplotBiDir   # used for fitting lorentzian
+import scipy.optimize # alternative for fitting lorentzian
 import common # common module for the calibcant package
 import config # config module for the calibcant package
 import data_logger
@@ -130,13 +131,6 @@ def fit_psd(freq_axis, psd_data, minFreq=500, maxFreq=7000) :
     assert imax >= imin + 10 , \
         "Less than 10 points in freq range (%g,%g)" % (minFreq, maxFreq)
 
-    # save about-to-be-fitted data to a temporary file
-    datafilename = "%s_%d" % (config.GNUFIT_DATA_BASE, time.time())
-    fd = file(datafilename, 'w')
-    for i in range(imin, imax) :
-        fd.write("%g\t%g\n" % (freq_axis[i], psd_data[i]))
-    fd.close()
-
     # generate guesses for Lorentzian parameters A,B,C
     max_psd_index = numpy.argmax(psd_data[imin:imax]) + imin
     max_psd = psd_data[max_psd_index]
@@ -178,31 +172,60 @@ def fit_psd(freq_axis, psd_data, minFreq=500, maxFreq=7000) :
     #  (a-w)**2 = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
     #  so w is a disaster ;)
     # For some values of A and B (non-underdamped), W is imaginary or negative.
+    #
+    # The mass m is given by m = k_B T / (pi A)
+    # The spring constant k is given by k = m (omega_0)**2
+    # The quality factor Q is given by Q = omega_0 m / gamma
     if config.TEXT_VERBOSE : 
         print "guesses : %g, %g, %g" % (A_guess, B_guess, C_guess)
     if PLOT_GUESSED_LORENTZIAN:
         vib_plot(None, freq_axis, psd_data, A_guess, B_guess, C_guess,
                  minFreq, maxFreq, plotVerbose=True)
 
-    # fit Lorentzian using Gnuplot's 'fit' command
-    g = GnuplotBiDir.Gnuplot()
-    # The Lorentzian is the predicted one-sided power spectral density
-    # For an oscillator whose position x obeys
-    # m x'' + gamma x' + kx = F_thermal(t)
-    #  A is the ?driving noise?
-    #  B is gamma, the damping term
-    #  C is omega_0, the resonant frequency
     # Fitting the PSD of V = photoSensitivity*x just rescales the parameters
-    g("f(x) = C / ((A**2-x**2)**2 + (B*x)**2)")
-    ## The mass m is given by m = k_B T / (pi A)
-    ## The spring constant k is given by k = m (omega_0)**2
-    ## The quality factor Q is given by Q = omega_0 m / gamma
-    g("A=%g; B=%g; C=%g" % (A_guess, B_guess, C_guess))
-    g("fit f(x) '%s' via A,B,C" % datafilename)
-    A=abs(float(g.getvar('A'))) # A and B only show up as squares in f(x)
-    B=abs(float(g.getvar('B'))) # so ensure we get positive values
-    C=float(g.getvar('C'))
-    os.remove(datafilename)
+
+    if False: # Gnuplot fit worse than scipy.optimize.leastsq fit.
+        # fit Lorentzian using Gnuplot's 'fit' command
+        
+        # save about-to-be-fitted data to a temporary file
+        datafilename = "%s_%d" % (config.GNUFIT_DATA_BASE, time.time())
+        fd = file(datafilename, 'w')
+        for i in range(imin, imax) :
+            fd.write("%g\t%g\n" % (freq_axis[i], psd_data[i]))
+        fd.close()
+        
+        g = GnuplotBiDir.Gnuplot()
+        g("f(x) = C / ((A**2-x**2)**2 + (B*x)**2)")
+        g("A=%g; B=%g; C=%g" % (A_guess, B_guess, C_guess))
+        g("fit f(x) '%s' via A,B,C" % datafilename)
+        A=abs(float(g.getvar('A'))) # A and B only show up as squares in f(x)
+        B=abs(float(g.getvar('B'))) # so ensure we get positive values
+        C=float(g.getvar('C'))
+        
+        os.remove(datafilename)
+    else:
+        # fit Lorenzian using scipy.optimize.leastsq
+        def residual(p, y, x):
+            A = p[0]
+            B = p[1]
+            C = p[2]
+            Y = C / ((A**2-x**2)**2 + (B*x)**2)
+            return Y-y
+        leastsq = scipy.optimize.leastsq
+        p,cov,info,mesg,ier = leastsq(residual, (A_guess, B_guess, C_guess),
+                                      args=(psd_data[imin:imax],
+                                            freq_axis[imin:imax]),
+                                      full_output=True, maxfev=10000)
+        if config.TEXT_VERBOSE :
+            print "Fitted params:",p
+            print "Covariance mx:",cov
+            print "Info:", info
+            print "mesg:", mesg
+            if ier == 1 :
+                print "Solution converged"
+            else :
+                print "Solution did not converge"
+        A,B,C = p
     return (A, B, C)
 
 def lorentzian_area(A, B, C) :