analyze.py: Set `changed = True` for tweaked vibration variance
[calibcant.git] / calibcant / vibration_analyze.py
index 1844dcf4e78be771f0e1645ed71078de68839a77..50bf87d3b36fce7818c14a3a58e93e19fe8ec8ae 100644 (file)
@@ -1,6 +1,6 @@
 # calibcant - tools for thermally calibrating AFM cantilevers
 #
-# Copyright (C) 2008-2012 W. Trevor King <wking@drexel.edu>
+# Copyright (C) 2008-2013 W. Trevor King <wking@tremily.us>
 #
 # This file is part of calibcant.
 #
@@ -306,10 +306,17 @@ def analyze(deflection, config, deflection_channel_config,
         config['chunk-size'], config['overlap'],
         config['window'])
 
+    if config['smooth-window'] and config['smooth-length']:
+        n = config['smooth-length']
+        smooth_window = config['smooth-window'](length=n)
+        smooth_window /= smooth_window.sum()
+    else:
+        smooth_window = None
     A,B,C,D = fit_psd(
         freq_axis, power,
         min_frequency=config['minimum-fit-frequency'],
         max_frequency=config['maximum-fit-frequency'],
+        smooth_window=smooth_window,
         offset=config['model'] == _OffsetBreitWigner)
 
     _LOG.debug('fit PSD(f) = C / ((A**2-f**2)**2 + (f*B)**2) with '
@@ -340,7 +347,7 @@ def breit_wigner(f, A, B, C, D=0):
     return abs(C) / ((A**2-f**2)**2 + (B*f)**2) + abs(D)
 
 def fit_psd(freq_axis, psd_data, min_frequency=500, max_frequency=25000,
-            offset=False):
+            smooth_window=_FFT_tools.window_hann(10), offset=False):
     """Fit the FFTed vibration data to a Breit-Wigner.
 
     Inputs
@@ -348,6 +355,7 @@ def fit_psd(freq_axis, psd_data, min_frequency=500, max_frequency=25000,
       psd_data       array of PSD amplitudes for the frequencies in freq_axis
       min_frequency  lower bound of Breit-Wigner fitting region
       max_frequency  upper bound of Breit-Wigner fitting region
+      smooth_window  array for smoothing the PSD before guessing params
       offset         add a white-noise offset to the Breit-Wigner model
     Output
       Breit-Wigner model fit parameters `A`, `B`, `C`, and `D`.
@@ -363,7 +371,11 @@ def fit_psd(freq_axis, psd_data, min_frequency=500, max_frequency=25000,
         min_frequency, max_frequency)
 
     # generate guesses for Breit-Wigner parameters A, B, C, and D
-    max_psd_index = _numpy.argmax(psd_data[imin:imax]) + imin
+    if hasattr(smooth_window, '__len__'):
+        smooth_psd_data = _numpy.convolve(psd_data, smooth_window, mode='same')
+    else:
+        smooth_psd_data = psd_data
+    max_psd_index = _numpy.argmax(smooth_psd_data[imin:imax]) + imin
     max_psd = psd_data[max_psd_index]
     res_freq = freq_axis[max_psd_index]
 
@@ -390,7 +402,7 @@ def fit_psd(freq_axis, psd_data, min_frequency=500, max_frequency=25000,
     A_guess = Q_guess*B_guess
     C_guess = max_psd * (-res_freq**4 + A_guess**4)
     if offset:
-        D_guess = psd_data[max_psd_index]
+        D_guess = psd_data[-1]
         C_guess -= D_guess
     else:
         D_guess = 0
@@ -454,7 +466,12 @@ def breit_wigner_area(A, B, C):
     #  <V(t)**2> = (pi*C) / (2*B*A**2)
     return (_numpy.pi * C) / (2 * B * A**2)
 
-def save(filename, group='/', raw=None, config=None,
+def breit_wigner_resonant_frequency(A, B):
+    if (B**2 >= 2*A**2):
+        return 0  # over- or critically-damped
+    return _numpy.sqrt(A**2 - B**2/2)
+
+def save(filename=None, group='/', raw=None, config=None,
          deflection_channel_config=None, processed=None):
     specs = [
         _SaveSpec(item=config, relpath='config/vibration',
@@ -531,21 +548,21 @@ def plot(deflection=None, freq_axis=None, power=None, A=None, B=None,
         freq_axes.autoscale(tight=True)
         xmin,xmax = freq_axes.get_xbound()
         ymin,ymax = freq_axes.get_ybound()
-    
+
         # highlight the region we're fitting
-        freq_axes.axvspan(
-            config['minimum-fit-frequency'],
-            config['maximum-fit-frequency'],
-            facecolor='g', alpha=0.1, zorder=-2)
+        if config:
+            freq_axes.axvspan(
+                config['minimum-fit-frequency'],
+                config['maximum-fit-frequency'],
+                facecolor='g', alpha=0.1, zorder=-2)
 
         if A is not None:
             fitdata = breit_wigner(freq_axis, A, B, C, D)
             freq_axes.plot(freq_axis, fitdata, 'b-')
             noisefloor = D + 0*freq_axis;
             freq_axes.plot(freq_axis, noisefloor)
-
-            if B**2 < 2*A**2:
-                res_freq = _numpy.sqrt(A**2 - B**2/2)
+            res_freq = breit_wigner_resonant_frequency(A=A, B=B)
+            if res_freq > 0:
                 freq_axes.axvline(res_freq, color='b', zorder=-1)
 
         freq_axes.set_title('power spectral density %s' % timestamp)