Add sine fit prototype for automated testing.

a662ac07 · Matthieu Cattin · 6d5f2088 · a662ac07
Commit a662ac07 authored Mar 13, 2013 by Matthieu Cattin
Hide whitespace changes
Inline Side-by-side

Showing with 143 additions and 0 deletions

test_sin_fit.py test/fmcadc100m14b4cha/python/test_sin_fit.py +143 -0

No files found.
--- a/test/fmcadc100m14b4cha/python/test_sin_fit.py
+++ b/test/fmcadc100m14b4cha/python/test_sin_fit.py
+#!   /usr/bin/env   python
+#    coding: utf8
+
+# Copyright CERN, 2013
+# Author: Matthieu Cattin <matthieu.cattin@cern.ch>
+# Licence: GPL v2 or later.
+# Website: http://www.ohwr.org
+
+
+# Import system modules
+import sys
+import time
+import os
+
+# Import specific modules
+from numpy import *
+from pylab import *
+from ctypes import *
+import scipy.optimize as optimize
+import scipy.fftpack as fftpack
+
+
+"""
+Sine wave fit test
+"""
+
+
+def gen_sine(frequency=1, amplitude=1, noise=0.0, cycles=1, rate=10, phase=0.0, offset=0.0):
+    # Sample a noisy sine wave and return a list of sample amplitudes
+
+    #print("frequency : %f"%frequency)
+    #print("amplitude : %f"%amplitude)
+    #print("noise     : %f"%noise)
+    #print("cycles    : %f"%cycles)
+    #print("rate      : %f"%rate)
+    #print("phase     : %f"%phase)
+    #print("offset    : %f"%offset)
+
+    # Calculated derived values
+    dt          = 1/float(frequency*rate)
+    #print("dt     : %f"%dt)
+    period      = 1/float(frequency)
+    #print("period : %f"%period)
+    steps       = arange(0, cycles*period, dt)
+
+    samples     = []    # create an empty list
+
+    # calculate amplitude at each sample point
+    # for given number of cycles
+    for t in steps:
+        angle   = 2*pi*t/period + phase     # in radians
+        a       = amplitude*sin(angle) + noise*randn() + offset
+        samples.append(a)
+
+    return samples
+
+class Parameter:
+    def __init__(self, value):
+            self.value = value
+
+    def set(self, value):
+            self.value = value
+
+    def __call__(self):
+            return self.value
+
+def fit(function, parameters, y, x = None):
+    def f(params):
+        i = 0
+        for p in parameters:
+            p.set(params[i])
+            i += 1
+        return y - function(x)
+
+    if x is None: x = np.arange(y.shape[0])
+    p = [param() for param in parameters]
+    return optimize.leastsq(f, p, full_output=True, ftol=1e-6, xtol=1e-6)
+
+def my_sine(x, o, a, f, p):
+    return o + a * sin(f*x + p)
+
+def plot_result(data, fit, sin_upper, sin_lower, ylimit, x_data):
+    clf()
+    plot(x_data, data, 'k', label='Data')
+    plot(x_data, fit, 'r', label='Fit')
+    plot(x_data, sin_upper, 'b', label='Upper limit')
+    plot(x_data, sin_lower, 'g', label='Lower limit')
+    ylim(-ylimit-(ylimit/10.0), ylimit+(ylimit/10.0))
+    grid(which='both')
+    legend()
+    draw()
+    show()
+    return 0
+
+
+def main (default_directory = '.'):
+
+    # Creates fake input data (noisy sine wave)
+    ampl = 2
+    freq = 1
+    rate = 500
+    phase = 0.5*pi
+    noise = 0.02
+    cycles = 2
+    offset = 0
+    yReal = gen_sine(freq, ampl, noise, cycles, rate, phase, offset)
+    xReal = arange(len(yReal))
+    print("Input sine params:\nOffset: %f, Amplitude: %f, Frequency: %f, Phase: %f"%(offset, ampl, freq, phase))
+
+    # rising edge detection
+    zero_cross = where(diff(sign(yReal)))[0]
+    print zero_cross
+    x = diff(sign(yReal))
+    rising_edge = where(select([x>0],[x]))[0]
+    print rising_edge
+
+    # give fit initial parameters
+    guess_ampl = max(yReal)
+    yhat = fftpack.rfft(yReal)
+    idx = (yhat**2).argmax()
+    freqs = fftpack.rfftfreq(len(yReal), d = (xReal[1]-xReal[0])/(2*pi))
+    guess_freq = freqs[idx]
+    guess_phase = 0.0
+    guess_offset = mean(yReal)
+
+    guess_params = [guess_offset, guess_ampl, guess_freq, guess_phase]
+
+    print "Guessed params:"
+    print guess_params
+
+    (fit_offset, fit_ampl, fit_freq, fit_phase), pcov = optimize.curve_fit(my_sine, xReal, yReal, guess_params)
+
+    print("Fit params:\nOffset: %f, Amplitude: %f, Frequency: %f, Phase: %f"%(fit_offset, fit_ampl, fit_freq, fit_phase))
+
+    fit_sin = my_sine(xReal, fit_offset, fit_ampl, fit_freq, fit_phase)
+
+    upper = my_sine(xReal, fit_offset+0.1, fit_ampl, fit_freq, fit_phase)
+    lower = my_sine(xReal, fit_offset-0.1, fit_ampl, fit_freq, fit_phase)
+    plot_result(yReal, fit_sin, upper, lower, 2, xReal)
+
+
+if __name__ == '__main__' :
+    main()