_spectral.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. # Author: Pim Schellart
  2. # 2010 - 2011
  3. """Tools for spectral analysis of unequally sampled signals."""
  4. import numpy as np
  5. #pythran export _lombscargle(float64[], float64[], float64[])
  6. def _lombscargle(x, y, freqs):
  7. """
  8. _lombscargle(x, y, freqs)
  9. Computes the Lomb-Scargle periodogram.
  10. Parameters
  11. ----------
  12. x : array_like
  13. Sample times.
  14. y : array_like
  15. Measurement values (must be registered so the mean is zero).
  16. freqs : array_like
  17. Angular frequencies for output periodogram.
  18. Returns
  19. -------
  20. pgram : array_like
  21. Lomb-Scargle periodogram.
  22. Raises
  23. ------
  24. ValueError
  25. If the input arrays `x` and `y` do not have the same shape.
  26. See also
  27. --------
  28. lombscargle
  29. """
  30. # Check input sizes
  31. if x.shape != y.shape:
  32. raise ValueError("Input arrays do not have the same size.")
  33. # Create empty array for output periodogram
  34. pgram = np.empty_like(freqs)
  35. c = np.empty_like(x)
  36. s = np.empty_like(x)
  37. for i in range(freqs.shape[0]):
  38. xc = 0.
  39. xs = 0.
  40. cc = 0.
  41. ss = 0.
  42. cs = 0.
  43. c[:] = np.cos(freqs[i] * x)
  44. s[:] = np.sin(freqs[i] * x)
  45. for j in range(x.shape[0]):
  46. xc += y[j] * c[j]
  47. xs += y[j] * s[j]
  48. cc += c[j] * c[j]
  49. ss += s[j] * s[j]
  50. cs += c[j] * s[j]
  51. if freqs[i] == 0:
  52. raise ZeroDivisionError()
  53. tau = np.arctan2(2 * cs, cc - ss) / (2 * freqs[i])
  54. c_tau = np.cos(freqs[i] * tau)
  55. s_tau = np.sin(freqs[i] * tau)
  56. c_tau2 = c_tau * c_tau
  57. s_tau2 = s_tau * s_tau
  58. cs_tau = 2 * c_tau * s_tau
  59. pgram[i] = 0.5 * (((c_tau * xc + s_tau * xs)**2 /
  60. (c_tau2 * cc + cs_tau * cs + s_tau2 * ss)) +
  61. ((c_tau * xs - s_tau * xc)**2 /
  62. (c_tau2 * ss - cs_tau * cs + s_tau2 * cc)))
  63. return pgram