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