Changed guessed Q from 5 to 1, which gives about the right breadth.
[calibcant.git] / 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 common # common module for the calibcant package
38 import config # config module for the calibcant package
39 import data_logger
40 import FFT_tools
41
42 PLOT_GUESSED_LORENTZIAN=False
43
44 def vib_analyze_naive(deflection_bits, Vphoto_in2V=config.Vphoto_in2V,
45                       zeroVphoto_bits=config.zeroVphoto_bits) :
46     """
47     Calculate the variance of the raw data, and convert it to Volts**2.
48     This method is simple and easy to understand, but it highly succeptible
49     to noise, drift, etc.
50     
51     Vphoto_in2V is a function converting Vphoto values from bits to Volts.
52     This function is assumed linear, see bump_analyze().
53     zeroVphoto_bits is the bit value corresponding to Vphoto = zero Volts.
54     """
55     Vphoto_std_bits = deflection_bits.std()
56     Vphoto_std = Vphoto_in2V(Vphoto_std_bits + zpiezo_zeroV_in)
57     return Vphoto_std**2
58     
59 def vib_analyze(deflection_bits, freq, minFreq=500, maxFreq=7000,
60                 chunk_size=4096, overlap=True, window=FFT_tools.window_hann,
61                 Vphoto_in2V=config.Vphoto_in2V, plotVerbose=False) :
62     """
63     Calculate the variance in the raw data due to the cantilever's 
64     thermal oscillation and convert it to Volts**2.
65
66     Improves on vib_analyze_naive() by first converting Vphoto(t) to 
67     frequency space, and fitting a Lorentzian in the relevant frequency
68     range (see cantilever_calib.pdf for derivation).
69
70     The conversion to frequency space generates an average power spectrum
71     by breaking the data into windowed chunks and averaging the power
72     spectrums for the chunks together.
73     See FFT_tools.avg_power_spectrum
74     
75     Vphoto_in2V is a function converting Vphoto values from bits to Volts.
76     This function is assumed linear, see bump_analyze().
77     zeroVphoto_bits is the bit value corresponding to Vphoto = zero Volts.
78     """
79     Vphoto_t = numpy.zeros((len(deflection_bits),),
80                            dtype=numpy.float)
81     # convert the data from bits to volts
82     if config.TEXT_VERBOSE : 
83         print "Converting %d bit values to voltages" % len(Vphoto_t)
84     for i in range(len(deflection_bits)) :
85         Vphoto_t[i] = Vphoto_in2V(deflection_bits[i])
86
87     # Compute the average power spectral density per unit time (in uV**2/Hz)
88     if config.TEXT_VERBOSE : 
89         print "Compute the averaged power spectral density in uV**2/Hz"
90     (freq_axis, power) = \
91         FFT_tools.unitary_avg_power_spectrum((Vphoto_t - Vphoto_t.mean())*1e6,
92                                              freq, chunk_size,overlap, window)
93
94     A,B,C = fit_psd(freq_axis, power, minFreq, maxFreq)
95
96     if config.TEXT_VERBOSE : 
97         print "Fitted f(x) = C / ((A**2-x**2)**2 + (x*B)**2) with"
98         print "A = %g, \t B = %g, \t C = %g" % (A, B, C)
99
100     vib_plot(deflection_bits, freq_axis, power, A, B, C,
101              plotVerbose=plotVerbose)
102
103     # Our A is in uV**2, so convert back to Volts**2
104     return lorentzian_area(A,B,C) * 1e-12
105
106 def fit_psd(freq_axis, psd_data, minFreq=500, maxFreq=7000) :
107     """
108     freq_axis : array of frequencies in Hz
109     psd_data  : array of PSD amplitudes for the frequencies in freq_axis
110
111     Calculate the variance in the raw data due to the cantilever's 
112     thermal oscillation and convert it to Volts**2.
113
114     Improves on vib_analyze_naive() by working on frequency space data 
115     and fitting a Lorentzian in the relevant frequency range (see
116     cantilever_calib.pdf for derivation).
117
118     The conversion to frequency space generates an average power spectrum
119     by breaking the data into windowed chunks and averaging the power
120     spectrums for the chunks together.
121     See FFT_tools.unitary_avg_power_spectrum().    
122     """
123     # cut out the relevent frequency range
124     if config.TEXT_VERBOSE : 
125         print "Cut the frequency range %g to %g Hz" % (minFreq, maxFreq)
126     imin = 0
127     while freq_axis[imin] < minFreq : imin += 1
128     imax = imin
129     while freq_axis[imax] < maxFreq : imax += 1
130     assert imax >= imin + 10 , \
131         "Less than 10 points in freq range (%g,%g)" % (minFreq, maxFreq)
132
133     # save about-to-be-fitted data to a temporary file
134     datafilename = "%s_%d" % (config.GNUFIT_DATA_BASE, time.time())
135     fd = file(datafilename, 'w')
136     for i in range(imin, imax) :
137         fd.write("%g\t%g\n" % (freq_axis[i], psd_data[i]))
138     fd.close()
139
140     # generate guesses for Lorentzian parameters A,B,C
141     max_psd_index = numpy.argmax(psd_data[imin:imax]) + imin
142     max_psd = psd_data[max_psd_index]
143     half_max_index = imin
144     while psd_data[half_max_index] < max_psd/2 :
145         half_max_index += 1
146     res_freq = freq_axis[max_psd_index]
147     half_width = freq_axis[max_psd_index] - freq_axis[half_max_index]
148     # Lorentzian L(x) = C / ((A**2-x**2)**2 + (B*x)**2)
149     # is expected power spectrum for
150     # x'' + B x' + A^2 x'' = F_external(t)/m
151     # (A = omega_0)
152     # C = (2 k_B T B) / (pi m)
153     # 
154     # A = resonant frequency
155     # peak at  x_res = sqrt(A^2 - B^2/2)
156     #  which could be complex if there isn't a peak (overdamped)
157     # peak height    = C / (3 x_res^4 + A^4)
158     # Q = A/B
159     #
160     # guess Q = 1, and adjust from there
161     Q_guess = 1
162     # so x_res^2 = B^2 Q^2 - B^2/2 = (Q^2-1/2)B^2 
163     #    B = x_res / sqrt(Q^2-1/2)
164     if config.TEXT_VERBOSE : 
165         print "params : %g, %g" % (res_freq, max_psd)
166     B_guess = res_freq / numpy.sqrt(Q_guess**2-0.5)
167     A_guess = Q_guess*B_guess
168     C_guess = max_psd * (-res_freq**4 + A_guess**4)
169     # 
170     # Half width w on lower side when L(a-w) = L(a)/2
171     #  (a**2 - (a-w)**2)**2 + (b*(a-w))**2 = 2*(b*a)**2
172     # Let W=(a-w)**2, A=a**2, and B=b**2
173     #  (A - W)**2 + BW = 2*AB
174     #  W**2 - 2AW + A**2 + BW = 2AB
175     #  W**2 + (B-2A)W + (A**2-2AB) = 0
176     #  W = (2A-B)/2 * [1 +/- sqrt(1 - 4(A**2-2AB)/(B-2A)**2]
177     #    = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
178     #  (a-w)**2 = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
179     #  so w is a disaster ;)
180     # For some values of A and B (non-underdamped), W is imaginary or negative.
181     if config.TEXT_VERBOSE : 
182         print "guesses : %g, %g, %g" % (A_guess, B_guess, C_guess)
183     if PLOT_GUESSED_LORENTZIAN:
184         vib_plot(None, freq_axis, psd_data, A_guess, B_guess, C_guess,
185                  plotVerbose=True)
186
187     # fit Lorentzian using Gnuplot's 'fit' command
188     g = GnuplotBiDir.Gnuplot()
189     # The Lorentzian is the predicted one-sided power spectral density
190     # For an oscillator whose position x obeys
191     # m x'' + gamma x' + kx = F_thermal(t)
192     #  A is the ?driving noise?
193     #  B is gamma, the damping term
194     #  C is omega_0, the resonant frequency
195     # Fitting the PSD of V = photoSensitivity*x just rescales the parameters
196     g("f(x) = C / ((A**2-x**2)**2 + (B*x)**2)")
197     ## The mass m is given by m = k_B T / (pi A)
198     ## The spring constant k is given by k = m (omega_0)**2
199     ## The quality factor Q is given by Q = omega_0 m / gamma
200     g("A=%g; B=%g; C=%g" % (A_guess, B_guess, C_guess))
201     g("fit f(x) '%s' via A,B,C" % datafilename)
202     A=abs(float(g.getvar('A'))) # A and B only show up as squares in f(x)
203     B=abs(float(g.getvar('B'))) # so ensure we get positive values
204     C=float(g.getvar('C'))
205     os.remove(datafilename)
206     return (A, B, C)
207
208 def lorentzian_area(A, B, C) :
209     # Integrating the the power spectral density per unit time (power) over the
210     # frequency axis [0, fN] returns the total signal power per unit time
211     #  int_0^fN power(f)df = 1/T int_0^T |x(t)**2| dt
212     #                      = <V(t)**2>, the variance for our AC signal.
213     # The variance from our fitted Lorentzian is the area under the Lorentzian
214     #  <V(t)**2> = (pi*C) / (2*B*A**2)
215     return (numpy.pi * C) / (2 * B * A**2)
216
217 def gnuplot_check_fit(datafilename, A, B, C) :
218     """
219     return a string containing a gnuplot script displaying the fit.
220     """
221     string = 'f(x) = C / ((A**2-x**2)**2 + (x*B)**2)\n'
222     string += 'A= %g ; B= %g ; C= %g\n' % (A, B, C)
223     string += 'set logscale y\n'
224     string += "plot '%s', f(x)\n" % datafilename
225     return string
226
227 def vib_save(data, log_dir=None) :
228     """Save the dictionary data, using data_logger.data_log()
229     data should be a dictionary of arrays with the fields
230       'Z piezo output'
231       'Z piezo input'
232       'Deflection input'
233     """
234     if log_dir :
235         log = data_logger.data_log(log_dir, noclobber_logsubdir=False,
236                                    log_name="vib_%gHz" % freq)
237         log.write_dict_of_arrays(data)
238
239 def vib_load(datafile=None) :
240     """Load the dictionary data, using data_logger.date_load()"
241     data should be a dictionary of arrays with the fields
242       'Z piezo output'
243       'Z piezo input'
244       'Deflection input'
245     """
246     dl = data_logger.data_load()
247     data = dl.read_dict_of_arrays(datafile)
248     return data
249
250 def vib_plot(deflection_bits, freq_axis, power, A, B, C,
251              plotVerbose=True) :
252     """
253     If plotVerbose or config.PYLAB_VERBOSE == True, plot 3 subfigures:
254      Time series (Vphoto vs sample index) (show any major drift events),
255      A histogram of Vphoto, with a gaussian fit (demonstrate randomness), and
256      FFTed Vphoto data (Vphoto vs Frequency) (show major noise components).
257     """
258     if plotVerbose or config.PYLAB_VERBOSE :
259         common._import_pylab()
260         common._pylab.figure(config.BASE_FIGNUM+2)
261
262         if deflection_bits != None:
263             # plot time series
264             common._pylab.subplot(311)
265             common._pylab.hold(False)
266             common._pylab.plot(deflection_bits, 'r.')
267             common._pylab.title("free oscillation")
268     
269             # plot histogram distribution and gaussian fit
270             common._pylab.subplot(312)
271             common._pylab.hold(False)
272             n, bins, patches = \
273                 common._pylab.hist(deflection_bits, bins=30,
274                                    normed=1, align='center')
275             gaus = numpy.zeros((len(bins),), dtype=numpy.float)
276             mean = deflection_bits.mean()
277             std = deflection_bits.std()
278             pi = numpy.pi
279             exp = numpy.exp
280             for i in range(len(bins)) :
281                 gaus[i] = (2*pi)**(-.5)/std * exp(-0.5*((bins[i]-mean)/std)**2)
282             common._pylab.hold(True)
283             common._pylab.plot(bins, gaus, 'r-');
284             common._pylab.hold(False)
285     
286             # plot FFTed data
287             axes = common._pylab.subplot(313)
288         else:
289             # use a nice big subplot just for the FFTed data
290             axes = common._pylab.subplot(111)
291         common._pylab.hold(False)
292         common._pylab.semilogy(freq_axis, power, 'r.-')
293         fitdata = C / ((A**2-freq_axis**2)**2 + (B*freq_axis)**2)
294         xmin,xmax = axes.get_xbound()
295         ymin,ymax = axes.get_ybound()
296         common._pylab.hold(True)
297         common._pylab.plot(freq_axis, fitdata, 'b-');
298         common._pylab.hold(False)
299         axes.axis([xmin,xmax,ymin,ymax])
300
301         common._flush_plot()
302     if (plotVerbose or config.GNUPLOT_VERBOSE) and False: # TODO: cleanup and test
303         # write all the ft data now
304         fd = file(datafilename, 'w')
305         for i in range(len(freq_axis)) :
306             fd.write("%g\t%g\n" % (freq_axis[i], power[i]))
307         fd.write("\n") # blank line to separate fit data.
308         # write the fit data
309         for i in range(imin, imax) :
310             x = freq_axis[i]
311             fd.write("%g\t%g\n" % (freq_axis[i],
312                                    C / ((A**2-x**2)**2 + (x*B)**2) ) )
313         fd.close()
314         print Vphoto_t.var(), A, B, C, ((numpy.pi * C)/(2 * B * A**2) * 1e-12)
315         g("set terminal epslatex color solid")
316         g("set output 'calibration.tex'")
317         g("set size 2,2") # multipliers on default 5"x3.5"
318         #g("set title 'Thermal calibration'")
319         g("set logscale y")
320         g("set xlabel 'Frequency (Hz)'")
321         g("set ylabel 'Power ($\mu$V$^2$/Hz)'")
322         g("set xrange [0:20000]")
323         g("g(x) = x < %g || x > %g ? 1/0 : f(x)" % (minFreq, maxFreq)) # only plot on fitted region
324         g("plot '%s' title 'Vibration data', g(x) title 'Lorentzian fit'" % datafilename)
325
326         g("set mouse")
327         g("pause mouse 'click with mouse'")
328         g.getvar("MOUSE_BUTTON")
329         del(g)
330
331 def vib_load_analyze_tweaked(tweak_file, minFreq=1000, maxFreq=7000,
332                              chunk_size=2048, overlap=True,
333                              window=FFT_tools.window_hann,
334                              Vphoto_in2V=config.Vphoto_in2V,
335                              plotVerbose=False) :
336     """
337     Read the vib files listed in tweak_file.
338     Return an array of Vphoto variances (due to the cantilever) in Volts**2
339     """
340     dl = data_logger.data_load()
341     Vphoto_var = []
342     for path in file(tweak_file, 'r') :
343         if path[-1] == '\n':
344             path = path.split('\n')[0]
345         # read the data
346         data = vib_load(path)
347         deflection_bits = data['Deflection input']
348         freq = float(path.split('_')[-1].split('Hz')[0])
349         if config.TEXT_VERBOSE :
350             print "Analyzing '%s' at %g Hz" % (path, freq)
351         # analyze
352         Vphoto_var.append(vib_analyze(deflection_bits, freq, minFreq, maxFreq,
353                                       chunk_size, overlap, window,
354                                       Vphoto_in2V, plotVerbose))
355     return numpy.array(Vphoto_var, dtype=numpy.float)
356
357 # commandline interface functions
358 import scipy.io, sys
359
360 def read_data(ifile):
361     """
362     ifile can be a filename string or open (seekable) file object.
363     returns an array of data values (1 column)
364     """
365     #returns (column 1 array, column 2 array)
366     #"""
367     if ifile == None :  ifile = sys.stdin
368     unlabeled_data=scipy.io.read_array(ifile)
369     return unlabeled_data #(unlabeled_data[:,0], unlabeled_data[:,1])
370
371 if __name__ == '__main__' :
372     # command line interface
373     from optparse import OptionParser
374     
375     usage_string = ('%prog <input-file>\n'
376                     '2008, W. Trevor King.\n'
377                     '\n'
378                     'Deflection power spectral density (PSD) data (X^2/Hz)\n'
379                     'and returns the variance in X units (X^2)'
380                     '<input-file> should be whitespace-delimited, 2 column ASCII\n'
381                     'without a header line.\n'
382                     'e.g: "<frequency_Hz>\\t<deflection_psd X^2/Hz>\\n"\n')
383     parser = OptionParser(usage=usage_string, version='%prog 0.1')
384     parser.add_option('-f', '--sample-freq', dest='freq', default=100e3,
385                       help='sample frequency in Hz (default %default)',
386                       type='float', metavar='FREQ')
387     parser.add_option('-m', '--min-freq', dest='min_freq', default=500,
388                       help='minimum frequency in Hz (default %default)',
389                       type='float', metavar='FREQ')
390     parser.add_option('-M', '--max-freq', dest='max_freq', default=7000,
391                       help='maximum frequency in Hz (default %default)',
392                       type='float', metavar='FREQ')
393     parser.add_option('-o', '--output-file', dest='ofilename',
394                       help='write output to FILE (default stdout)',
395                       type='string', metavar='FILE')
396     parser.add_option('-c', '--comma-out', dest='comma_out', action='store_true',
397                       help='Output comma-seperated values (default %default)',
398                       default=False)
399     parser.add_option('-g', '--gnuplot', dest='gnuplot', action='store_true',
400                       help='Print gnuplot fit check script to stderr',
401                       default=False)
402     parser.add_option('-p', '--pylab', dest='pylab', action='store_true',
403                       help='Produce pylab fit checks during execution',
404                       default=False)
405     parser.add_option('-G', '--plot-guess', dest='plot_guess', action='store_true',
406                       help='Produce plots of the pre-fit Lorentzian',
407                       default=False)
408     parser.add_option('-t', '--tweak-mode', dest='tweakmode', action='store_true',
409                       help='Run in tweak-file mode',
410                       default=False)
411     parser.add_option('-v', '--verbose', dest='verbose', action='store_true',
412                       help='Print lots of debugging information',
413                       default=False)
414
415     options,args = parser.parse_args()
416     parser.destroy()
417     assert len(args) >= 1, "Need an input file"
418         
419     ifilename = args[0]
420
421     if options.ofilename != None :
422         ofile = file(options.ofilename, 'w')
423     else :
424         ofile = sys.stdout
425     if options.verbose == True :
426         vfile = sys.stderr
427     else :
428         vfile = None
429     config.TEXT_VERBOSE = options.verbose
430     config.PYLAB_INTERACTIVE = False
431     config.PYLAB_VERBOSE = options.pylab
432     config.GNUPLOT_VERBOSE = options.gnuplot
433     PLOT_GUESSED_LORENTZIAN = options.plot_guess
434
435     if options.tweakmode == False :
436         data = read_data(ifilename)
437         deflection_bits = data['Deflection input']
438         Vphoto_var = vib_analyze(deflection_bits, freq=options.freq,
439                                  minFreq=options.min_freq,
440                                  maxFreq=options.max_freq)
441         print >> ofile, Vphoto_var
442     else : # tweak mode
443         Vphoto_vars = vib_load_analyze_tweaked(ifilename,
444                                                minFreq=options.min_freq,
445                                                maxFreq=options.max_freq)
446         if options.comma_out :
447             sep = ','
448         else :
449             sep = '\n'
450         common.write_array(ofile, Vphoto_vars, sep)
451
452     if common._final_flush_plot != None:
453         common._final_flush_plot()
454     
455     if options.ofilename != None :
456         ofile.close()
457