vibration_analyze: Add smooth_window to PSD fitting

[calibcant.git] / calibcant / vibration_analyze.py

diff --git a/calibcant/vibration_analyze.py b/calibcant/vibration_analyze.py
index f5b83dec8cfa7f338262fb90f41de697d8fe5d2b..a0d383e6370fba1b733e89685af7c906da2ca40a 100644 (file)
--- a/calibcant/vibration_analyze.py
+++ b/calibcant/vibration_analyze.py
@@ -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]