Add Fast SCAN polar codec

Author: grigory

python_polar_coding/polar_codes/__init__.py

@@ -1,3 +1,4 @@
+from .fast_scan import FastSCANCodec
 from .fast_ssc import FastSSCPolarCodec
 from .g_fast_ssc import GeneralizedFastSSCPolarCodec
 from .rc_scan import RCSCANPolarCodec

python_polar_coding/polar_codes/fast_scan/__init__.py

@@ -0,0 +1,4 @@
+from .codec import *
+from .decoder import *
+from .functions import *
+from .node import *

python_polar_coding/polar_codes/fast_scan/codec.py

@@ -0,0 +1,7 @@
+from python_polar_coding.polar_codes.rc_scan import RCSCANPolarCodec
+
+from .decoder import FastSCANDecoder
+
+
+class FastSCANCodec(RCSCANPolarCodec):
+    decoder_class = FastSCANDecoder

python_polar_coding/polar_codes/fast_scan/decoder.py

@@ -0,0 +1,7 @@
+from python_polar_coding.polar_codes.rc_scan import RCSCANDecoder
+
+from .node import FastSCANNode
+
+
+class FastSCANDecoder(RCSCANDecoder):
+    node_class = FastSCANNode

python_polar_coding/polar_codes/fast_scan/functions.py

@@ -0,0 +1,26 @@
+import numba
+import numpy as np
+
+
+@numba.njit
+def compute_repetition_beta(alpha) -> np.array:
+    """Compute beta value for Repetition node."""
+    alpha_sum = np.sum(alpha)
+    return -1 * alpha + alpha_sum
+
+
+@numba.njit
+def compute_spc_beta(alpha) -> np.array:
+    """Compute beta value for Single parity node."""
+    all_sign = np.sign(np.prod(alpha))
+    abs_alpha = np.fabs(alpha)
+    first_min_idx, second_min_idx = np.argsort(abs_alpha)[:2]
+
+    result = np.sign(alpha) * all_sign
+    for i in range(result.size):
+        if i == first_min_idx:
+            result[i] *= abs_alpha[second_min_idx]
+        else:
+            result[i] *= abs_alpha[first_min_idx]
+
+    return result

python_polar_coding/polar_codes/fast_scan/node.py

@@ -0,0 +1,56 @@
+from python_polar_coding.polar_codes.rc_scan import (
+    RCSCANNode,
+    compute_beta_one_node,
+    compute_beta_zero_node,
+)
+
+from .functions import compute_repetition_beta, compute_spc_beta
+
+
+class FastSCANNode(RCSCANNode):
+
+    def compute_leaf_beta(self):
+        """Compute leaf beta."""
+        if self.is_repetition:
+            self.beta = compute_repetition_beta(self.alpha)
+        if self.is_parity:
+            self.beta = compute_spc_beta(self.alpha)
+
+    def initialize_leaf_beta(self):
+        """Initialize BETA values on tree building.
+
+        Initialize ZERO and ONE nodes following to Section III
+        doi:10.1109/jsac.2014.140515
+
+        """
+        if not self.is_leaf:
+            return
+
+        if self.is_parity or self.is_repetition:
+            return
+
+        if self.is_zero:
+            self._beta = compute_beta_zero_node(self.alpha)
+        if self.is_one:
+            self._beta = compute_beta_one_node(self.alpha)
+
+    def get_node_type(self):
+        """Get the type of Fast SCAN Node.
+
+        * Zero node - [0, 0, 0, 0, 0, 0, 0, 0];
+        * One node - [1, 1, 1, 1, 1, 1, 1, 1];
+        * Single parity check node - [0, 1, 1, 1, 1, 1, 1, 1];
+        * Repetition node - [0, 0, 0, 0, 0, 0, 0, 1].
+
+        Or other type.
+
+        """
+        if self._check_is_zero(self._mask):
+            return self.ZERO_NODE
+        if self._check_is_one(self._mask):
+            return self.ONE_NODE
+        if self._check_is_rep(self._mask):
+            return self.REPETITION
+        if self._check_is_spc(self._mask):
+            return self.SINGLE_PARITY_CHECK
+        return self.OTHER

python_polar_coding/polar_codes/rc_scan/__init__.py

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

python_polar_coding/polar_codes/rc_scan/decoder.py

@@ -48,7 +48,7 @@ def clean_before_decoding(self):
 
         """
         for node in PreOrderIter(self._decoding_tree):
-            if not node.is_leaf:
+            if not (node.is_zero or node.is_one):
                 node.beta *= 0
 
     def compute_intermediate_alpha(self, leaf):
@@ -81,8 +81,9 @@ def compute_intermediate_beta(self, node):
 
     @property
     def result(self):
-        if self.is_systematic:
-            return make_hard_decision(self.root.alpha + self._compute_result_beta())
+        if not self.is_systematic:
+            raise TypeError('Code must be systematic')
+        return make_hard_decision(self.root.alpha + self._compute_result_beta())
 
     @staticmethod
     def compute_left_alpha(parent_alpha, beta):

python_polar_coding/polar_codes/rc_scan/node.py

@@ -14,8 +14,6 @@ def compute_leaf_beta(self):
         in leaves.
 
         """
-        if self.is_zero or self.is_one:
-            return
 
     def initialize_leaf_beta(self):
         """Initialize BETA values on tree building.

python_polar_coding/simulation/simulation.py

@@ -9,6 +9,7 @@
     GeneralizedFastSSCPolarCodec,
     RCSCANPolarCodec,
 )
+from ..polar_codes.fast_scan import FastSCANCodec
 from . import functions, http
 
 
@@ -17,7 +18,13 @@ class CodeTypes:
     FAST_SSC = 'fast-ssc'
     RC_SCAN = 'rc_scan'
     G_FAST_SSC = 'g-fast-ssc'
-    ALL = [FAST_SSC, RC_SCAN]
+    FAST_SCAN = 'fast-scan'
+    ALL = [
+        FAST_SSC,
+        RC_SCAN,
+        G_FAST_SSC,
+        FAST_SCAN,
+    ]
 
 
 class ChannelTypes:
@@ -28,6 +35,7 @@ class ChannelTypes:
     CodeTypes.FAST_SSC: FastSSCPolarCodec,
     CodeTypes.RC_SCAN: RCSCANPolarCodec,
     CodeTypes.G_FAST_SSC: GeneralizedFastSSCPolarCodec,
+    CodeTypes.FAST_SCAN: FastSCANCodec,
 }