Added comparison to vib_analyze_naive() in vib_analyze().
[calibcant.git] / calibcant / vib_analyze.py
index cfa2c3e90e45663f409010cf18e1b9ef80359ac6..72c2140fecf6c9183408065506082e402c6b943b 100755 (executable)
@@ -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
@@ -50,10 +51,9 @@ def vib_analyze_naive(deflection_bits, Vphoto_in2V=config.Vphoto_in2V,
     
     Vphoto_in2V is a function converting Vphoto values from bits to Volts.
     This function is assumed linear, see bump_analyze().
-    zeroVphoto_bits is the bit value corresponding to Vphoto = zero Volts.
     """
     Vphoto_std_bits = deflection_bits.std()
-    Vphoto_std = Vphoto_in2V(Vphoto_std_bits + zpiezo_zeroV_in)
+    Vphoto_std = Vphoto_in2V(Vphoto_std_bits + zeroVphoto_bits)
     return Vphoto_std**2
     
 def vib_analyze(deflection_bits, freq, minFreq=500, maxFreq=7000,
@@ -63,18 +63,16 @@ def vib_analyze(deflection_bits, freq, minFreq=500, maxFreq=7000,
     Calculate the variance in the raw data due to the cantilever's 
     thermal oscillation and convert it to Volts**2.
 
-    Improves on vib_analyze_naive() by first converting Vphoto(t) to 
-    frequency space, and fitting a Lorentzian in the relevant frequency
-    range (see cantilever_calib.pdf for derivation).
+    Improves on vib_analyze_naive() by first converting Vphoto(t) to
+    frequency space, and fitting a Lorentzian in the relevant
+    frequency range (see cantilever_calib.pdf for derivation).
+    However, there may be cases where the fit is thrown off by noise
+    spikes in frequency space.  To protect from errors, the fitted
+    variance is compared to the naive variance (where all noise is
+    included), and the minimum variance is returned.
 
-    The conversion to frequency space generates an average power spectrum
-    by breaking the data into windowed chunks and averaging the power
-    spectrums for the chunks together.
-    See FFT_tools.avg_power_spectrum
-    
     Vphoto_in2V is a function converting Vphoto values from bits to Volts.
     This function is assumed linear, see bump_analyze().
-    zeroVphoto_bits is the bit value corresponding to Vphoto = zero Volts.
     """
     Vphoto_t = numpy.zeros((len(deflection_bits),),
                            dtype=numpy.float)
@@ -98,10 +96,18 @@ def vib_analyze(deflection_bits, freq, minFreq=500, maxFreq=7000,
         print "A = %g, \t B = %g, \t C = %g" % (A, B, C)
 
     vib_plot(deflection_bits, freq_axis, power, A, B, C,
-             plotVerbose=plotVerbose)
+             minFreq, maxFreq, plotVerbose=plotVerbose)
 
     # Our A is in uV**2, so convert back to Volts**2
-    return lorentzian_area(A,B,C) * 1e-12
+    fitted_variance = lorentzian_area(A,B,C) * 1e-12
+
+    naive_variance = vib_analyze_naive(deflection_bits,
+                                       Vphoto_in2V=Vphoto_in2V)
+    if config.TEXT_VERBOSE:
+        print "Fitted variance: %g V**2", fitted_variance
+        print "Naive variance : %g V**2", naive_variance
+    
+    return min(fitted_variance, naive_variance)
 
 def fit_psd(freq_axis, psd_data, minFreq=500, maxFreq=7000) :
     """
@@ -111,14 +117,10 @@ def fit_psd(freq_axis, psd_data, minFreq=500, maxFreq=7000) :
     Calculate the variance in the raw data due to the cantilever's 
     thermal oscillation and convert it to Volts**2.
 
-    Improves on vib_analyze_naive() by working on frequency space data 
-    and fitting a Lorentzian in the relevant frequency range (see
-    cantilever_calib.pdf for derivation).
-
     The conversion to frequency space generates an average power spectrum
     by breaking the data into windowed chunks and averaging the power
     spectrums for the chunks together.
-    See FFT_tools.unitary_avg_power_spectrum().    
+    See FFT_tools.unitary_avg_power_spectrum().
     """
     # cut out the relevent frequency range
     if config.TEXT_VERBOSE : 
@@ -130,13 +132,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 +173,62 @@ 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,
-                 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
+                 minFreq, maxFreq, plotVerbose=True)
+
     # 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
+        A=abs(A) # A and B only show up as squares in f(x)
+        B=abs(B) # so ensure we get positive values
     return (A, B, C)
 
 def lorentzian_area(A, B, C) :
@@ -248,7 +274,7 @@ def vib_load(datafile=None) :
     return data
 
 def vib_plot(deflection_bits, freq_axis, power, A, B, C,
-             plotVerbose=True) :
+             minFreq=None, maxFreq=None, plotVerbose=True) :
     """
     If plotVerbose or config.PYLAB_VERBOSE == True, plot 3 subfigures:
      Time series (Vphoto vs sample index) (show any major drift events),
@@ -290,10 +316,16 @@ def vib_plot(deflection_bits, freq_axis, power, A, B, C,
             axes = common._pylab.subplot(111)
         common._pylab.hold(False)
         common._pylab.semilogy(freq_axis, power, 'r.-')
-        fitdata = C / ((A**2-freq_axis**2)**2 + (B*freq_axis)**2)
+        common._pylab.hold(True)
         xmin,xmax = axes.get_xbound()
         ymin,ymax = axes.get_ybound()
-        common._pylab.hold(True)
+        
+        if minFreq is not None and maxFreq is not None:
+            # highlight the region we're fitting in
+            common._pylab.axvspan(minFreq, maxFreq, facecolor='g', alpha=0.1,
+                                  zorder = -1)
+        
+        fitdata = C / ((A**2-freq_axis**2)**2 + (B*freq_axis)**2)
         common._pylab.plot(freq_axis, fitdata, 'b-');
         common._pylab.hold(False)
         axes.axis([xmin,xmax,ymin,ymax])
@@ -376,10 +408,9 @@ if __name__ == '__main__' :
                     '2008, W. Trevor King.\n'
                     '\n'
                     'Deflection power spectral density (PSD) data (X^2/Hz)\n'
-                    'and returns the variance in X units (X^2)'
-                    '<input-file> should be whitespace-delimited, 2 column ASCII\n'
-                    'without a header line.\n'
-                    'e.g: "<frequency_Hz>\\t<deflection_psd X^2/Hz>\\n"\n')
+                    'and returns the variance in X units (X^2)\n'
+                    '<input-file> should be 1 column ASCII without a header line.\n'
+                    'e.g: "<deflection bits>\\n..."\n')
     parser = OptionParser(usage=usage_string, version='%prog 0.1')
     parser.add_option('-f', '--sample-freq', dest='freq', default=100e3,
                       help='sample frequency in Hz (default %default)',
@@ -408,6 +439,9 @@ if __name__ == '__main__' :
     parser.add_option('-t', '--tweak-mode', dest='tweakmode', action='store_true',
                       help='Run in tweak-file mode',
                       default=False)
+    parser.add_option('-d', '--datalogger-mode', dest='datalogger_mode', action='store_true',
+                      help='Read input files with datalogger.read_dict_of_arrays().  This is useful, for example, to test a single line from a tweakfile.',
+                      default=False)
     parser.add_option('-v', '--verbose', dest='verbose', action='store_true',
                       help='Print lots of debugging information',
                       default=False)
@@ -433,8 +467,11 @@ if __name__ == '__main__' :
     PLOT_GUESSED_LORENTZIAN = options.plot_guess
 
     if options.tweakmode == False :
-        data = read_data(ifilename)
-        deflection_bits = data['Deflection input']
+        if options.datalogger_mode:
+            data = vib_load(ifilename)
+            deflection_bits = data['Deflection input']
+        else:
+            deflection_bits = read_data(ifilename)
         Vphoto_var = vib_analyze(deflection_bits, freq=options.freq,
                                  minFreq=options.min_freq,
                                  maxFreq=options.max_freq)