Add sub sample accurate fft lag finder

1 files changed, 32 insertions, 64 deletions

diff --git a/src/dab_util.py b/src/dab_util.py
index aa2030f..8ab1e03 100644
--- a/src/dab_util.py
+++ b/src/dab_util.py
@@ -1,6 +1,7 @@
 import numpy as np
 import scipy
 import matplotlib.pyplot as plt
+import fftconvolve
 
 c = {
         "bw":1536000
@@ -41,67 +42,34 @@ def crop_signal(signal, n_window = 1000, n_zeros = 120000, debug = False):
     signal = signal[max(0,idx_start - n_zeros): min(idx_end + n_zeros, signal.shape[0] -1)]
     return signal
 
-def lagcorr(x,y,lag=None,verbose=False):
-    '''Compute lead-lag correlations between 2 time series.
-
-    <x>,<y>: 1-D time series.
-    <lag>: lag option, could take different forms of <lag>:
-          if 0 or None, compute ordinary correlation and p-value;
-          if positive integer, compute lagged correlation with lag
-          upto <lag>;
-          if negative integer, compute lead correlation with lead
-          upto <-lag>;
-          if pass in an list or tuple or array of integers, compute
-          lead/lag correlations at different leads/lags.
-
-    Note: when talking about lead/lag, uses <y> as a reference.
-    Therefore positive lag means <x> lags <y> by <lag>, computation is
-    done by shifting <x> to the left hand side by <lag> with respect to
-    <y>.
-    Similarly negative lag means <x> leads <y> by <lag>, computation is
-    done by shifting <x> to the right hand side by <lag> with respect to
-    <y>.
-
-    Return <result>: a (n*2) array, with 1st column the correlation
-    coefficients, 2nd column correpsonding p values.
-
-    Currently only works for 1-D arrays.
-    '''
-
-    import numpy
-    from scipy.stats import pearsonr
-
-    if len(x)!=len(y):
-        raise Exception('Input variables of different lengths.')
-
-    #--------Unify types of <lag>-------------
-    if numpy.isscalar(lag):
-        if abs(lag)>=len(x):
-            raise Exception('Maximum lag equal or larger than array.')
-        if lag<0:
-            lag=-numpy.arange(abs(lag)+1)
-        elif lag==0:
-            lag=[0,]
-        else:
-            lag=numpy.arange(lag+1)
-    elif lag is None:
-        lag=[0,]
-    else:
-        lag=numpy.asarray(lag)
-
-    #-------Loop over lags---------------------
-    result=[]
-    if verbose:
-        print '\n#<lagcorr>: Computing lagged-correlations at lags:',lag
-
-    for ii in lag:
-        if ii<0:
-            result.append(pearsonr(x[:ii],y[-ii:]))
-        elif ii==0:
-            result.append(pearsonr(x,y))
-        elif ii>0:
-            result.append(pearsonr(x[ii:],y[:-ii]))
-
-    result=numpy.asarray(result)
-
-    return result
+#def fftlag(signal_original, signal_rec):
+#    """
+#    Efficient way to find lag between two signals
+#    Args:
+#        signal_original: The signal that has been sent
+#        signal_rec: The signal that has been recored
+#    """
+#    c = np.flipud(scipy.signal.fftconvolve(signal_original,np.flipud(signal_rec)))
+#    #plt.plot(c)
+#    return np.argmax(c) - signal_original.shape[0] + 1
+
+def fftlag(signal_original, signal_rec, n_upsampling = 1):
+    """
+    Efficient way to find lag between two signals
+    Args:
+        signal_original: The signal that has been sent
+        signal_rec: The signal that has been recored
+    """
+    c = np.flipud(fftconvolve.fftconvolve(signal_original,np.flipud(signal_rec), n_upsampling))
+    #plt.plot(c)
+    return (np.argmax(c) - signal_original.shape[0] + 1)
+
+def get_amp_ratio(ampl_1, ampl_2, a_out_abs, a_in_abs):
+    idxs = (a_in_abs > ampl_1) & (a_in_abs < ampl_2)
+    ratio = a_out_abs[idxs] / a_in_abs[idxs]
+    return ratio.mean(), ratio.var()
+
+def get_phase(ampl_1, ampl_2, a_out, a_in):
+    idxs = (np.abs(a_in) > ampl_1) & (np.abs(a_in) < ampl_2)
+    ratio = np.angle(a_out[idxs], deg=True) - np.angle(a_in[idxs], deg=True)
+    return ratio.mean(), ratio.var()