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