fix tests for custom burg implementation

timescales/autoreg/spectral.py

@@ -60,12 +60,11 @@ def compute_ar_spectrum(sig, fs, order, f_range=None, method='burg', nfft=4096,
 
     # 1d
     if method == 'burg':
-        ar, rho = burg(sig, order=order)
+        ar = burg(sig, order=order)
     else:
-        ar, rho = yule_walker(sig, order=order)
-
-    powers = arma2psd(A=-ar, rho=rho, T=fs, NFFT=nfft)
+        ar, _ = yule_walker(sig, order=order)
 
+    powers = arma2psd(A=-ar, rho=1., T=fs, NFFT=nfft)
     freqs = fftfreq(nfft, 1/fs)
     powers = powers[:len(freqs)//2]
     freqs = freqs[:len(freqs)//2]
@@ -77,7 +76,10 @@ def compute_ar_spectrum(sig, fs, order, f_range=None, method='burg', nfft=4096,
 
     return freqs, powers
 
-def burg(sig, order):
+def burg(sig, order, demean=True):
+
+    if demean:
+        sig = sig - sig.mean()
 
     # Initalize arrays for reflection coefficients
     #   and ar coefficients

timescales/decompose/decompose.py

@@ -10,10 +10,9 @@
 from scipy.signal import find_peaks
 from scipy.fft import fftfreq
 
-from statsmodels.regression.linear_model import burg
-
 from spectrum import eigen
 
+from timescales.autoreg import burg
 from timescales.sim import sim_asine_oscillation, sim_autoregressive
 
 class CAD:
@@ -378,7 +377,9 @@ def _fit(xs, sig, fs, osc_order, ar_order, *fit_args, return_params=False):
         return osc_fit, params
 
     # Fit AR Burg
-    ar_params, _ = burg(sig-osc_fit, order=ar_order, demean=True)
+    y = sig - osc_fit
+    y -= y.mean()
+    ar_params = burg(y, order=ar_order)
 
     params['params_ar'] = {'ar_' + str(ind):v for ind, v in enumerate(ar_params)}
 

timescales/tests/pipe/test_pipe.py

@@ -106,17 +106,6 @@ def test_pipe_run():
     # AR-PSD is more Lorentzian than ACF
     assert pipe_psd.results[:, 1].mean() > pipe_acf.results[:, 1].mean()
 
-    print(pipe_psd.results[:, 1])
-    print(pipe_acf.results[:, 1])
-    print()
-    print(tau)
-    print(pipe_psd.results[:, 2])
-    print(pipe_acf.results[:, 2])
-    print()
-    print(10)
-    print(pipe_psd.results[:, 0])
-    print(pipe_acf.results[:, 0])
-
     # AR-PSD timescales are more accurate (MAE)
     assert np.abs(pipe_psd.results[:, 2] - tau).mean() < \
         np.abs(pipe_acf.results[:, 2] - tau).mean()