From cbb6036bebb5813cebdd183008daefd9c010f1d6 Mon Sep 17 00:00:00 2001
From: Andreas Tsouchlos <an.tsouchlos@gmail.com>
Date: Mon, 7 Nov 2022 14:39:09 +0100
Subject: [PATCH] Added unit tests for the calculation of the A nd B sets;
 Fixed broken unit tests

---
 sw/test/test_proximal.py      | 46 ++++++++++++++++++++++++++++++++---
 sw/test/test_soft_decision.py |  6 ++++-
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/sw/test/test_proximal.py b/sw/test/test_proximal.py
index d5f6cd6..6c87cf6 100644
--- a/sw/test/test_proximal.py
+++ b/sw/test/test_proximal.py
@@ -14,8 +14,12 @@ class CheckParityTestCase(unittest.TestCase):
         H = np.array([[1, 0, 1, 0, 1, 0, 1],
                       [0, 1, 1, 0, 0, 1, 1],
                       [0, 0, 0, 1, 1, 1, 1]])
+        R = np.array([[0, 0, 1, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 1, 0, 0],
+                      [0, 0, 0, 0, 0, 1, 0],
+                      [0, 0, 0, 0, 0, 0, 1]])
 
-        decoder = proximal.ProximalDecoder(H)
+        decoder = proximal.ProximalDecoder(H, R)
 
         d1 = np.array([0, 1, 0, 1])
         c1 = np.dot(np.transpose(G), d1) % 2
@@ -37,19 +41,53 @@ class CheckParityTestCase(unittest.TestCase):
 class GradientTestCase(unittest.TestCase):
     """Test case for the calculation of the gradient of the code-constraint-polynomial."""
     def test_grad_h(self):
+        """Test the gradient of the code-constraint polynomial."""
         H = np.array([[1, 0, 0],
                       [0, 1, 0]])
         x = np.array([1, 2, 2])
+        R = np.array([0])
 
-        decoder = proximal.ProximalDecoder(H)
+        decoder = proximal.ProximalDecoder(H, R)
         grad = decoder._grad_h(x)
 
         expected = 4 * (x**2 - 1)*x + 2 / x * np.array([0, 2, 0])
 
-        print(f"expected: {expected}")
-
         self.assertEqual(np.array_equal(grad, expected), True)
 
+    def test_gen_A_B(self):
+        """Test the generation of the A and B sets used for the gradient calculation."""
+        # Hamming(7,4) code
+        G = np.array([[1, 1, 1, 0, 0, 0, 0],
+                      [1, 0, 0, 1, 1, 0, 0],
+                      [0, 1, 0, 1, 0, 1, 0],
+                      [1, 1, 0, 1, 0, 0, 1]])
+        H = np.array([[1, 0, 1, 0, 1, 0, 1],
+                      [0, 1, 1, 0, 0, 1, 1],
+                      [0, 0, 0, 1, 1, 1, 1]])
+        R = np.array([[0, 0, 1, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 1, 0, 0],
+                      [0, 0, 0, 0, 0, 1, 0],
+                      [0, 0, 0, 0, 0, 0, 1]])
+
+        decoder = proximal.ProximalDecoder(H, R)
+
+        expected_A = [np.array([0, 2, 4, 6]),
+                      np.array([1, 2, 5, 6]),
+                      np.array([3, 4, 5, 6])]
+        expected_B = [np.array([0]),
+                      np.array([1]),
+                      np.array([0, 1]),
+                      np.array([2]),
+                      np.array([0, 2]),
+                      np.array([1, 2]),
+                      np.array([0, 1, 2])]
+
+        for A_i, expected_A_i in zip(decoder._A, expected_A):
+            self.assertEqual(np.array_equal(A_i, expected_A_i), True)
+
+        for B_k, expected_B_k in zip(decoder._B, expected_B):
+            self.assertEqual(np.array_equal(B_k, expected_B_k), True)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/sw/test/test_soft_decision.py b/sw/test/test_soft_decision.py
index f77240b..ed84ef4 100644
--- a/sw/test/test_soft_decision.py
+++ b/sw/test/test_soft_decision.py
@@ -14,8 +14,12 @@ class CodewordGenerationTestCase(unittest.TestCase):
         H = np.array([[1, 0, 1, 0, 1, 0, 1],
                       [0, 1, 1, 0, 0, 1, 1],
                       [0, 0, 0, 1, 1, 1, 1]])
+        R = np.array([[0, 0, 1, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 1, 0, 0],
+                      [0, 0, 0, 0, 0, 1, 0],
+                      [0, 0, 0, 0, 0, 0, 1]])
 
-        decoder = naive_soft_decision.SoftDecisionDecoder(G, H)
+        decoder = naive_soft_decision.SoftDecisionDecoder(G, H, R)
 
         expected_datawords = np.array([[0, 0, 0, 0],
                                        [0, 0, 0, 1],