8c6791e20aae3e9e71fc9020bf73f3198a516ea5
[calibcant.git] / calibcant / vib_analyze.py
1 # calibcant - tools for thermally calibrating AFM cantilevers
2 #
3 # Copyright (C) 2008-2011 W. Trevor King <wking@drexel.edu>
4 #
5 # This file is part of calibcant.
6 #
7 # calibcant is free software: you can redistribute it and/or
8 # modify it under the terms of the GNU Lesser General Public
9 # License as published by the Free Software Foundation, either
10 # version 3 of the License, or (at your option) any later version.
11 #
12 # calibcant is distributed in the hope that it will be useful,
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 # GNU Lesser General Public License for more details.
16 #
17 # You should have received a copy of the GNU Lesser General Public
18 # License along with calibcant.  If not, see
19 # <http://www.gnu.org/licenses/>.
20
21 """Thermal vibration analysis.
22
23 Separate the more general `vib_analyze()` from the other `vib_*()`
24 functions in calibcant.
25
26 The relevent physical quantities are :
27   Vphoto   The photodiode vertical deflection voltage (what we measure)
28
29 >>> import os
30 >>> from pprint import pprint
31 >>> import random
32 >>> import tempfile
33 >>> import numpy
34 >>> from .config import VibrationConfig
35 >>> from h5config.storage.hdf5 import pprint_HDF5
36 >>> from pypiezo.test import get_piezo_config
37 >>> from pypiezo.base import convert_volts_to_bits
38
39 >>> fd,filename = tempfile.mkstemp(suffix='.h5', prefix='calibcant-')
40 >>> os.close(fd)
41
42 >>> piezo_config = get_piezo_config()
43 >>> vibration_config = VibrationConfig()
44 >>> vibration_config['frequency'] = 50e3
45
46 We'll be generating a test vibration time series with the following
47 parameters.  Make sure these are all floats to avoid accidental
48 integer division in later steps.
49
50 >>> m = 5e-11       # kg
51 >>> gamma = 1.6e-6  # N*s/m
52 >>> k = 0.05        # N/m
53 >>> T = 1/vibration_config['frequency']
54 >>> T  # doctest: +ELLIPSIS
55 2...e-05
56 >>> N = int(2**15)  # count
57 >>> F_sigma = 1e3   # N
58
59 where `T` is the sampling period, `N` is the number of samples, and
60 `F_sigma` is the standard deviation of the white-noise external force.
61 Note that the resonant frequency is less than the Nyquist frequency so
62 we don't have to worry too much about aliasing.
63
64 >>> w0 = numpy.sqrt(k/m)
65 >>> f0 = w0/(2*numpy.pi)
66 >>> f0  # doctest: +ELLIPSIS
67 5032.9...
68 >>> f_nyquist = vibration_config['frequency']/2
69 >>> f_nyquist  # doctest: +ELLIPSIS
70 25000.0...
71
72 The damping ratio is
73
74 >>> damping = gamma / (2*m*w0)
75 >>> damping  # doctest: +ELLIPSIS
76 0.505...
77
78 The quality factor is 
79
80 >>> Q = m*w0 / gamma
81 >>> Q  # doctest: +ELLIPSIS
82 0.988...
83 >>> (1 / (2*damping)) / Q  # doctest: +ELLIPSIS
84 1.000...
85
86 We expect the white-noise power spectral density (PSD) to be a flat
87 line at
88
89 >>> F0 = F_sigma**2 * 2 * T
90
91 because the integral from `0` `1/2T` should be `F_sigma**2`.
92
93 The expected time series PSD parameters are
94
95 >>> A = f0
96 >>> B = gamma/(m*2*numpy.pi)
97 >>> C = F0/(m**2*(2*numpy.pi)**4)
98
99 Simulate a time series with the proper PSD using center-differencing.
100
101   m\ddot{x} + \gamma \dot{x} + kx = F
102
103   m \frac{x_{i+1} - 2x_i + x_{i-1}}{T**2}
104     + \gamma \frac{x_{i+1}-x_{i-1}}{T}
105     + kx_i = F_i
106
107   a x_{i+1} + b x_{i} + c x_{i-1} = F_i
108
109 where `T` is the sampling period, `i=t/T` is the measurement index,
110 `a=m/T**2+gamma/2T`, `b=-2m/T**2+k`, and `c=m/T**2-gamma/2T`.
111 Rearranging and shifting to `j=i-1`
112
113   x_j = \frac{F_{i-1} - bx_{i-1} - cx_{i-2}}{a}
114
115 >>> a = m/T**2 + gamma/(2*T)
116 >>> b = -2*m/T**2 + k
117 >>> c = m/T**2 - gamma/(2*T)
118 >>> x = numpy.zeros((N+2,), dtype=numpy.float)  # two extra initial points
119 >>> F = numpy.zeros((N,), dtype=numpy.float)
120 >>> for i in range(2, x.size):
121 ...     Fp = random.gauss(mu=0, sigma=F_sigma)
122 ...     F[i-2] = Fp
123 ...     xp = x[i-1]
124 ...     xpp = x[i-2]
125 ...     x[i] = (Fp - b*xp - c*xpp)/a
126 >>> x = x[2:]  # drop the initial points
127
128 Convert the simulated data to bits.
129
130 >>> deflection = x
131 >>> deflection_bits = convert_volts_to_bits(
132 ...     piezo_config.select_config('inputs', 'deflection'), x)
133
134 Analyze the simulated data.
135
136 >>> naive_vibration = vib_analyze_naive(deflection)
137 >>> naive_vibration  # doctest: +SKIP
138 136871517.43486854
139 >>> abs(naive_vibration / 136.9e6 - 1) < 0.1
140 True
141
142 >>> processed_vibration = vib_analyze(
143 ...     deflection_bits, vibration_config,
144 ...     piezo_config.select_config('inputs', 'deflection'))
145 >>> processed_vibration  # doctest: +SKIP
146 136457906.25574699
147
148 >>> vib_plot(deflection=deflection_bits, vibration_config=vibration_config)
149 >>> vib_save(filename=filename, group='/vibration/',
150 ...     raw_vibration=deflection_bits, vibration_config=vibration_config,
151 ...     deflection_channel_config=piezo_config.select_config(
152 ...         'inputs', 'deflection'),
153 ...     processed_vibration=processed_vibration)
154
155 >>> pprint_HDF5(filename)  # doctest: +ELLIPSIS, +REPORT_UDIFF
156 /
157   /vibration
158     /vibration/config
159       /vibration/config/deflection
160         <HDF5 dataset "channel": shape (), type "<i4">
161           0
162         <HDF5 dataset "conversion-coefficients": shape (2,), type "<i4">
163           [0 1]
164         <HDF5 dataset "conversion-origin": shape (), type "<i4">
165           0
166         <HDF5 dataset "device": shape (), type "|S12">
167           /dev/comedi0
168         <HDF5 dataset "inverse-conversion-coefficients": shape (2,), type "<i4">
169           [0 1]
170         <HDF5 dataset "inverse-conversion-origin": shape (), type "<i4">
171           0
172         <HDF5 dataset "maxdata": shape (), type "<i4">
173           100
174         <HDF5 dataset "name": shape (), type "|S10">
175           deflection
176         <HDF5 dataset "range": shape (), type "<i4">
177           1
178         <HDF5 dataset "subdevice": shape (), type "<i4">
179           -1
180       /vibration/config/vibration
181         <HDF5 dataset "chunk-size": shape (), type "<i4">
182           2048
183         <HDF5 dataset "frequency": shape (), type "<f8">
184           50000.0
185         <HDF5 dataset "maximum-fit-frequency": shape (), type "<f8">
186           25000.0
187         <HDF5 dataset "minimum-fit-frequency": shape (), type "<f8">
188           500.0
189         <HDF5 dataset "model": shape (), type "|S12">
190           Breit-Wigner
191         <HDF5 dataset "overlap": shape (), type "|b1">
192           False
193         <HDF5 dataset "sample-time": shape (), type "<i4">
194           1
195         <HDF5 dataset "window": shape (), type "|S4">
196           Hann
197     <HDF5 dataset "processed": shape (), type "<f8">
198       ...
199     /vibration/raw
200       <HDF5 dataset "deflection": shape (32768,), type "<f8">
201         [...]
202
203 >>> (raw_vibration,vibration_config,deflection_channel_config,
204 ...  processed_vibration) = vib_load(
205 ...     filename=filename, group='/vibration/')
206
207 >>> processed_vibration  # doctest: +SKIP
208 136457906.25574699
209 >>> abs(processed_vibration / 136.5e6 - 1) < 0.1
210 True
211
212 >>> os.remove(filename)
213 """
214
215 import os as _os
216 import time as _time
217
218 import h5py as _h5py
219 import numpy as _numpy
220 from scipy.optimize import leastsq as _leastsq
221
222 try:
223     import matplotlib as _matplotlib
224     import matplotlib.pyplot as _matplotlib_pyplot
225     import time as _time  # for timestamping lines on plots
226 except (ImportError, RuntimeError), e:
227     _matplotlib = None
228     _matplotlib_import_error = e
229
230 from h5config.storage.hdf5 import HDF5_Storage as _HDF5_Storage
231 from h5config.storage.hdf5 import h5_create_group as _h5_create_group
232 import FFT_tools as _FFT_tools
233 from pypiezo.base import convert_bits_to_volts as _convert_bits_to_volts
234 from pypiezo.config import ChannelConfig as _ChannelConfig
235
236 from . import LOG as _LOG
237 from . import package_config as _package_config
238 from .config import Variance as _Variance
239 from .config import BreitWigner as _BreitWigner
240 from .config import OffsetBreitWigner as _OffsetBreitWigner
241 from .config import VibrationConfig as _VibrationConfig
242
243
244 def vib_analyze_naive(deflection):
245     """Calculate the deflection variance in Volts**2.
246
247     This method is simple and easy to understand, but it highly
248     succeptible to noise, drift, etc.
249     
250     Inputs
251       deflection : numpy array with deflection timeseries in Volts.
252     """
253     std = deflection.std()
254     var = std**2
255     _LOG.debug('naive deflection variance: %g V**2' % var)
256     return var
257
258 def vib_analyze(deflection, vibration_config, deflection_channel_config,
259                 plot=False):
260     """Calculate the deflection variance in Volts**2.
261
262     Improves on `vib_analyze_naive()` by first converting `Vphoto(t)`
263     to frequency space, and fitting a Breit-Wigner in the relevant
264     frequency range (see cantilever_calib.pdf for derivation).
265     However, there may be cases where the fit is thrown off by noise
266     spikes in frequency space.  To protect from errors, the fitted
267     variance is compared to the naive variance (where all noise is
268     included), and the minimum variance is returned.
269
270     Inputs:
271       deflection        Vphoto deflection input in bits.
272       vibration_config  `.config._VibrationConfig` instance
273       deflection_channel_config
274                         deflection `pypiezo.config.ChannelConfig` instance
275       plot              boolean overriding matplotlib config setting.
276
277     The conversion to frequency space generates an average power
278     spectrum by breaking the data into windowed chunks and averaging
279     the power spectrums for the chunks together. See
280     `FFT_tools.unitary_avg_power_spectrum()` for details.
281     """
282     # convert the data from bits to volts
283     deflection_v = _convert_bits_to_volts(
284         deflection_channel_config, deflection)
285     mean = deflection_v.mean()
286     _LOG.debug('average thermal deflection (Volts): %g' % mean)
287
288     naive_variance = vib_analyze_naive(deflection_v)
289     if vibration_config['model'] == _Variance:
290         return naive_variance
291     
292     # Compute the average power spectral density per unit time (in uV**2/Hz)
293     _LOG.debug('compute the averaged power spectral density in uV**2/Hz')
294     freq_axis,power = _FFT_tools.unitary_avg_power_spectrum(
295         (deflection_v - mean)*1e6, vibration_config['frequency'],
296         vibration_config['chunk-size'], vibration_config['overlap'],
297         vibration_config['window'])
298
299     A,B,C,D = fit_psd(
300         freq_axis, power,
301         min_frequency=vibration_config['minimum-fit-frequency'],
302         max_frequency=vibration_config['maximum-fit-frequency'],
303         offset=vibration_config['model'] == _OffsetBreitWigner)
304
305     _LOG.debug('fit PSD(f) = C / ((A**2-f**2)**2 + (f*B)**2) with '
306                'A = %g, B = %g, C = %g, D = %g' % (A, B, C, D))
307
308     if plot or _package_config['matplotlib']:
309         vib_plot(deflection, freq_axis, power, A, B, C, D,
310                  vibration_config=vibration_config)
311
312     # Our A is in uV**2, so convert back to Volts**2
313     fitted_variance = breit_wigner_area(A,B,C) * 1e-12
314
315     _LOG.debug('fitted deflection variance: %g V**2' % fitted_variance)
316
317     if _package_config['matplotlib']:
318         vib_plot(deflection, freq_axis, power, A, B, C, D,
319                  vibration_config=vibration_config)
320
321     return min(fitted_variance, naive_variance)
322
323 def breit_wigner(f, A, B, C, D=0):
324     """Breit-Wigner (sortof).
325
326     Inputs
327       f  Frequency
328       A  Resonant frequency
329       B  Quality Q = A/B
330       C  Scaling factor
331       D  Optional white-noise offset
332
333     All parameters must be postive.
334     """
335     return abs(C) / ((A**2-f**2)**2 + (B*f)**2) + abs(D)
336
337 def fit_psd(freq_axis, psd_data, min_frequency=500, max_frequency=25000,
338             offset=False):
339     """Fit the FFTed vibration data to a Breit-Wigner.
340
341     Inputs
342       freq_axis      array of frequencies in Hz
343       psd_data       array of PSD amplitudes for the frequencies in freq_axis
344       min_frequency  lower bound of Breit-Wigner fitting region
345       max_frequency  upper bound of Breit-Wigner fitting region
346       offset         add a white-noise offset to the Breit-Wigner model
347     Output
348       Breit-Wigner model fit parameters `A`, `B`, `C`, and `D`.
349     """
350     # cut out the relevent frequency range
351     _LOG.debug('cut the frequency range %g to %g Hz'
352                % (min_frequency, max_frequency))
353     imin = 0
354     while freq_axis[imin] < min_frequency : imin += 1
355     imax = imin
356     while freq_axis[imax] < max_frequency : imax += 1
357     assert imax >= imin + 10 , 'less than 10 points in freq range (%g,%g)' % (
358         min_frequency, max_frequency)
359
360     # generate guesses for Breit-Wigner parameters A, B, C, and D
361     max_psd_index = _numpy.argmax(psd_data[imin:imax]) + imin
362     max_psd = psd_data[max_psd_index]
363     res_freq = freq_axis[max_psd_index]
364
365     # Breit-Wigner L(x) = C / ((A**2-x**2)**2 + (B*x)**2)
366     # is expected power spectrum for
367     # x'' + B x' + A^2 x'' = F_external(t)/m
368     # (A = omega_0)
369     # C = (2 k_B T B) / (pi m)
370     # 
371     # A = resonant frequency
372     # peak at  x_res = sqrt(A^2 - B^2/2)  (by differentiating)
373     #  which could be complex if there isn't a peak (overdamped)
374     # peak height    = C / (3 x_res^4 + A^4)
375     #
376     # Q = A/B
377     #
378     # Guessing Q = 1 is pretty safe.
379
380     Q_guess = 1
381
382     # so x_res^2 = B^2 Q^2 - B^2/2 = (Q^2-1/2)B^2 
383     #    B = x_res / sqrt(Q^2-1/2)
384     B_guess = res_freq / _numpy.sqrt(Q_guess**2-0.5)
385     A_guess = Q_guess*B_guess
386     C_guess = max_psd * (-res_freq**4 + A_guess**4)
387     if offset:
388         D_guess = psd_data[max_psd_index]
389         C_guess -= D_guess
390     else:
391         D_guess = 0
392     _LOG.debug(('guessed params: resonant freq %g, max psd %g, Q %g, A %g, '
393                 'B %g, C %g, D %g') % (
394             res_freq, max_psd, Q_guess, A_guess, B_guess, C_guess, D_guess))
395     # Half width w on lower side when L(a-w) = L(a)/2
396     #  (a**2 - (a-w)**2)**2 + (b*(a-w))**2 = 2*(b*a)**2
397     # Let W=(a-w)**2, A=a**2, and B=b**2
398     #  (A - W)**2 + BW = 2*AB
399     #  W**2 - 2AW + A**2 + BW = 2AB
400     #  W**2 + (B-2A)W + (A**2-2AB) = 0
401     #  W = (2A-B)/2 * [1 +/- sqrt(1 - 4(A**2-2AB)/(B-2A)**2]
402     #    = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
403     #  (a-w)**2 = (2A-B)/2 * [1 +/- sqrt(1 - 4A(A-2B)/(B-2A)**2]
404     #  so w is a disaster ;)
405     # For some values of A and B (non-underdamped), W is imaginary or negative.
406     #
407     # The mass m is given by m = k_B T / (pi A)
408     # The spring constant k is given by k = m (omega_0)**2
409     # The quality factor Q is given by Q = omega_0 m / gamma
410
411     # Fitting the PSD of V = photoSensitivity*x just rescales the parameters
412
413     # fit Breit-Wigner using scipy.optimize.leastsq
414     def residual(p, y, x):
415         return breit_wigner(x, *p) - y
416     if offset:
417         guess = _numpy.array((A_guess, B_guess, C_guess, D_guess))
418     else:
419         guess = _numpy.array((A_guess, B_guess, C_guess))
420
421     p,cov,info,mesg,ier = _leastsq(
422         residual, guess,
423         args=(psd_data[imin:imax], freq_axis[imin:imax]),
424         full_output=True, maxfev=10000)
425     _LOG.debug('fitted params: %s' % p)
426     _LOG.debug('covariance matrix: %s' % cov)
427     #_LOG.debug('info: %s' % info)
428     _LOG.debug('message: %s' % mesg)
429     if ier == 1:
430         _LOG.debug('solution converged')
431     else:
432         _LOG.debug('solution did not converge')
433     if offset:
434         A,B,C,D = p
435     else:
436         A,B,C = p
437         D = 0
438     A=abs(A) # A and B only show up as squares in f(x)
439     B=abs(B) # so ensure we get positive values.
440     C=abs(C) # Only abs(C) is used in breit_wigner().
441     return (A, B, C, D)
442
443 def breit_wigner_area(A, B, C):
444     # Integrating the the power spectral density per unit time (power) over the
445     # frequency axis [0, fN] returns the total signal power per unit time
446     #  int_0^fN power(f)df = 1/T int_0^T |x(t)**2| dt
447     #                      = <V(t)**2>, the variance for our AC signal.
448     # The variance from our fitted Breit-Wigner is the area under the Breit-Wigner
449     #  <V(t)**2> = (pi*C) / (2*B*A**2)
450     return (_numpy.pi * C) / (2 * B * A**2)
451
452 def vib_save(filename, group='/', raw_vibration=None, vibration_config=None,
453              deflection_channel_config=None, processed_vibration=None):
454     with _h5py.File(filename, 'a') as f:
455         cwg = _h5_create_group(f, group)
456         if raw_vibration is not None:
457             try:
458                 del cwg['raw/deflection']
459             except KeyError:
460                 pass
461             cwg['raw/deflection'] = raw_vibration
462         storage = _HDF5_Storage()
463         for config,key in [(vibration_config, 'config/vibration'),
464                            (deflection_channel_config,
465                             'config/deflection')]:
466             if config is None:
467                 continue
468             config_cwg = _h5_create_group(cwg, key)
469             storage.save(config=config, group=config_cwg)
470         if processed_vibration is not None:
471             try:
472                 del cwg['processed']
473             except KeyError:
474                 pass
475             cwg['processed'] = processed_vibration
476
477 def vib_load(filename, group='/'):
478     assert group.endswith('/')
479     raw_vibration = processed_vibration = None
480     configs = []
481     with _h5py.File(filename, 'a') as f:
482         try:
483             raw_vibration = f[group+'raw/deflection'][...]
484         except KeyError:
485             pass
486         for Config,key in [(_VibrationConfig, 'config/vibration'),
487                            (_ChannelConfig, 'config/deflection')]:
488             config = Config(storage=_HDF5_Storage(
489                     filename=filename, group=group+key))
490             configs.append(config)
491         try:
492             processed_vibration = float(f[group+'processed'][...])
493         except KeyError:
494             pass
495     ret = [raw_vibration]
496     ret.extend(configs)
497     ret.append(processed_vibration)
498     for config in configs:
499         config.load()
500     return tuple(ret)
501
502 def vib_plot(deflection=None, freq_axis=None, power=None, A=None, B=None,
503              C=None, D=0, vibration_config=None, analyze=False):
504     """Plot 3 subfigures displaying vibration data and analysis.
505
506      Time series (Vphoto vs sample index) (show any major drift events),
507      A histogram of Vphoto, with a gaussian fit (demonstrate randomness), and
508      FFTed Vphoto data (Vphoto vs Frequency) (show major noise components).
509     """
510     if not _matplotlib:
511         raise _matplotlib_import_error
512     figure = _matplotlib_pyplot.figure()
513
514     if power is None:
515         assert deflection != None, (
516             'must set at least one of `deflection` or `power`.')
517         time_axes = figure.add_subplot(2, 1, 1)
518         hist_axes = figure.add_subplot(2, 1, 2)
519         freq_axes = None
520     elif deflection is None:
521         time_axes = hist_axes = None
522         freq_axes = figure.add_subplot(1, 1, 1)
523     else:
524         time_axes = figure.add_subplot(3, 1, 1)
525         hist_axes = figure.add_subplot(3, 1, 2)
526         freq_axes = figure.add_subplot(3, 1, 3)
527         
528     timestamp = _time.strftime('%H%M%S')
529
530     if deflection is not None:
531         time_axes.plot(deflection, 'r.')
532         time_axes.set_title('free oscillation')
533
534         # plot histogram distribution and gaussian fit
535         hist_axes.hold(True)
536         n,bins,patches = hist_axes.hist(
537             deflection, bins=30, normed=True, align='mid')
538         gauss = _numpy.zeros((len(bins),), dtype=_numpy.float)
539         mean = deflection.mean()
540         std = deflection.std()
541         pi = _numpy.pi
542         exp = _numpy.exp
543         gauss = _numpy.sqrt(2*pi)/std * exp(-0.5*((bins-mean)/std)**2)
544         # Matplotlib's normed histogram uses bin heights of n/(len(x)*dbin)
545         dbin = bins[1]-bins[0]
546         hist_axes.plot(bins, gauss/dbin, 'r-')
547     if power is not None:
548         freq_axes.hold(True)
549         freq_axes.set_yscale('log')
550         freq_axes.plot(freq_axis, power, 'r.-')
551         xmin,xmax = freq_axes.get_xbound()
552         ymin,ymax = freq_axes.get_ybound()
553     
554         # highlight the region we're fitting
555         freq_axes.axvspan(
556             vibration_config['minimum-fit-frequency'],
557             vibration_config['maximum-fit-frequency'],
558             facecolor='g', alpha=0.1, zorder=-2)
559
560         if A is not None:
561             fitdata = breit_wigner(freq_axis, A, B, C, D)
562             freq_axes.plot(freq_axis, fitdata, 'b-')
563             noisefloor = D + 0*freq_axis;
564             freq_axes.plot(freq_axis, noisefloor)
565
566             if B**2 < 2*A**2:
567                 res_freq = _numpy.sqrt(A**2 - B**2/2)
568                 freq_axes.axvline(res_freq, color='b', zorder=-1)
569
570         freq_axes.set_title('power spectral density %s' % timestamp)
571         freq_axes.axis([xmin,xmax,ymin,ymax])
572         freq_axes.set_xlabel('frequency (Hz)')
573         freq_axes.set_ylabel('PSD')
574     if hasattr(figure, 'show'):
575         figure.show()