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.

math.py 3.5 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. #!/usr/bin/python
  2. # Miscellaenous useful mathematical functions
  3. from nilmdb.utils.printf import *
  4. from numpy import *
  5. from scipy import *
  6. def sfit4(data, fs):
  7. """(A, f0, phi, C) = sfit4(data, fs)
  8. Compute 4-parameter (unknown-frequency) least-squares fit to
  9. sine-wave data, according to IEEE Std 1241-2010 Annex B
  10. Input:
  11. data vector of input samples
  12. fs sampling rate (Hz)
  13. Output:
  14. Parameters [A, f0, phi, C] to fit the equation
  15. x[n] = A * sin(f0/fs * 2 * pi * n + phi) + C
  16. where n is sample number. Or, as a function of time:
  17. x(t) = A * sin(f0 * 2 * pi * t + phi) + C
  18. by Jim Paris
  19. (Verified to match sfit4.m)
  20. """
  21. N = len(data)
  22. t = linspace(0, (N-1) / float(fs), N)
  23. ## Estimate frequency using FFT (step b)
  24. Fc = fft(data)
  25. F = abs(Fc)
  26. F[0] = 0 # eliminate DC
  27. # Find pair of spectral lines with largest amplitude:
  28. # resulting values are in F(i) and F(i+1)
  29. i = argmax(F[0:int(N/2)] + F[1:int(N/2+1)])
  30. # Interpolate FFT to get a better result (from Markus [B37])
  31. try:
  32. U1 = real(Fc[i])
  33. U2 = real(Fc[i+1])
  34. V1 = imag(Fc[i])
  35. V2 = imag(Fc[i+1])
  36. n = 2 * pi / N
  37. ni1 = n * i
  38. ni2 = n * (i+1)
  39. K = ((V2-V1)*sin(ni1) + (U2-U1)*cos(ni1)) / (U2-U1)
  40. Z1 = V1 * (K - cos(ni1)) / sin(ni1) + U1
  41. Z2 = V2 * (K - cos(ni2)) / sin(ni2) + U2
  42. i = arccos((Z2*cos(ni2) - Z1*cos(ni1)) / (Z2-Z1)) / n
  43. except Exception:
  44. # Just go with the biggest FFT peak
  45. i = argmax(F[0:int(N/2)])
  46. # Convert to Hz
  47. f0 = i * float(fs) / N
  48. # Fit it. We'll catch exceptions here and just returns zeros
  49. # if something fails with the least squares fit, etc.
  50. try:
  51. # first guess for A0, B0 using 3-parameter fit (step c)
  52. s = zeros(3)
  53. w = 2*pi*f0
  54. # Now iterate 7 times (step b, plus 6 iterations of step i)
  55. for idx in range(7):
  56. D = c_[cos(w*t), sin(w*t), ones(N),
  57. -s[0] * t * sin(w*t) + s[1] * t * cos(w*t) ] # eqn B.16
  58. s = linalg.lstsq(D, data, rcond=None)[0] # eqn B.18
  59. w = w + s[3] # update frequency estimate
  60. ## Extract results
  61. A = sqrt(s[0]*s[0] + s[1]*s[1]) # eqn B.21
  62. f0 = w / (2*pi)
  63. phi = arctan2(s[0], s[1]) # eqn B.22 (flipped for sin instead of cos)
  64. C = s[2]
  65. return (A, f0, phi, C)
  66. except Exception as e:
  67. # something broke down; just return zeros
  68. return (0, 0, 0, 0)
  69. def peak_detect(data, delta = 0.1):
  70. """Simple min/max peak detection algorithm, taken from my code
  71. in the disagg.m from the 10-8-5 paper.
  72. Returns an array of peaks: each peak is a tuple
  73. (n, p, is_max)
  74. where n is the row number in 'data', and p is 'data[n]',
  75. and is_max is True if this is a maximum, False if it's a minimum,
  76. """
  77. peaks = [];
  78. cur_min = (None, inf)
  79. cur_max = (None, -inf)
  80. lookformax = False
  81. for (n, p) in enumerate(data):
  82. if p > cur_max[1]:
  83. cur_max = (n, p)
  84. if p < cur_min[1]:
  85. cur_min = (n, p)
  86. if lookformax:
  87. if p < (cur_max[1] - delta):
  88. peaks.append((cur_max[0], cur_max[1], True))
  89. cur_min = (n, p)
  90. lookformax = False
  91. else:
  92. if p > (cur_min[1] + delta):
  93. peaks.append((cur_min[0], cur_min[1], False))
  94. cur_max = (n, p)
  95. lookformax = True
  96. return peaks