Merge branch 'feature/extended-g-fast-ssc-codec'

Author: grigory

e_g_fast_ssc_analisys.ipynb

@@ -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/__init__.py

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

python_polar_coding/polar_codes/base/constants.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/base/functions.py

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

python_polar_coding/polar_codes/e_g_fast_ssc/__init__.py

@@ -0,0 +1,10 @@
+from python_polar_coding.polar_codes.g_fast_ssc import (
+    GeneralizedFastSSCPolarCodec,
+)
+
+from .decoder import EGFastSSCDecoder
+
+
+class EGFastSSCPolarCodec(GeneralizedFastSSCPolarCodec):
+    """Extended Generalized Fast SSC codec."""
+    decoder_class = EGFastSSCDecoder

python_polar_coding/polar_codes/e_g_fast_ssc/codec.py

@@ -0,0 +1,10 @@
+from python_polar_coding.polar_codes.g_fast_ssc import (
+    GeneralizedFastSSCDecoder,
+)
+
+from .node import EGFastSSCNode
+
+
+class EGFastSSCDecoder(GeneralizedFastSSCDecoder):
+    """"""
+    node_class = EGFastSSCNode

python_polar_coding/polar_codes/e_g_fast_ssc/decoder.py

@@ -0,0 +1,20 @@
+import numba
+import numpy as np
+
+from ..base import compute_left_alpha
+
+
+@numba.njit
+def compute_left_alpha_sign(alpha: np.array) -> np.array:
+    """"""
+    left_alpha = compute_left_alpha(alpha)
+    return np.sign(np.sum(left_alpha))
+
+
+@numba.njit
+def compute_right_alpha(alpha: np.array, left_sign: int = 1) -> np.array:
+    """`left_sign` is 1 or -1"""
+    N = alpha.size // 2
+    left_alpha = alpha[:N]
+    right_alpha = alpha[N:]
+    return right_alpha + left_alpha * left_sign

python_polar_coding/polar_codes/e_g_fast_ssc/functions.py

@@ -0,0 +1,117 @@
+import numpy as np
+
+from python_polar_coding.polar_codes.g_fast_ssc import GeneralizedFastSSCNode
+
+from .functions import compute_left_alpha_sign, compute_right_alpha
+
+
+class EGFastSSCNode(GeneralizedFastSSCNode):
+    """Decoder for Generalized Fast SSC code.
+
+    Based on: https://arxiv.org/pdf/1804.09508.pdf
+
+    """
+    ZERO_ANY = 'ZERO-ANY'
+    REP_ANY = 'REP-ANY'
+
+    def __init__(self, *args, **kwargs):
+        # Contains `ANY` node for ZERO_ANY or REP_ANY
+        self.inner_node = None
+        super().__init__(*args, **kwargs)
+
+    @property
+    def is_any(self):
+        return (
+            self.is_zero or
+            self.is_one or
+            self.is_repetition or
+            self.is_parity or
+            self.is_g_repetition or
+            self.is_rg_parity
+        )
+
+    @property
+    def is_zero_any(self):
+        return self._node_type == self.ZERO_ANY
+
+    @property
+    def is_rep_any(self):
+        return self._node_type == self.REP_ANY
+
+    def get_node_type(self):
+        ntype = super().get_node_type()
+        if ntype != self.OTHER:
+            return ntype
+        if self._check_is_zero_any(self._mask):
+            return self.ZERO_ANY
+        if self._check_is_rep_any(self._mask):
+            return self.REP_ANY
+        return self.OTHER
+
+    def _check_is_zero_any(self, mask):
+        """"""
+        left, right = np.split(mask, 2)
+        if not self._check_is_zero(left):
+            return False
+        inner_node = self.__class__(
+            mask=right,
+            name=self.ROOT,
+            N_min=self.N_min,
+            AF=self.AF
+        )
+        if not inner_node.is_any:
+            return False
+
+        self.inner_node = inner_node
+        return True
+
+    def _check_is_rep_any(self, mask):
+        """"""
+        left, right = np.split(mask, 2)
+        if not self._check_is_rep(left):
+            return False
+        right_node = self.__class__(
+            mask=right,
+            name=self.ROOT,
+            N_min=self.N_min,
+            AF=self.AF
+        )
+        if not right_node.is_any:
+            return False
+
+        self.inner_node = right_node
+        return True
+
+    def compute_leaf_beta(self):
+        super().compute_leaf_beta()
+        klass = self.__class__
+
+        if self._node_type == klass.ZERO_ANY:
+            self._beta = self.compute_zero_any()
+        if self._node_type == klass.REP_ANY:
+            self._beta = self.compute_rep_any()
+
+    def compute_zero_any(self):
+        """"""
+        right_alpha = compute_right_alpha(self.alpha, left_sign=1)
+
+        self.inner_node.alpha = right_alpha
+        self.inner_node.compute_leaf_beta()
+
+        beta = np.zeros(self.N, dtype=np.int8)
+        beta[:self.inner_node.N] = self.inner_node.beta
+        beta[self.inner_node.N:] = self.inner_node.beta
+        return beta
+
+    def compute_rep_any(self):
+        """"""
+        left_sign = compute_left_alpha_sign(self.alpha)
+        right_alpha = compute_right_alpha(self.alpha, left_sign)
+
+        self.inner_node.alpha = right_alpha
+        self.inner_node.compute_leaf_beta()
+
+        beta = np.zeros(self.N, dtype=np.int8)
+        beta[:self.inner_node.N] = self.inner_node.beta
+        beta[self.inner_node.N:] = self.inner_node.beta
+        return beta

python_polar_coding/polar_codes/e_g_fast_ssc/node.py

@@ -1,6 +1,10 @@
 import numpy as np
 from anytree import PreOrderIter
 
+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
@@ -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)
@@ -70,15 +75,19 @@ 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:
-                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):