Reorganize code of RC SCAN Polar codec

Author: grigory

python_polar_coding/polar_codes/rc_scan.py

@@ -0,0 +1,4 @@
+from .codec import RCSCANPolarCodec
+from .decoder import RCSCANDecoder
+from .functions import *
+from .node import INFINITY, RCSCANNode

python_polar_coding/polar_codes/rc_scan/__init__.py

@@ -0,0 +1,40 @@
+from typing import Union
+
+from python_polar_coding.polar_codes.base import BasePolarCodec
+
+from .decoder import RCSCANDecoder
+
+
+class RCSCANPolarCodec(BasePolarCodec):
+    """Polar code with RC-SCAN decoding algorithm.
+
+    Based on: https://arxiv.org/pdf/1510.06495.pdf,
+    DOI: 10.1109/DASIP.2015.7367252
+
+    """
+    decoder_class = RCSCANDecoder
+
+    def __init__(
+            self,
+            N: int,
+            K: int,
+            design_snr: float = 0.0,
+            mask: Union[str, None] = None,
+            pcc_method: str = BasePolarCodec.BHATTACHARYYA,
+            I: int = 1,
+            * args, **kwargs,
+    ):
+        self.I = I
+        super().__init__(N=N, K=K,
+                         is_systematic=True,
+                         design_snr=design_snr,
+                         mask=mask,
+                         pcc_method=pcc_method)
+
+    def init_decoder(self):
+        return self.decoder_class(n=self.n, mask=self.mask, I=self.I)
+
+    def to_dict(self):
+        d = super().to_dict()
+        d.update({'I': self.I})
+        return d

python_polar_coding/polar_codes/rc_scan/codec.py

@@ -1,72 +1,11 @@
-import numba
 import numpy as np
 from anytree import PreOrderIter
 
-from ..base.functions import function_1, function_2, make_hard_decision
-from .fast_ssc_decoder import FastSSCDecoder, FastSSCNode
+from python_polar_coding.polar_codes.fast_ssc import FastSSCDecoder
 
-# LLR = 1000 is high enough to be considered as +∞ for RC-SCAN decoding
-INFINITY = 1000
-
-
-class RCSCANNode(FastSSCNode):
-
-    def compute_leaf_beta(self):
-        """Do nothing.
-
-        Unlike SC-based decoders SCAN decoders does not make decisions
-        in leaves.
-
-        """
-
-    def initialize_leaf_beta(self):
-        """Initialize BETA values on tree building.
-
-        Initialize Leaves following to Section III doi:10.1109/jsac.2014.140515
-
-        """
-        if not self.is_leaf:
-            return
-
-        if self._node_type == RCSCANNode.ZERO_NODE:
-            self._beta = self._compute_zero_node_beta(self.alpha)
-        if self._node_type == RCSCANNode.ONE_NODE:
-            self._beta = self._compute_one_node_beta(self.alpha)
-
-    def get_node_type(self):
-        """Get the type of RC SCAN Node.
-
-        * Zero node - [0, 0, 0, 0, 0, 0, 0, 0];
-        * One node - [1, 1, 1, 1, 1, 1, 1, 1];
-
-        Or other type.
-
-        """
-        if np.all(self._mask == 0):
-            return RCSCANNode.ZERO_NODE
-        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)
+from ..base import make_hard_decision
+from .functions import function_1, function_2
+from .node import RCSCANNode
 
 
 class RCSCANDecoder(FastSSCDecoder):
@@ -79,12 +18,14 @@ class RCSCANDecoder(FastSSCDecoder):
     """
     node_class = RCSCANNode
 
-    def __init__(self, n: int,
-                 mask: np.array,
-                 is_systematic: bool = True,
-                 code_min_size: int = 0,
-                 I: int = 1):
-        super().__init__(n=n, mask=mask, is_systematic=is_systematic,
+    def __init__(
+            self,
+            n: int,
+            mask: np.array,
+            code_min_size: int = 0,
+            I: int = 1,
+    ):
+        super().__init__(n=n, mask=mask, is_systematic=True,
                          code_min_size=code_min_size)
         self.I = I
 
@@ -105,8 +46,6 @@ def clean_before_decoding(self):
 
         Run this before calling `__call__` method.
 
-        TODO: Add JIT
-
         """
         for node in PreOrderIter(self._decoding_tree):
             if not node.is_leaf:

…/polar_codes/decoders/rc_scan_decoder.py …ar_coding/polar_codes/rc_scan/decoder.pypython_polar_coding/polar_codes/decoders/rc_scan_decoder.py renamed to python_polar_coding/polar_codes/rc_scan/decoder.py python_polar_coding/polar_codes/decoders/rc_scan_decoder.py renamed to python_polar_coding/polar_codes/rc_scan/decoder.py

@@ -0,0 +1,23 @@
+import numba
+
+from ..base import compute_alpha
+
+
+@numba.njit
+def function_1(a, b, c):
+    """Function 1.
+
+    Source: doi:10.1007/s12243-018-0634-7, formula 1.
+
+    """
+    return compute_alpha(a, b + c)
+
+
+@numba.njit
+def function_2(a, b, c):
+    """Function 2.
+
+    Source: doi:10.1007/s12243-018-0634-7, formula 2.
+
+    """
+    return compute_alpha(a, b) + c

python_polar_coding/polar_codes/rc_scan/functions.py

@@ -0,0 +1,67 @@
+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
+
+
+class RCSCANNode(FastSSCNode):
+
+    def compute_leaf_beta(self):
+        """Do nothing.
+
+        Unlike SC-based decoders SCAN decoders does not make decisions
+        in leaves.
+
+        """
+
+    def initialize_leaf_beta(self):
+        """Initialize BETA values on tree building.
+
+        Initialize Leaves following to Section III doi:10.1109/jsac.2014.140515
+
+        """
+        if not self.is_leaf:
+            return
+
+        if self._node_type == RCSCANNode.ZERO_NODE:
+            self._beta = self._compute_zero_node_beta(self.alpha)
+        if self._node_type == RCSCANNode.ONE_NODE:
+            self._beta = self._compute_one_node_beta(self.alpha)
+
+    def get_node_type(self):
+        """Get the type of RC SCAN Node.
+
+        * Zero node - [0, 0, 0, 0, 0, 0, 0, 0];
+        * One node - [1, 1, 1, 1, 1, 1, 1, 1];
+
+        Or other type.
+
+        """
+        if np.all(self._mask == 0):
+            return RCSCANNode.ZERO_NODE
+        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)