From 61b8b6015efc3029ea11ed8ed03dfad09e1045cb Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Sun, 17 Aug 2014 18:12:23 +0300
Subject: [PATCH] add equalization tests

---
 tests/test_equalizer.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)
 create mode 100644 tests/test_equalizer.py

diff --git a/tests/test_equalizer.py b/tests/test_equalizer.py
new file mode 100644
index 0000000..ce50a32
--- /dev/null
+++ b/tests/test_equalizer.py
@@ -0,0 +1,27 @@
+from amodem import train
+from amodem import dsp
+from numpy.linalg import norm
+import numpy as np
+
+
+def test_fir():
+    a = [1, 0.8, -0.1, 0, 0]
+    tx = train.equalizer
+    rx = dsp.lfilter(x=tx, b=[1], a=a)
+    h_ = dsp.estimate(x=rx, y=tx, order=len(a))
+    tx_ = dsp.lfilter(x=rx, b=h_, a=[1])
+    assert norm(h_ - a) < 1e-12
+    assert (norm(tx - tx_) / norm(tx)) < 1e-12
+
+
+def test_iir():
+    alpha = 0.1
+    b = [1, -alpha]
+    tx = train.equalizer
+    rx = dsp.lfilter(x=tx, b=b, a=[1])
+    h_ = dsp.estimate(x=rx, y=tx, order=20)
+    tx_ = dsp.lfilter(x=rx, b=h_, a=[1])
+
+    h_expected = np.array([alpha ** i for i in range(len(h_))])
+    assert norm(h_ - h_expected) < 1e-12
+    assert (norm(tx - tx_) / norm(tx)) < 1e-12