From b285deb00eddff8c84a50a22d694d8c438ff1f68 Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Thu, 14 Aug 2014 09:45:05 +0300
Subject: [PATCH] support sigproc.estimate()

---
 amodem/sigproc.py     | 14 ++++++++++++++
 tests/test_sigproc.py | 28 +++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/amodem/sigproc.py b/amodem/sigproc.py
index dfe818a..ddd6751 100644
--- a/amodem/sigproc.py
+++ b/amodem/sigproc.py
@@ -33,6 +33,20 @@ def lfilter(b, a, x):
     return np.array(y)
 
 
+def estimate(x, y, order, lookahead=0):
+    offset = order - 1
+    assert offset >= lookahead
+    b = y[offset-lookahead:len(x)-lookahead]
+
+    A = []  # columns of x
+    N = len(x) - order + 1
+    for i in range(order):
+        A.append(x[i:N+i])
+
+    # switch to rows for least-squares
+    h = linalg.lstsq(np.array(A).T, b)[0]
+    return h[::-1]
+
 def train(S, training):
     A = np.array([S[1:], S[:-1], training[:-1]]).T
     b = training[1:]
diff --git a/tests/test_sigproc.py b/tests/test_sigproc.py
index cbad696..d4608b2 100644
--- a/tests/test_sigproc.py
+++ b/tests/test_sigproc.py
@@ -1,8 +1,10 @@
-from amodem import sigproc
-from amodem import config
-import numpy as np
 import random
 import itertools
+import numpy as np
+from numpy.linalg import norm
+
+from amodem import sigproc
+from amodem import config
 
 
 def test_qam():
@@ -33,3 +35,23 @@ def test_filter():
     x = [1] + [0] * 10
     y = sigproc.lfilter(b=[0.5], a=[1, -0.5], x=x)
     assert list(y) == [0.5 ** (i+1) for i in range(len(x))]
+
+def test_estimate():
+    r = np.random.RandomState(seed=0)
+    x = r.uniform(-1, 1, [1000])
+    x[:10] = 0
+    x[len(x)-10:] = 0
+
+    h = [0.1, 0.6, 0.9, 0.7, -0.2]
+    L = len(h) / 2
+
+    y = sigproc.lfilter(b=h, a=[1], x=x)
+    h_ = sigproc.estimate(
+        x=x[:len(x)-L], y=y[L:],
+        order=len(h), lookahead=L
+    )
+    y_ = sigproc.lfilter(b=h_, a=[1], x=x)
+
+    assert norm(y - y_) < 1e-12
+    assert norm(h - h_) < 1e-12
+