Clarified usage information for vib_analyze.py
[calibcant.git] / calibcant / vib_analyze.py
1 #!/usr/bin/python
2 #
3 # calibcant - tools for thermally calibrating AFM cantilevers
4 #
5 # Copyright (C) 2007-2009 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,D = 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, \t D = %g" % (A, B, C, D)
111
112     vib_plot(deflection_bits, freq_axis, power, A, B, C, D, **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=25000) :
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     D_guess = 0 # allow a noise floor offset (DISABLED in fitting below)
179     # 
180     # Half width w on lower side when L(a-w) = L(a)/2
181     #  (a**2 - (a-w)**2)**2 + (b*(a-w))**2 = 2*(b*a)**2
182     # Let W=(a-w)**2, A=a**2, and B=b**2
183     #  (A - W)**2 + BW = 2*AB
184     #  W**2 - 2AW + A**2 + BW = 2AB
185     #  W**2 + (B-2A)W + (A**2-2AB) = 0
186     #  W = (2A-B)/2 * [1 +/- sqrt(1 - 4(A**2-2AB)/(B-2A)**2]
187     #    = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
188     #  (a-w)**2 = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
189     #  so w is a disaster ;)
190     # For some values of A and B (non-underdamped), W is imaginary or negative.
191     #
192     # The mass m is given by m = k_B T / (pi A)
193     # The spring constant k is given by k = m (omega_0)**2
194     # The quality factor Q is given by Q = omega_0 m / gamma
195     if config.TEXT_VERBOSE : 
196         print "guesses : %g, %g, %g" % (A_guess, B_guess, C_guess)
197     if PLOT_GUESSED_LORENTZIAN:
198         vib_plot(None, freq_axis, psd_data, A_guess, B_guess, C_guess,
199                  minFreq, maxFreq, plotVerbose=True)
200
201     # Fitting the PSD of V = photoSensitivity*x just rescales the parameters
202
203     if False: # Gnuplot fit worse than scipy.optimize.leastsq fit.
204         # fit Lorentzian using Gnuplot's 'fit' command
205         
206         # save about-to-be-fitted data to a temporary file
207         datafilename = "%s_%d" % (config.GNUFIT_DATA_BASE, time.time())
208         fd = file(datafilename, 'w')
209         for i in range(imin, imax) :
210             fd.write("%g\t%g\n" % (freq_axis[i], psd_data[i]))
211         fd.close()
212         
213         g = GnuplotBiDir.Gnuplot()
214         g("f(x) = C / ((A**2-x**2)**2 + (B*x)**2)")
215         g("A=%g; B=%g; C=%g" % (A_guess, B_guess, C_guess))
216         g("fit f(x) '%s' via A,B,C" % datafilename)
217         A=abs(float(g.getvar('A'))) # A and B only show up as squares in f(x)
218         B=abs(float(g.getvar('B'))) # so ensure we get positive values
219         C=float(g.getvar('C'))
220         
221         os.remove(datafilename)
222     else:
223         # fit Lorenzian using scipy.optimize.leastsq
224         def residual(p, y, x):
225             A = p[0]
226             B = p[1]
227             C = p[2]
228             Y = C / ((A**2-x**2)**2 + (B*x)**2)
229             return Y-y
230         leastsq = scipy.optimize.leastsq
231         p,cov,info,mesg,ier = leastsq(residual, (A_guess, B_guess, C_guess),
232                                       args=(psd_data[imin:imax],
233                                             freq_axis[imin:imax]),
234                                       full_output=True, maxfev=10000)
235         if config.TEXT_VERBOSE:
236             print "Fitted params:",p
237             print "Covariance mx:",cov
238             #print "Info:", info  # too much information :p
239             print "mesg:", mesg
240             if ier == 1 :
241                 print "Solution converged"
242             else :
243                 print "Solution did not converge"
244         A,B,C = p
245         D = 0 # do not fit the noise floor (fit doesn't look very convincing)
246         A=abs(A) # A and B only show up as squares in f(x)
247         B=abs(B) # so ensure we get positive values
248     return (A, B, C, D)
249
250 def lorentzian_area(A, B, C) :
251     # Integrating the the power spectral density per unit time (power) over the
252     # frequency axis [0, fN] returns the total signal power per unit time
253     #  int_0^fN power(f)df = 1/T int_0^T |x(t)**2| dt
254     #                      = <V(t)**2>, the variance for our AC signal.
255     # The variance from our fitted Lorentzian is the area under the Lorentzian
256     #  <V(t)**2> = (pi*C) / (2*B*A**2)
257     return (numpy.pi * C) / (2 * B * A**2)
258
259 def gnuplot_check_fit(datafilename, A, B, C) :
260     """
261     return a string containing a gnuplot script displaying the fit.
262     """
263     string = 'f(x) = C / ((A**2-x**2)**2 + (x*B)**2)\n'
264     string += 'A= %g ; B= %g ; C= %g\n' % (A, B, C)
265     string += 'set logscale y\n'
266     string += "plot '%s', f(x)\n" % datafilename
267     return string
268
269 @splittableKwargsFunction()
270 def vib_save(data, log_dir=None) :
271     """Save the dictionary data, using data_logger.data_log()
272     data should be a dictionary of arrays with the fields
273       'Z piezo output'
274       'Z piezo input'
275       'Deflection input'
276     """
277     if log_dir :
278         freq = data['sample frequency Hz']
279         log = data_logger.data_log(log_dir, noclobber_logsubdir=False,
280                                    log_name="vib_%gHz" % freq)
281         log.write_dict_of_arrays(data)
282
283 def vib_load(datafile=None) :
284     """Load the dictionary data, using data_logger.date_load()"
285     data should be a dictionary of arrays with the fields
286       'Z piezo output'
287       'Z piezo input'
288       'Deflection input'
289     """
290     dl = data_logger.data_load()
291     data = dl.read_dict_of_arrays(datafile)
292     return data
293
294 @splittableKwargsFunction()
295 def vib_plot(deflection_bits, freq_axis, power, A, B, C, D,
296              minFreq=None, maxFreq=None, plotVerbose=False) :
297     """
298     If plotVerbose or config.PYLAB_VERBOSE == True, plot 3 subfigures:
299      Time series (Vphoto vs sample index) (show any major drift events),
300      A histogram of Vphoto, with a gaussian fit (demonstrate randomness), and
301      FFTed Vphoto data (Vphoto vs Frequency) (show major noise components).
302     """
303     if plotVerbose or config.PYLAB_VERBOSE :
304         common._import_pylab()
305         common._pylab.figure(config.BASE_FIGNUM+2)
306
307         if deflection_bits != None:
308             # plot time series
309             common._pylab.subplot(311)
310             common._pylab.hold(False)
311             common._pylab.plot(deflection_bits, 'r.')
312             common._pylab.title("free oscillation")
313     
314             # plot histogram distribution and gaussian fit
315             common._pylab.subplot(312)
316             common._pylab.hold(False)
317             n, bins, patches = \
318                 common._pylab.hist(deflection_bits, bins=30,
319                                    normed=1, align='center')
320             gaus = numpy.zeros((len(bins),), dtype=numpy.float)
321             mean = deflection_bits.mean()
322             std = deflection_bits.std()
323             pi = numpy.pi
324             exp = numpy.exp
325             for i in range(len(bins)) :
326                 gaus[i] = (2*pi)**(-.5)/std * exp(-0.5*((bins[i]-mean)/std)**2)
327             common._pylab.hold(True)
328             common._pylab.plot(bins, gaus, 'r-');
329             common._pylab.hold(False)
330     
331             # plot FFTed data
332             axes = common._pylab.subplot(313)
333         else:
334             # use a nice big subplot just for the FFTed data
335             axes = common._pylab.subplot(111)
336         common._pylab.hold(False)
337         common._pylab.semilogy(freq_axis, power, 'r.-')
338         common._pylab.hold(True)
339         xmin,xmax = axes.get_xbound()
340         ymin,ymax = axes.get_ybound()
341         
342         if minFreq is not None and maxFreq is not None:
343             # highlight the region we're fitting in
344             common._pylab.axvspan(minFreq, maxFreq, facecolor='g', alpha=0.1,
345                                   zorder = -1)
346         
347         fitdata = C / ((A**2-freq_axis**2)**2 + (B*freq_axis)**2) + D
348         common._pylab.plot(freq_axis, fitdata, 'b-')
349         noisefloor = D + 0*freq_axis;
350         common._pylab.plot(freq_axis, noisefloor)
351         common._pylab.hold(False)
352         axes.axis([xmin,xmax,ymin,ymax])
353
354         common._flush_plot()
355     if (plotVerbose or config.GNUPLOT_VERBOSE) and False: # TODO: cleanup and test
356         # write all the ft data now
357         fd = file(datafilename, 'w')
358         for i in range(len(freq_axis)) :
359             fd.write("%g\t%g\n" % (freq_axis[i], power[i]))
360         fd.write("\n") # blank line to separate fit data.
361         # write the fit data
362         for i in range(imin, imax) :
363             x = freq_axis[i]
364             fd.write("%g\t%g\n" % (freq_axis[i],
365                                    C / ((A**2-x**2)**2 + (x*B)**2) ) )
366         fd.close()
367         print Vphoto_t.var(), A, B, C, ((numpy.pi * C)/(2 * B * A**2) * 1e-12)
368         g("set terminal epslatex color solid")
369         g("set output 'calibration.tex'")
370         g("set size 2,2") # multipliers on default 5"x3.5"
371         #g("set title 'Thermal calibration'")
372         g("set logscale y")
373         g("set xlabel 'Frequency (Hz)'")
374         g("set ylabel 'Power ($\mu$V$^2$/Hz)'")
375         g("set xrange [0:20000]")
376         g("g(x) = x < %g || x > %g ? 1/0 : f(x)" % (minFreq, maxFreq)) # only plot on fitted region
377         g("plot '%s' title 'Vibration data', g(x) title 'Lorentzian fit'" % datafilename)
378
379         g("set mouse")
380         g("pause mouse 'click with mouse'")
381         g.getvar("MOUSE_BUTTON")
382         del(g)
383
384 # Make vib_analyze splittable (was postponed until the child functions
385 # were defined).
386 make_splittable_kwargs_function(vib_analyze,
387     (fit_psd, 'freq_axis', 'psd_data'),
388     (vib_plot, 'deflection_bits', 'freq_axis', 'power', 'A', 'B', 'C', 'D',
389      'minFreq', 'maxFreq'))
390
391 @splittableKwargsFunction((vib_analyze, 'deflection_bits', 'freq'))
392 def vib_load_analyze_tweaked(tweak_file, **kwargs) :
393     """
394     Read the vib files listed in tweak_file.
395     Return an array of Vphoto variances (due to the cantilever) in Volts**2
396     """
397     vib_analyze_kwargs, = \
398         vib_load_analyze_tweaked._splitargs(vib_load_analyze_tweaked, kwargs)
399     dl = data_logger.data_load()
400     Vphoto_var = []
401     for path in file(tweak_file, 'r') :
402         path = path.strip()
403         if path[0] == '#': # a comment
404             continue
405         # read the data
406         data = vib_load(path)
407         deflection_bits = data['Deflection input']
408         freq = float(path.split('_')[-1].split('Hz')[0])
409         if config.TEXT_VERBOSE :
410             print "Analyzing '%s' at %g Hz" % (path, freq)
411         # analyze
412         Vphoto_var.append(vib_analyze(deflection_bits, freq,
413                                       **vib_analyze_kwargs))
414     return numpy.array(Vphoto_var, dtype=numpy.float)
415
416 # commandline interface functions
417 import scipy.io, sys
418
419 def read_data(ifile):
420     """
421     ifile can be a filename string or open (seekable) file object.
422     returns an array of data values (1 column)
423     """
424     #returns (column 1 array, column 2 array)
425     #"""
426     if ifile == None :  ifile = sys.stdin
427     unlabeled_data=scipy.io.read_array(ifile)
428     return unlabeled_data #(unlabeled_data[:,0], unlabeled_data[:,1])
429
430 if __name__ == '__main__' :
431     # command line interface
432     from optparse import OptionParser
433     
434     usage_string = ('%prog <input-file>\n'
435                     '2007-2009 W. Trevor King.\n'
436                     '\n'
437                     'There are several operation modes, each handling a different <input-file>\n'
438                     'type.  In each case, the output is the fitted variance (in square Volts).\n'
439                     '\n'
440                     'Single file mode without datalogger mode (the default) :\n'
441                     '<input-file> should be a 1 column ASCII without a header line of cantilever\n'
442                     'deflection (in bits)\n'
443                     'e.g: "<deflection bits>\\n..."\n'
444                     '\n'
445                     'Single file mode with datalogger mode :\n'
446                     'Same as without datalogger mode, except the input should be a datalogger file\n'
447                     'with data stored in bits (e.g. as saved by the unfold module).\n'
448                     '\n'
449                     'Tweak file mode:\n'
450                     'Runs the same analysis as in single file mode for each vib file in a tweak\n'
451                     'file.  Each line in the tweak file specifies a single vib file.  Blank lines\n'
452                     'and those beginning with a pound sign (#) are ignored.  Vib files are valid\n'
453                     'datalogger files as per single-file-with-datalogger-mode.\n'
454                     )
455     parser = OptionParser(usage=usage_string, version='%prog 0.1')
456     parser.add_option('-f', '--sample-freq', dest='freq', default=100e3,
457                       help='sample frequency in Hz (default %default)',
458                       type='float', metavar='FREQ')
459     parser.add_option('-m', '--min-freq', dest='min_freq', default=500,
460                       help='minimum frequency in Hz (default %default)',
461                       type='float', metavar='FREQ')
462     parser.add_option('-M', '--max-freq', dest='max_freq', default=7000,
463                       help='maximum frequency in Hz (default %default)',
464                       type='float', metavar='FREQ')
465     parser.add_option('-o', '--output-file', dest='ofilename',
466                       help='write output to FILE (default stdout)',
467                       type='string', metavar='FILE')
468     parser.add_option('-c', '--comma-out', dest='comma_out', action='store_true',
469                       help='Output comma-seperated values (default %default)',
470                       default=False)
471     parser.add_option('-g', '--gnuplot', dest='gnuplot', action='store_true',
472                       help='Print gnuplot fit check script to stderr',
473                       default=False)
474     parser.add_option('-p', '--pylab', dest='pylab', action='store_true',
475                       help='Produce pylab fit checks during execution',
476                       default=False)
477     parser.add_option('-G', '--plot-guess', dest='plot_guess', action='store_true',
478                       help='Produce plots of the pre-fit Lorentzian',
479                       default=False)
480     parser.add_option('-t', '--tweak-mode', dest='tweakmode', action='store_true',
481                       help='Run in tweak-file mode',
482                       default=False)
483     parser.add_option('-d', '--datalogger-mode', dest='datalogger_mode', action='store_true',
484                       help='Read input files with datalogger.read_dict_of_arrays().  This is useful, for example, to test a single line from a tweakfile.',
485                       default=False)
486     parser.add_option('-v', '--verbose', dest='verbose', action='store_true',
487                       help='Print lots of debugging information',
488                       default=False)
489
490     options,args = parser.parse_args()
491     parser.destroy()
492     assert len(args) >= 1, "Need an input file"
493         
494     ifilename = args[0]
495
496     if options.ofilename != None :
497         ofile = file(options.ofilename, 'w')
498     else :
499         ofile = sys.stdout
500     config.TEXT_VERBOSE = options.verbose
501     config.PYLAB_INTERACTIVE = False
502     config.PYLAB_VERBOSE = options.pylab
503     config.GNUPLOT_VERBOSE = options.gnuplot
504     PLOT_GUESSED_LORENTZIAN = options.plot_guess
505
506     if options.tweakmode == False : # single file mode
507         if options.datalogger_mode:
508             data = vib_load(ifilename)
509             deflection_bits = data['Deflection input']
510         else:
511             deflection_bits = read_data(ifilename)
512         Vphoto_var = vib_analyze(deflection_bits, freq=options.freq,
513                                  minFreq=options.min_freq,
514                                  maxFreq=options.max_freq)
515         print >> ofile, Vphoto_var
516     else : # tweak mode
517         Vphoto_vars = vib_load_analyze_tweaked(ifilename,
518                                                minFreq=options.min_freq,
519                                                maxFreq=options.max_freq)
520         if options.comma_out :
521             sep = ','
522         else :
523             sep = '\n'
524         common.write_array(ofile, Vphoto_vars, sep)
525
526     if common._final_flush_plot != None:
527         common._final_flush_plot()
528     
529     if options.ofilename != None :
530         ofile.close()