Small refactoring of RC SCAN

Author: grigory

python_polar_coding/polar_codes/base/__init__.py

@@ -1,4 +1,6 @@
 from .codec import BaseCRCPolarCodec, BasePolarCodec
+from .constants import *
 from .decoder import BaseDecoder
 from .decoding_path import DecodingPathMixin
+from .encoder import *
 from .functions import *

python_polar_coding/polar_codes/base/constants.py

@@ -0,0 +1,2 @@
+# LLR = 1000 is high enough to be considered as +∞ for SCAN decoding
+INFINITY = 1000

python_polar_coding/polar_codes/fast_ssc/decoder.py

@@ -1,11 +1,11 @@
 import numpy as np
 from anytree import PreOrderIter
 
-from python_polar_coding.polar_codes.sc import SCDecoder
 from python_polar_coding.polar_codes.base.functions import (
     compute_left_alpha,
     compute_right_alpha,
 )
+from python_polar_coding.polar_codes.sc import SCDecoder
 
 from .node import FastSSCNode
 
@@ -75,6 +75,10 @@ def compute_intermediate_alpha(self, leaf):
             if node.is_computed:
                 continue
 
+            # No need to compute zero node because output is vector of zeros
+            if node.is_zero:
+                continue
+
             parent_alpha = node.parent.alpha
 
             if node.is_left:

python_polar_coding/polar_codes/fast_ssc/node.py

@@ -43,6 +43,9 @@ def __init__(self, mask, name=ROOT, N_min=None, **kwargs):
     def __str__(self):
         return ''.join([str(m) for m in self._mask])
 
+    def __len__(self):
+        return self._mask.size
+
     @property
     def N(self):
         return self._mask.size
@@ -100,6 +103,26 @@ def repetition_min_size(self):
     def spc_min_size(self):
         return self.M or self.__class__.SPC_MIN_SIZE
 
+    @property
+    def is_zero(self):
+        return self._node_type == self.ZERO_NODE
+
+    @property
+    def is_one(self):
+        return self._node_type == self.ONE_NODE
+
+    @property
+    def is_repetition(self):
+        return self._node_type == self.REPETITION
+
+    @property
+    def is_parity(self):
+        return self._node_type == self.SINGLE_PARITY_CHECK
+
+    @property
+    def type(self):
+        return self._node_type
+
     def to_dict(self):
         return {
             'type': self._node_type,

python_polar_coding/polar_codes/g_fast_ssc/node.py

@@ -23,6 +23,14 @@ def __init__(self, AF=1, *args, **kwargs):
         self.mask_steps = None
         super().__init__(*args, **kwargs)
 
+    @property
+    def is_g_repetition(self):
+        return self._node_type == self.G_REPETITION
+
+    @property
+    def is_rg_parity(self):
+        return self._node_type == self.RG_PARITY
+
     def get_node_type(self):
         ntype = super().get_node_type()
         if ntype != self.OTHER:

python_polar_coding/polar_codes/rc_scan/decoder.py

@@ -54,10 +54,9 @@ def clean_before_decoding(self):
     def compute_intermediate_alpha(self, leaf):
         """Compute intermediate Alpha values (LLR)."""
         for node in leaf.path[1:]:
-            if node.is_computed or node.is_leaf:
+            if node.is_computed or node.is_zero or node.is_one:
                 continue
 
-            node.is_computed = True
             parent_alpha = node.parent.alpha
 
             if node.is_left:
@@ -68,6 +67,8 @@ def compute_intermediate_alpha(self, leaf):
                 left_beta = node.siblings[0].beta
                 node.alpha = self.compute_right_alpha(parent_alpha, left_beta)
 
+            node.is_computed = True
+
     def compute_intermediate_beta(self, node):
         """Compute intermediate BETA values."""
         parent = node.parent

python_polar_coding/polar_codes/rc_scan/functions.py

@@ -1,6 +1,7 @@
 import numba
+import numpy as np
 
-from ..base import compute_alpha
+from ..base import INFINITY, compute_alpha
 
 
 @numba.njit
@@ -21,3 +22,23 @@ def function_2(a, b, c):
 
     """
     return compute_alpha(a, b) + c
+
+
+@numba.njit
+def compute_beta_zero_node(alpha):
+    """Compute beta values for ZERO node.
+
+    https://arxiv.org/pdf/1510.06495.pdf Section III.C.
+
+    """
+    return np.ones(alpha.size, dtype=np.double) * INFINITY
+
+
+@numba.njit
+def compute_beta_one_node(alpha):
+    """Compute beta values for ONE node.
+
+    https://arxiv.org/pdf/1510.06495.pdf Section III.C.
+
+    """
+    return np.zeros(alpha.size, dtype=np.double)

python_polar_coding/polar_codes/rc_scan/node.py

@@ -1,21 +1,21 @@
-import numba
 import numpy as np
 
 from python_polar_coding.polar_codes.fast_ssc import FastSSCNode
 
-# LLR = 1000 is high enough to be considered as +∞ for RC-SCAN decoding
-INFINITY = 1000
+from .functions import compute_beta_one_node, compute_beta_zero_node
 
 
 class RCSCANNode(FastSSCNode):
 
     def compute_leaf_beta(self):
-        """Do nothing.
+        """Do nothing for ZERO and ONE nodes.
 
         Unlike SC-based decoders SCAN decoders does not make decisions
         in leaves.
 
         """
+        if self.is_zero or self.is_one:
+            return
 
     def initialize_leaf_beta(self):
         """Initialize BETA values on tree building.
@@ -27,9 +27,9 @@ def initialize_leaf_beta(self):
             return
 
         if self._node_type == RCSCANNode.ZERO_NODE:
-            self._beta = self._compute_zero_node_beta(self.alpha)
+            self._beta = compute_beta_zero_node(self.alpha)
         if self._node_type == RCSCANNode.ONE_NODE:
-            self._beta = self._compute_one_node_beta(self.alpha)
+            self._beta = compute_beta_one_node(self.alpha)
 
     def get_node_type(self):
         """Get the type of RC SCAN Node.
@@ -45,23 +45,3 @@ def get_node_type(self):
         if np.all(self._mask == 1):
             return RCSCANNode.ONE_NODE
         return RCSCANNode.OTHER
-
-    @staticmethod
-    @numba.njit
-    def _compute_zero_node_beta(llr):
-        """Compute beta values for ZERO node.
-
-        https://arxiv.org/pdf/1510.06495.pdf Section III.C.
-
-        """
-        return np.ones(llr.size, dtype=np.double) * INFINITY
-
-    @staticmethod
-    @numba.njit
-    def _compute_one_node_beta(llr):
-        """Compute beta values for ONE node.
-
-        https://arxiv.org/pdf/1510.06495.pdf Section III.C.
-
-        """
-        return np.zeros(llr.size, dtype=np.double)

python_polar_coding/tests/test_fast_ssc/test_decoder.py

@@ -169,9 +169,9 @@ def test_complex(self):
             np.array([-0.8674, 0.9677, ]),
             np.array([-0.7723, 1.8675, -0.2039, -2.5321, ]),
             np.array([-0.2514, ]),
-            np.array([0.8554, ]),
+            np.array([0, ]),
             np.array([-1.2404, ]),
-            np.array([2.9912, ]),
+            np.array([0, ]),
             np.array([-2.3952, -1.3818, 4.7891, -6.4949, ]),
         ]