Refactor alpha computation

Author: grigory

python_polar_coding/polar_codes/base/functions.py

@@ -32,3 +32,21 @@ def compute_alpha(a, b):
 def make_hard_decision(soft_input):
     """Makes hard decision based on soft input values (LLR)."""
     return np.array([s < 0 for s in soft_input], dtype=np.int8)
+
+
+@numba.njit
+def compute_left_alpha(llr):
+    """Compute Alpha for left node during SC-based decoding."""
+    N = llr.size // 2
+    left = llr[:N]
+    right = llr[N:]
+    return compute_alpha(left, right)
+
+
+@numba.njit
+def compute_right_alpha(llr, left_beta):
+    """Compute Alpha for right node during SC-based decoding."""
+    N = llr.size // 2
+    left = llr[:N]
+    right = llr[N:]
+    return right - (2 * left_beta - 1) * left

python_polar_coding/polar_codes/fast_ssc/decoder.py

@@ -2,6 +2,10 @@
 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 .node import FastSSCNode
 
@@ -18,12 +22,13 @@ def __init__(
             code_min_size: int = 0,
     ):
         super().__init__(n=n, mask=mask, is_systematic=is_systematic)
-        self._decoding_tree = self.node_class(
-            mask=self.mask,
-            N_min=code_min_size,
-        )
+        self._decoding_tree = self.setup_decoding_tree(code_min_size)
         self._position = 0
 
+    def setup_decoding_tree(self, N_min, **kwargs):
+        """Setup decoding tree."""
+        return self.node_class(mask=self.mask, N_min=N_min)
+
     def _set_initial_state(self, received_llr):
         """Initialize decoder with received message."""
         self.current_state = np.zeros(self.n, dtype=np.int8)
@@ -73,12 +78,12 @@ def compute_intermediate_alpha(self, leaf):
             parent_alpha = node.parent.alpha
 
             if node.is_left:
-                node.alpha = self._compute_left_alpha(parent_alpha)
+                node.alpha = compute_left_alpha(parent_alpha)
                 continue
 
             left_node = node.siblings[0]
             left_beta = left_node.beta
-            node.alpha = self._compute_right_alpha(parent_alpha, left_beta)
+            node.alpha = compute_right_alpha(parent_alpha, left_beta)
             node.is_computed = True
 
     def compute_intermediate_beta(self, node):

python_polar_coding/polar_codes/fast_ssc/node.py

@@ -138,7 +138,7 @@ 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.repetition_min_size and self._check_is_parity(self._mask):  # noqa
+        if self.N >= self.repetition_min_size and self._check_is_rep(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
@@ -153,7 +153,7 @@ def _check_is_zero(self, mask):
     def _check_is_spc(self, mask):
         return mask[0] == 0 and np.sum(mask) == mask.size - 1
 
-    def _check_is_parity(self, mask):
+    def _check_is_rep(self, mask):
         return mask[-1] == 1 and np.sum(mask) == 1
 
     def _build_decoding_tree(self):

python_polar_coding/polar_codes/g_fast_ssc/decoder.py

@@ -16,10 +16,14 @@ def __init__(
             code_min_size: int = 0,
             AF: int = 1,
     ):
-        super().__init__(n=n, mask=mask, is_systematic=is_systematic)
-        self._decoding_tree = self.node_class(
-            mask=self.mask,
-            N_min=code_min_size,
-            AF=AF,
+        self.AF = AF
+        super().__init__(
+            n=n,
+            mask=mask,
+            is_systematic=is_systematic,
+            code_min_size=code_min_size,
         )
-        self._position = 0
+
+    def setup_decoding_tree(self, N_min, **kwargs):
+        """Setup decoding tree."""
+        return self.node_class(mask=self.mask, N_min=N_min, AF=self.AF)

python_polar_coding/polar_codes/g_fast_ssc/functions.py

@@ -5,7 +5,7 @@
 from ..fast_ssc import compute_single_parity_check
 
 
-@numba.njit
+# @numba.njit
 def compute_g_repetition(llr, mask_steps, last_chunk_type, N):
     """Compute bits for Generalized Repetition node.
 

python_polar_coding/polar_codes/sc/decoder.py

@@ -88,13 +88,13 @@ def _compute_intermediate_alpha(self, position):
                 continue
 
             if self.current_state[i - 1] == 0:
-                self.intermediate_llr[i] = self._compute_left_alpha(llr)
+                self.intermediate_llr[i] = functions.compute_left_alpha(llr)
                 continue
 
             end = position
             start = end - np.power(2, self.n - i)
             left_bits = self.intermediate_bits[i][start: end]
-            self.intermediate_llr[i] = self._compute_right_alpha(llr, left_bits)
+            self.intermediate_llr[i] = functions.compute_right_alpha(llr, left_bits)
 
     def _compute_beta(self, position):
         """Make decision about current decoding value."""