calibcant/vib_analyze.py

   1 #!/usr/bin/python
   2 #
   3 # calibcant - tools for thermally calibrating AFM cantilevers
   4 #
   5 # Copyright (C) 2007,2008, William Trevor King
   6 #
   7 # This program is free software; you can redistribute it and/or
   8 # modify it under the terms of the GNU General Public License as
   9 # published by the Free Software Foundation; either version 3 of the
  10 # License, or (at your option) any later version.
  11 #
  12 # This program is distributed in the hope that it will be useful, but
  13 # WITHOUT ANY WARRANTY; without even the implied warranty of
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
  15 # See the GNU General Public License for more details.
  16 #
  17 # You should have received a copy of the GNU General Public License
  18 # along with this program; if not, write to the Free Software
  19 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
  20 # 02111-1307, USA.
  21 #
  22 # The author may be contacted at <wking@drexel.edu> on the Internet, or
  23 # write to Trevor King, Drexel University, Physics Dept., 3141 Chestnut St.,
  24 # Philadelphia PA 19104, USA.
  25
  26 """
  27 Separate the more general vib_analyze() from the other vib_*()
  28 functions in calibcant.  Also provide a command line interface
  29 for analyzing data acquired through other workflows.
  30
  31 The relevent physical quantities are :
  32  Vphoto   The photodiode vertical deflection voltage (what we measure)
  33 """
  34
  35 import os, time, numpy
  36 import GnuplotBiDir   # used for fitting lorentzian
  37 import scipy.optimize # alternative for fitting lorentzian
  38 import common # common module for the calibcant package
  39 import config # config module for the calibcant package
  40 import data_logger
  41 import FFT_tools
  42 from splittable_kwargs import splittableKwargsFunction, \
  43     make_splittable_kwargs_function
  44
  45 PLOT_GUESSED_LORENTZIAN=False
  46
  47 @splittableKwargsFunction()
  48 def vib_analyze_naive(deflection_bits, Vphoto_in2V=config.Vphoto_in2V,
  49                       zeroVphoto_bits=config.zeroVphoto_bits) :
  50     """
  51     Calculate the variance of the raw data, and convert it to Volts**2.
  52     This method is simple and easy to understand, but it highly succeptible
  53     to noise, drift, etc.
  54
  55     Vphoto_in2V is a function converting Vphoto values from bits to Volts.
  56     This function is assumed linear, see bump_analyze().
  57     """
  58     Vphoto_std_bits = deflection_bits.std()
  59     Vphoto_std = Vphoto_in2V(Vphoto_std_bits + zeroVphoto_bits)
  60     return Vphoto_std**2
  61
  62 #@splittableKwargsFunction((fit_psd, 'freq_axis', 'psd_data'),
  63 #                          (vib_plot, 'deflection_bits', 'freq_axis', 'power',
  64 #                           'A', 'B', 'C', 'minFreq', 'maxFreq'))
  65 # Some of the child functions aren't yet defined, so postpone
  66 # make-splittable until later in the module.
  67 def vib_analyze(deflection_bits, freq,
  68                 chunk_size=4096, overlap=True, window=FFT_tools.window_hann,
  69                 Vphoto_in2V=config.Vphoto_in2V, **kwargs) :
  70     """
  71     Calculate the variance in the raw data due to the cantilever's
  72     thermal oscillation and convert it to Volts**2.
  73
  74     Improves on vib_analyze_naive() by first converting Vphoto(t) to
  75     frequency space, and fitting a Lorentzian in the relevant
  76     frequency range (see cantilever_calib.pdf for derivation).
  77     However, there may be cases where the fit is thrown off by noise
  78     spikes in frequency space.  To protect from errors, the fitted
  79     variance is compared to the naive variance (where all noise is
  80     included), and the minimum variance is returned.
  81
  82     Vphoto_in2V is a function converting Vphoto values from bits to Volts.
  83     This function is assumed linear, see bump_analyze().
  84     """
  85     fit_psd_kwargs,vib_plot_kwargs = vib_analyze._splitargs(vib_analyze,kwargs)
  86     if 'minFreq' in fit_psd_kwargs:
  87         vib_plot_kwargs['minFreq'] = fit_psd_kwargs['minFreq']
  88     if 'maxFreq' in fit_psd_kwargs:
  89         vib_plot_kwargs['maxFreq'] = fit_psd_kwargs['maxFreq']
  90
  91     Vphoto_t = numpy.zeros((len(deflection_bits),),
  92                            dtype=numpy.float)
  93     # convert the data from bits to volts
  94     if config.TEXT_VERBOSE :
  95         print "Converting %d bit values to voltages" % len(Vphoto_t)
  96     for i in range(len(deflection_bits)) :
  97         Vphoto_t[i] = Vphoto_in2V(deflection_bits[i])
  98
  99     # Compute the average power spectral density per unit time (in uV**2/Hz)
 100     if config.TEXT_VERBOSE :
 101         print "Compute the averaged power spectral density in uV**2/Hz"
 102     (freq_axis, power) = \
 103         FFT_tools.unitary_avg_power_spectrum((Vphoto_t - Vphoto_t.mean())*1e6,
 104                                              freq, chunk_size,overlap, window)
 105
 106     A,B,C = fit_psd(freq_axis, power, **kwargs)
 107
 108     if config.TEXT_VERBOSE :
 109         print "Fitted f(x) = C / ((A**2-x**2)**2 + (x*B)**2) with"
 110         print "A = %g, \t B = %g, \t C = %g" % (A, B, C)
 111
 112     vib_plot(deflection_bits, freq_axis, power, A, B, C, **kwargs)
 113
 114     # Our A is in uV**2, so convert back to Volts**2
 115     fitted_variance = lorentzian_area(A,B,C) * 1e-12
 116
 117     naive_variance = vib_analyze_naive(deflection_bits,
 118                                        Vphoto_in2V=Vphoto_in2V)
 119     if config.TEXT_VERBOSE:
 120         print "Fitted variance: %g V**2", fitted_variance
 121         print "Naive variance : %g V**2", naive_variance
 122
 123     return min(fitted_variance, naive_variance)
 124
 125 @splittableKwargsFunction()
 126 def fit_psd(freq_axis, psd_data, minFreq=500, maxFreq=7000) :
 127     """
 128     freq_axis : array of frequencies in Hz
 129     psd_data  : array of PSD amplitudes for the frequencies in freq_axis
 130
 131     Calculate the variance in the raw data due to the cantilever's
 132     thermal oscillation and convert it to Volts**2.
 133
 134     The conversion to frequency space generates an average power spectrum
 135     by breaking the data into windowed chunks and averaging the power
 136     spectrums for the chunks together.
 137     See FFT_tools.unitary_avg_power_spectrum().
 138     """
 139     # cut out the relevent frequency range
 140     if config.TEXT_VERBOSE :
 141         print "Cut the frequency range %g to %g Hz" % (minFreq, maxFreq)
 142     imin = 0
 143     while freq_axis[imin] < minFreq : imin += 1
 144     imax = imin
 145     while freq_axis[imax] < maxFreq : imax += 1
 146     assert imax >= imin + 10 , \
 147         "Less than 10 points in freq range (%g,%g)" % (minFreq, maxFreq)
 148
 149     # generate guesses for Lorentzian parameters A,B,C
 150     max_psd_index = numpy.argmax(psd_data[imin:imax]) + imin
 151     max_psd = psd_data[max_psd_index]
 152     half_max_index = imin
 153     while psd_data[half_max_index] < max_psd/2 :
 154         half_max_index += 1
 155     res_freq = freq_axis[max_psd_index]
 156     half_width = freq_axis[max_psd_index] - freq_axis[half_max_index]
 157     # Lorentzian L(x) = C / ((A**2-x**2)**2 + (B*x)**2)
 158     # is expected power spectrum for
 159     # x'' + B x' + A^2 x'' = F_external(t)/m
 160     # (A = omega_0)
 161     # C = (2 k_B T B) / (pi m)
 162     #
 163     # A = resonant frequency
 164     # peak at  x_res = sqrt(A^2 - B^2/2)
 165     #  which could be complex if there isn't a peak (overdamped)
 166     # peak height    = C / (3 x_res^4 + A^4)
 167     # Q = A/B
 168     #
 169     # guess Q = 1, and adjust from there
 170     Q_guess = 1
 171     # so x_res^2 = B^2 Q^2 - B^2/2 = (Q^2-1/2)B^2
 172     #    B = x_res / sqrt(Q^2-1/2)
 173     if config.TEXT_VERBOSE :
 174         print "params : %g, %g" % (res_freq, max_psd)
 175     B_guess = res_freq / numpy.sqrt(Q_guess**2-0.5)
 176     A_guess = Q_guess*B_guess
 177     C_guess = max_psd * (-res_freq**4 + A_guess**4)
 178     #
 179     # Half width w on lower side when L(a-w) = L(a)/2
 180     #  (a**2 - (a-w)**2)**2 + (b*(a-w))**2 = 2*(b*a)**2
 181     # Let W=(a-w)**2, A=a**2, and B=b**2
 182     #  (A - W)**2 + BW = 2*AB
 183     #  W**2 - 2AW + A**2 + BW = 2AB
 184     #  W**2 + (B-2A)W + (A**2-2AB) = 0
 185     #  W = (2A-B)/2 * [1 +/- sqrt(1 - 4(A**2-2AB)/(B-2A)**2]
 186     #    = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
 187     #  (a-w)**2 = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
 188     #  so w is a disaster ;)
 189     # For some values of A and B (non-underdamped), W is imaginary or negative.
 190     #
 191     # The mass m is given by m = k_B T / (pi A)
 192     # The spring constant k is given by k = m (omega_0)**2
 193     # The quality factor Q is given by Q = omega_0 m / gamma
 194     if config.TEXT_VERBOSE :
 195         print "guesses : %g, %g, %g" % (A_guess, B_guess, C_guess)
 196     if PLOT_GUESSED_LORENTZIAN:
 197         vib_plot(None, freq_axis, psd_data, A_guess, B_guess, C_guess,
 198                  minFreq, maxFreq, plotVerbose=True)
 199
 200     # Fitting the PSD of V = photoSensitivity*x just rescales the parameters
 201
 202     if False: # Gnuplot fit worse than scipy.optimize.leastsq fit.
 203         # fit Lorentzian using Gnuplot's 'fit' command
 204
 205         # save about-to-be-fitted data to a temporary file
 206         datafilename = "%s_%d" % (config.GNUFIT_DATA_BASE, time.time())
 207         fd = file(datafilename, 'w')
 208         for i in range(imin, imax) :
 209             fd.write("%g\t%g\n" % (freq_axis[i], psd_data[i]))
 210         fd.close()
 211
 212         g = GnuplotBiDir.Gnuplot()
 213         g("f(x) = C / ((A**2-x**2)**2 + (B*x)**2)")
 214         g("A=%g; B=%g; C=%g" % (A_guess, B_guess, C_guess))
 215         g("fit f(x) '%s' via A,B,C" % datafilename)
 216         A=abs(float(g.getvar('A'))) # A and B only show up as squares in f(x)
 217         B=abs(float(g.getvar('B'))) # so ensure we get positive values
 218         C=float(g.getvar('C'))
 219
 220         os.remove(datafilename)
 221     else:
 222         # fit Lorenzian using scipy.optimize.leastsq
 223         def residual(p, y, x):
 224             A = p[0]
 225             B = p[1]
 226             C = p[2]
 227             Y = C / ((A**2-x**2)**2 + (B*x)**2)
 228             return Y-y
 229         leastsq = scipy.optimize.leastsq
 230         p,cov,info,mesg,ier = leastsq(residual, (A_guess, B_guess, C_guess),
 231                                       args=(psd_data[imin:imax],
 232                                             freq_axis[imin:imax]),
 233                                       full_output=True, maxfev=10000)
 234         if config.TEXT_VERBOSE:
 235             print "Fitted params:",p
 236             print "Covariance mx:",cov
 237             #print "Info:", info  # too much information :p
 238             print "mesg:", mesg
 239             if ier == 1 :
 240                 print "Solution converged"
 241             else :
 242                 print "Solution did not converge"
 243         A,B,C = p
 244         A=abs(A) # A and B only show up as squares in f(x)
 245         B=abs(B) # so ensure we get positive values
 246     return (A, B, C)
 247
 248 def lorentzian_area(A, B, C) :
 249     # Integrating the the power spectral density per unit time (power) over the
 250     # frequency axis [0, fN] returns the total signal power per unit time
 251     #  int_0^fN power(f)df = 1/T int_0^T |x(t)**2| dt
 252     #                      = <V(t)**2>, the variance for our AC signal.
 253     # The variance from our fitted Lorentzian is the area under the Lorentzian
 254     #  <V(t)**2> = (pi*C) / (2*B*A**2)
 255     return (numpy.pi * C) / (2 * B * A**2)
 256
 257 def gnuplot_check_fit(datafilename, A, B, C) :
 258     """
 259     return a string containing a gnuplot script displaying the fit.
 260     """
 261     string = 'f(x) = C / ((A**2-x**2)**2 + (x*B)**2)\n'
 262     string += 'A= %g ; B= %g ; C= %g\n' % (A, B, C)
 263     string += 'set logscale y\n'
 264     string += "plot '%s', f(x)\n" % datafilename
 265     return string
 266
 267 @splittableKwargsFunction()
 268 def vib_save(data, log_dir=None) :
 269     """Save the dictionary data, using data_logger.data_log()
 270     data should be a dictionary of arrays with the fields
 271       'Z piezo output'
 272       'Z piezo input'
 273       'Deflection input'
 274     """
 275     if log_dir :
 276         freq = data['sample frequency Hz']
 277         log = data_logger.data_log(log_dir, noclobber_logsubdir=False,
 278                                    log_name="vib_%gHz" % freq)
 279         log.write_dict_of_arrays(data)
 280
 281 def vib_load(datafile=None) :
 282     """Load the dictionary data, using data_logger.date_load()"
 283     data should be a dictionary of arrays with the fields
 284       'Z piezo output'
 285       'Z piezo input'
 286       'Deflection input'
 287     """
 288     dl = data_logger.data_load()
 289     data = dl.read_dict_of_arrays(datafile)
 290     return data
 291
 292 @splittableKwargsFunction()
 293 def vib_plot(deflection_bits, freq_axis, power, A, B, C,
 294              minFreq=None, maxFreq=None, plotVerbose=False) :
 295     """
 296     If plotVerbose or config.PYLAB_VERBOSE == True, plot 3 subfigures:
 297      Time series (Vphoto vs sample index) (show any major drift events),
 298      A histogram of Vphoto, with a gaussian fit (demonstrate randomness), and
 299      FFTed Vphoto data (Vphoto vs Frequency) (show major noise components).
 300     """
 301     if plotVerbose or config.PYLAB_VERBOSE :
 302         common._import_pylab()
 303         common._pylab.figure(config.BASE_FIGNUM+2)
 304
 305         if deflection_bits != None:
 306             # plot time series
 307             common._pylab.subplot(311)
 308             common._pylab.hold(False)
 309             common._pylab.plot(deflection_bits, 'r.')
 310             common._pylab.title("free oscillation")
 311
 312             # plot histogram distribution and gaussian fit
 313             common._pylab.subplot(312)
 314             common._pylab.hold(False)
 315             n, bins, patches = \
 316                 common._pylab.hist(deflection_bits, bins=30,
 317                                    normed=1, align='center')
 318             gaus = numpy.zeros((len(bins),), dtype=numpy.float)
 319             mean = deflection_bits.mean()
 320             std = deflection_bits.std()
 321             pi = numpy.pi
 322             exp = numpy.exp
 323             for i in range(len(bins)) :
 324                 gaus[i] = (2*pi)**(-.5)/std * exp(-0.5*((bins[i]-mean)/std)**2)
 325             common._pylab.hold(True)
 326             common._pylab.plot(bins, gaus, 'r-');
 327             common._pylab.hold(False)
 328
 329             # plot FFTed data
 330             axes = common._pylab.subplot(313)
 331         else:
 332             # use a nice big subplot just for the FFTed data
 333             axes = common._pylab.subplot(111)
 334         common._pylab.hold(False)
 335         common._pylab.semilogy(freq_axis, power, 'r.-')
 336         common._pylab.hold(True)
 337         xmin,xmax = axes.get_xbound()
 338         ymin,ymax = axes.get_ybound()
 339
 340         if minFreq is not None and maxFreq is not None:
 341             # highlight the region we're fitting in
 342             common._pylab.axvspan(minFreq, maxFreq, facecolor='g', alpha=0.1,
 343                                   zorder = -1)
 344
 345         fitdata = C / ((A**2-freq_axis**2)**2 + (B*freq_axis)**2)
 346         common._pylab.plot(freq_axis, fitdata, 'b-');
 347         common._pylab.hold(False)
 348         axes.axis([xmin,xmax,ymin,ymax])
 349
 350         common._flush_plot()
 351     if (plotVerbose or config.GNUPLOT_VERBOSE) and False: # TODO: cleanup and test
 352         # write all the ft data now
 353         fd = file(datafilename, 'w')
 354         for i in range(len(freq_axis)) :
 355             fd.write("%g\t%g\n" % (freq_axis[i], power[i]))
 356         fd.write("\n") # blank line to separate fit data.
 357         # write the fit data
 358         for i in range(imin, imax) :
 359             x = freq_axis[i]
 360             fd.write("%g\t%g\n" % (freq_axis[i],
 361                                    C / ((A**2-x**2)**2 + (x*B)**2) ) )
 362         fd.close()
 363         print Vphoto_t.var(), A, B, C, ((numpy.pi * C)/(2 * B * A**2) * 1e-12)
 364         g("set terminal epslatex color solid")
 365         g("set output 'calibration.tex'")
 366         g("set size 2,2") # multipliers on default 5"x3.5"
 367         #g("set title 'Thermal calibration'")
 368         g("set logscale y")
 369         g("set xlabel 'Frequency (Hz)'")
 370         g("set ylabel 'Power ($\mu$V$^2$/Hz)'")
 371         g("set xrange [0:20000]")
 372         g("g(x) = x < %g || x > %g ? 1/0 : f(x)" % (minFreq, maxFreq)) # only plot on fitted region
 373         g("plot '%s' title 'Vibration data', g(x) title 'Lorentzian fit'" % datafilename)
 374
 375         g("set mouse")
 376         g("pause mouse 'click with mouse'")
 377         g.getvar("MOUSE_BUTTON")
 378         del(g)
 379
 380 # Make vib_analyze splittable (was postponed until the child functions
 381 # were defined).
 382 make_splittable_kwargs_function(vib_analyze,
 383     (fit_psd, 'freq_axis', 'psd_data'),
 384     (vib_plot, 'deflection_bits', 'freq_axis', 'power', 'A', 'B', 'C',
 385      'minFreq', 'maxFreq'))
 386
 387 @splittableKwargsFunction((vib_analyze, 'deflection_bits', 'freq'))
 388 def vib_load_analyze_tweaked(tweak_file, **kwargs) :
 389     """
 390     Read the vib files listed in tweak_file.
 391     Return an array of Vphoto variances (due to the cantilever) in Volts**2
 392     """
 393     vib_analyze_kwargs, = \
 394         vib_load_analyze_tweaked._splitargs(vib_load_analyze_tweaked, kwargs)
 395     dl = data_logger.data_load()
 396     Vphoto_var = []
 397     for path in file(tweak_file, 'r') :
 398         if path[-1] == '\n':
 399             path = path.split('\n')[0]
 400         # read the data
 401         data = vib_load(path)
 402         deflection_bits = data['Deflection input']
 403         freq = float(path.split('_')[-1].split('Hz')[0])
 404         if config.TEXT_VERBOSE :
 405             print "Analyzing '%s' at %g Hz" % (path, freq)
 406         # analyze
 407         Vphoto_var.append(vib_analyze(deflection_bits, freq,
 408                                       **vib_analyze_kwargs))
 409     return numpy.array(Vphoto_var, dtype=numpy.float)
 410
 411 # commandline interface functions
 412 import scipy.io, sys
 413
 414 def read_data(ifile):
 415     """
 416     ifile can be a filename string or open (seekable) file object.
 417     returns an array of data values (1 column)
 418     """
 419     #returns (column 1 array, column 2 array)
 420     #"""
 421     if ifile == None :  ifile = sys.stdin
 422     unlabeled_data=scipy.io.read_array(ifile)
 423     return unlabeled_data #(unlabeled_data[:,0], unlabeled_data[:,1])
 424
 425 if __name__ == '__main__' :
 426     # command line interface
 427     from optparse import OptionParser
 428
 429     usage_string = ('%prog <input-file>\n'
 430                     '2008, W. Trevor King.\n'
 431                     '\n'
 432                     'Deflection power spectral density (PSD) data (X^2/Hz)\n'
 433                     'and returns the variance in X units (X^2)\n'
 434                     '<input-file> should be 1 column ASCII without a header line.\n'
 435                     'e.g: "<deflection bits>\\n..."\n')
 436     parser = OptionParser(usage=usage_string, version='%prog 0.1')
 437     parser.add_option('-f', '--sample-freq', dest='freq', default=100e3,
 438                       help='sample frequency in Hz (default %default)',
 439                       type='float', metavar='FREQ')
 440     parser.add_option('-m', '--min-freq', dest='min_freq', default=500,
 441                       help='minimum frequency in Hz (default %default)',
 442                       type='float', metavar='FREQ')
 443     parser.add_option('-M', '--max-freq', dest='max_freq', default=7000,
 444                       help='maximum frequency in Hz (default %default)',
 445                       type='float', metavar='FREQ')
 446     parser.add_option('-o', '--output-file', dest='ofilename',
 447                       help='write output to FILE (default stdout)',
 448                       type='string', metavar='FILE')
 449     parser.add_option('-c', '--comma-out', dest='comma_out', action='store_true',
 450                       help='Output comma-seperated values (default %default)',
 451                       default=False)
 452     parser.add_option('-g', '--gnuplot', dest='gnuplot', action='store_true',
 453                       help='Print gnuplot fit check script to stderr',
 454                       default=False)
 455     parser.add_option('-p', '--pylab', dest='pylab', action='store_true',
 456                       help='Produce pylab fit checks during execution',
 457                       default=False)
 458     parser.add_option('-G', '--plot-guess', dest='plot_guess', action='store_true',
 459                       help='Produce plots of the pre-fit Lorentzian',
 460                       default=False)
 461     parser.add_option('-t', '--tweak-mode', dest='tweakmode', action='store_true',
 462                       help='Run in tweak-file mode',
 463                       default=False)
 464     parser.add_option('-d', '--datalogger-mode', dest='datalogger_mode', action='store_true',
 465                       help='Read input files with datalogger.read_dict_of_arrays().  This is useful, for example, to test a single line from a tweakfile.',
 466                       default=False)
 467     parser.add_option('-v', '--verbose', dest='verbose', action='store_true',
 468                       help='Print lots of debugging information',
 469                       default=False)
 470
 471     options,args = parser.parse_args()
 472     parser.destroy()
 473     assert len(args) >= 1, "Need an input file"
 474
 475     ifilename = args[0]
 476
 477     if options.ofilename != None :
 478         ofile = file(options.ofilename, 'w')
 479     else :
 480         ofile = sys.stdout
 481     config.TEXT_VERBOSE = options.verbose
 482     config.PYLAB_INTERACTIVE = False
 483     config.PYLAB_VERBOSE = options.pylab
 484     config.GNUPLOT_VERBOSE = options.gnuplot
 485     PLOT_GUESSED_LORENTZIAN = options.plot_guess
 486
 487     if options.tweakmode == False :
 488         if options.datalogger_mode:
 489             data = vib_load(ifilename)
 490             deflection_bits = data['Deflection input']
 491         else:
 492             deflection_bits = read_data(ifilename)
 493         Vphoto_var = vib_analyze(deflection_bits, freq=options.freq,
 494                                  minFreq=options.min_freq,
 495                                  maxFreq=options.max_freq)
 496         print >> ofile, Vphoto_var
 497     else : # tweak mode
 498         Vphoto_vars = vib_load_analyze_tweaked(ifilename,
 499                                                minFreq=options.min_freq,
 500                                                maxFreq=options.max_freq)
 501         if options.comma_out :
 502             sep = ','
 503         else :
 504             sep = '\n'
 505         common.write_array(ofile, Vphoto_vars, sep)
 506
 507     if common._final_flush_plot != None:
 508         common._final_flush_plot()
 509
 510     if options.ofilename != None :
 511         ofile.close()
 512