You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

#### 131 lines 3.8 KiB Raw Permalink Blame History

 ``````#!/usr/bin/env python3 # Miscellaenous useful mathematical functions from numpy import * import scipy def numpy_raise_errors(func): def wrap(*args, **kwargs): old = seterr('raise') try: return func(*args, **kwargs) finally: seterr(**old) return wrap @numpy_raise_errors def sfit4(data, fs): """(A, f0, phi, C) = sfit4(data, fs) Compute 4-parameter (unknown-frequency) least-squares fit to sine-wave data, according to IEEE Std 1241-2010 Annex B Input: data vector of input samples fs sampling rate (Hz) Output: Parameters [A, f0, phi, C] to fit the equation x[n] = A * sin(f0/fs * 2 * pi * n + phi) + C where n is sample number. Or, as a function of time: x(t) = A * sin(f0 * 2 * pi * t + phi) + C by Jim Paris (Verified to match sfit4.m) """ N = len(data) if N < 2: raise ValueError("bad data") t = linspace(0, (N-1) / float(fs), N) # # Estimate frequency using FFT (step b) # Fc = scipy.fft.fft(data) F = abs(Fc) F[0] = 0 # eliminate DC # Find pair of spectral lines with largest amplitude: # resulting values are in F(i) and F(i+1) i = argmax(F[0:int(N/2)] + F[1:int(N/2+1)]) # Interpolate FFT to get a better result (from Markus [B37]) try: U1 = real(Fc[i]) U2 = real(Fc[i+1]) V1 = imag(Fc[i]) V2 = imag(Fc[i+1]) n = 2 * pi / N ni1 = n * i ni2 = n * (i+1) K = ((V2-V1)*sin(ni1) + (U2-U1)*cos(ni1)) / (U2-U1) Z1 = V1 * (K - cos(ni1)) / sin(ni1) + U1 Z2 = V2 * (K - cos(ni2)) / sin(ni2) + U2 i = arccos((Z2*cos(ni2) - Z1*cos(ni1)) / (Z2-Z1)) / n except Exception: # Just go with the biggest FFT peak i = argmax(F[0:int(N/2)]) # Convert to Hz f0 = i * float(fs) / N # Fit it. We'll catch exceptions here and just returns zeros # if something fails with the least squares fit, etc. try: # first guess for A0, B0 using 3-parameter fit (step c) s = zeros(3) w = 2*pi*f0 # Now iterate 7 times (step b, plus 6 iterations of step i) for idx in range(7): D = c_[cos(w*t), sin(w*t), ones(N), -s[0] * t * sin(w*t) + s[1] * t * cos(w*t)] # eqn B.16 s = linalg.lstsq(D, data, rcond=None)[0] # eqn B.18 w = w + s[3] # update frequency estimate # # Extract results # A = sqrt(s[0]*s[0] + s[1]*s[1]) # eqn B.21 f0 = w / (2*pi) phi = arctan2(s[0], s[1]) # eqn B.22 (flipped for sin instead of cos) C = s[2] return (A, f0, phi, C) except Exception: # pragma: no cover (not sure if we can hit this?) # something broke down; just return zeros return (0, 0, 0, 0) def peak_detect(data, delta=0.1): """Simple min/max peak detection algorithm, taken from my code in the disagg.m from the 10-8-5 paper. Returns an array of peaks: each peak is a tuple (n, p, is_max) where n is the row number in 'data', and p is 'data[n]', and is_max is True if this is a maximum, False if it's a minimum, """ peaks = [] cur_min = (None, inf) cur_max = (None, -inf) lookformax = False for (n, p) in enumerate(data): if p > cur_max[1]: cur_max = (n, p) if p < cur_min[1]: cur_min = (n, p) if lookformax: if p < (cur_max[1] - delta): peaks.append((cur_max[0], cur_max[1], True)) cur_min = (n, p) lookformax = False else: if p > (cur_min[1] + delta): peaks.append((cur_min[0], cur_min[1], False)) cur_max = (n, p) lookformax = True return peaks``````