Add small improvements

Author: grigory

README.md

@@ -43,7 +43,7 @@ for snr in snr_range:
     for _ in range(messages):
         msg = generate_binary_message(size=K)
         encoded = codec.encode(msg)
-        transmitted = bpsk.transmit(message=msg, snr_db=snr)
+        transmitted = bpsk.transmit(message=encoded, snr_db=snr)
         decoded = codec.decode(transmitted)
 
         bit_errors, frame_error = compute_fails(msg, decoded)

python_polar_coding/polar_codes/fast_ssc/decoder.py

@@ -10,10 +10,13 @@ class FastSSCDecoder(SCDecoder):
     """Implements Fast SSC decoding algorithm."""
     node_class = FastSSCNode
 
-    def __init__(self, n: int,
-                 mask: np.array,
-                 is_systematic: bool = True,
-                 code_min_size: int = 0):
+    def __init__(
+            self,
+            n: int,
+            mask: np.array,
+            is_systematic: bool = True,
+            code_min_size: int = 0,
+    ):
         super().__init__(n=n, mask=mask, is_systematic=is_systematic)
         self._decoding_tree = self.node_class(
             mask=self.mask,

python_polar_coding/polar_codes/fast_ssc/node.py

@@ -40,6 +40,9 @@ def __init__(self, mask, name=ROOT, N_min=None, **kwargs):
         self.is_computed = False
         self._build_decoding_tree()
 
+    def __str__(self):
+        return ''.join([str(m) for m in self._mask])
+
     @property
     def N(self):
         return self._mask.size
@@ -135,10 +138,10 @@ def get_node_type(self):
             return FastSSCNode.ZERO_NODE
         if self._check_is_one(self._mask) and self.N >= self.one_min_size:
             return FastSSCNode.ONE_NODE
-        if self.N >= self.spc_min_size and self._check_is_spc(self._mask):
-            return FastSSCNode.SINGLE_PARITY_CHECK
         if self.N >= self.repetition_min_size and self._check_is_parity(self._mask):  # noqa
             return FastSSCNode.REPETITION
+        if self.N >= self.spc_min_size and self._check_is_spc(self._mask):
+            return FastSSCNode.SINGLE_PARITY_CHECK
         return FastSSCNode.OTHER
 
     def _check_is_one(self, mask):

python_polar_coding/polar_codes/g_fast_ssc/node.py

@@ -27,9 +27,9 @@ def get_node_type(self):
         ntype = super().get_node_type()
         if ntype != self.OTHER:
             return ntype
-        if self._check_is_g_repetition():
+        if self._check_is_g_repetition(self._mask):
             return self.G_REPETITION
-        if self._check_is_rg_parity():
+        if self._check_is_rg_parity(self._mask):
             return self.RG_PARITY
         return self.OTHER
 
@@ -46,15 +46,15 @@ def _build_decoding_tree(self):
         cls(mask=left_mask, name=self.LEFT, N_min=self.M, parent=self, AF=self.AF)
         cls(mask=right_mask, name=self.RIGHT, N_min=self.M, parent=self, AF=self.AF)
 
-    def _check_is_g_repetition(self):
+    def _check_is_g_repetition(self, mask):
         """Check the node is Generalized Repetition node.
 
         Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, A.
 
         """
         # 1. Split mask into T chunks, T in range [2, 4, ..., N/2]
         for t in splits(self.MIN_CHUNKS, self.N // 2):
-            chunks = np.split(self._mask, t)
+            chunks = np.split(mask, t)
 
             last = chunks[-1]
             last_ok = (
@@ -74,15 +74,15 @@ def _check_is_g_repetition(self):
 
         return False
 
-    def _check_is_rg_parity(self):
+    def _check_is_rg_parity(self, mask):
         """Check the node is Relaxed Generalized Parity Check node.
 
         Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, B.
 
         """
         # 1. Split mask into T chunks, T in range [2, 4, ..., N/2]
         for t in splits(self.MIN_CHUNKS, self.N // 2):
-            chunks = np.split(self._mask, t)
+            chunks = np.split(mask, t)
 
             first = chunks[0]
             if not self._check_is_zero(first):