ba-thesis/sw/test/test_proximal.py

56 lines
1.7 KiB
Python

import unittest
import numpy as np
from decoders import proximal
class CheckParityTestCase(unittest.TestCase):
"""Test case for the check_parity function."""
def test_check_parity(self):
# 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]])
decoder = proximal.ProximalDecoder(H)
d1 = np.array([0, 1, 0, 1])
c1 = np.dot(np.transpose(G), d1) % 2
d2 = np.array([0, 0, 0, 0])
c2 = np.dot(np.transpose(G), d2) % 2
d3 = np.array([1, 1, 1, 1])
c3 = np.dot(np.transpose(G), d3) % 2
invalid_codeword = np.array([0, 1, 1, 0, 1, 1, 1])
self.assertEqual(decoder._check_parity(c1), True)
self.assertEqual(decoder._check_parity(c2), True)
self.assertEqual(decoder._check_parity(c3), True)
self.assertEqual(decoder._check_parity(invalid_codeword), False)
class GradientTestCase(unittest.TestCase):
"""Test case for the calculation of the gradient of the code-constraint-polynomial"""
def test_grad_h(self):
H = np.array([[1, 0, 0],
[0, 1, 0]])
x = np.array([1, 2, 2])
decoder = proximal.ProximalDecoder(H)
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)
if __name__ == "__main__":
unittest.main()