Jim Paris
cfc66b6847
This found a small number of real bugs too, for example, this one that looked weird because of a 2to3 conversion, but was wrong both before and after: - except IndexError as TypeError: + except (IndexError, TypeError):
131 lines
3.8 KiB
Python
131 lines
3.8 KiB
Python
#!/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
|