Align the signal

1 files changed, 36 insertions, 5 deletions

diff --git a/align/GenerateExampleTxRxIQ.py b/align/GenerateExampleTxRxIQ.py
index 448b859..fc105fd 100755
--- a/align/GenerateExampleTxRxIQ.py
+++ b/align/GenerateExampleTxRxIQ.py
@@ -43,6 +43,16 @@ oversample = 8
 
 do_plot = False
 
+def gen_omega(length):
+    if not length % 2 == 0:
+        raise ValueError('Needs an even length array.')
+    halflength = int(length/2)
+
+    omega = np.zeros(length)
+    omega[0:halflength] = 2*np.pi*np.arange(halflength)
+    omega[halflength+1:] = 2*np.pi*( np.arange(halflength+1, length)-length )
+    return omega / length
+
 def gen_signals(delay):
     ## Generate signal
 
@@ -90,10 +100,7 @@ def arg_max_corr(a, b):
     # Start by finding the coarse discretised arg_max
     coarse_max = np.argmax(np.correlate(a, b, mode='full')) - length+1
 
-    omega = np.zeros(length)
-    omega[0:halflength] = (2*np.pi*np.arange(halflength))/length
-    omega[halflength+1:] = (2*np.pi*
-            (np.arange(halflength+1, length)-length))/length
+    omega = gen_omega(length)
 
     fft_a = np.fft.fft(a)
 
@@ -110,9 +117,33 @@ def arg_max_corr(a, b):
 
     return np.real(max_arg)
 
+def delay_signal(sig, delay):
+    """Apply the delay calculated by arg_max_corr to sig"""
+    frac_delay, int_delay = np.modf(delay)
+    int_delay = int(int_delay)
+
+    print("Correcting integer delay {}".format(int_delay))
+    sig = np.roll(sig, int_delay)
+
+    print("Correcting fractional delay {}".format(frac_delay))
+    tau = -frac_delay
+
+    omega = gen_omega(len(sig))
+
+    sig_fft = np.fft.fft(sig)
+
+    rotate_vec = np.exp(1j*tau*omega)
+    rotate_vec[int(len(sig)/2)] = np.cos(np.pi*tau)
+
+    return np.fft.ifft(sig_fft * rotate_vec)
+
 for d in [2, 5, 9, 15, 34, 120]:
     a, b = gen_signals(d)
-    print("{} {}".format(d / oversample, arg_max_corr(a,b)))
+    delay = arg_max_corr(a,b)
+    print("{} {}".format(d / oversample, delay))
+
+    print("{} {}".format(d / oversample, arg_max_corr(a, delay_signal(b, delay))))
+
 
 if do_plot:
     plt.show()