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