Refactor G-Fast SSC codec

Author: grigory

.gitignore

@@ -8,4 +8,5 @@
 apak/*
 firestore/*credentials*
 src*
-tests/_trial_temp*
+tests/_trial_temp*
+venv/

python_polar_coding/polar_codes/base/functions/node_types.py

@@ -24,6 +24,10 @@ class NodeTypeDetector:
     # Minimal number of chunks in generalized nodes
     MIN_CHUNKS = 2
 
+    def __init__(self, *args, **kwargs):
+        self.last_chunk_type = None
+        self.mask_steps = None
+
     def __call__(
             self,
             supported_nodes: list,

python_polar_coding/polar_codes/base/node.py

@@ -41,6 +41,10 @@ def __init__(self, mask: np.array, name: str = ROOT, AF: int = 0, **kwargs):  #
         self._alpha = np.zeros(self.N, dtype=np.double)
         self._beta = np.zeros(self.N, dtype=np.int8)
 
+        # For generalized decoders
+        self.last_chunk_type = get_node_type.last_chunk_type
+        self.mask_steps = get_node_type.mask_steps
+
         self.build_decoding_tree()
 
     def __str__(self):

python_polar_coding/polar_codes/g_fast_ssc/__init__.py

@@ -1,4 +1,3 @@
 from .codec import GeneralizedFastSSCPolarCodec
 from .decoder import GeneralizedFastSSCDecoder
-from .functions import *
-from .node import GeneralizedFastSSCNode
+from .node import GFastSSCNode

python_polar_coding/polar_codes/g_fast_ssc/codec.py

@@ -18,26 +18,22 @@ def __init__(
             N: int,
             K: int,
             design_snr: float = 0.0,
-            is_systematic: bool = True,
             mask: Union[str, None] = None,
             pcc_method: str = FastSSCPolarCodec.BHATTACHARYYA,
-            Ns: int = 1,
-            AF: int = 1,
+            AF: int = 0,
     ):
 
-        self.Ns = Ns
         self.AF = AF
-        super().__init__(N=N, K=K,
-                         is_systematic=is_systematic,
-                         design_snr=design_snr,
-                         mask=mask,
-                         pcc_method=pcc_method)
+        super().__init__(
+            N=N,
+            K=K,
+            design_snr=design_snr,
+            mask=mask,
+            pcc_method=pcc_method,
+        )
 
     def init_decoder(self):
-        return self.decoder_class(n=self.n, mask=self.mask,
-                                  is_systematic=self.is_systematic,
-                                  code_min_size=self.Ns,
-                                  AF=self.AF)
+        return self.decoder_class(n=self.n, mask=self.mask, AF=self.AF)
 
     def to_dict(self):
         d = super().to_dict()

python_polar_coding/polar_codes/g_fast_ssc/decoder.py

@@ -2,28 +2,16 @@
 
 from python_polar_coding.polar_codes.fast_ssc import FastSSCDecoder
 
-from .node import GeneralizedFastSSCNode
+from .node import GFastSSCNode
 
 
 class GeneralizedFastSSCDecoder(FastSSCDecoder):
-    node_class = GeneralizedFastSSCNode
+    node_class = GFastSSCNode
 
-    def __init__(
-            self,
-            n: int,
-            mask: np.array,
-            is_systematic: bool = True,
-            code_min_size: int = 0,
-            AF: int = 1,
-    ):
+    def __init__(self, n: int, mask: np.array, AF: int = 1):
         self.AF = AF
-        super().__init__(
-            n=n,
-            mask=mask,
-            is_systematic=is_systematic,
-            code_min_size=code_min_size,
-        )
+        super().__init__(n=n, mask=mask)
 
-    def setup_decoding_tree(self, N_min, **kwargs):
+    def _setup_decoding_tree(self):
         """Setup decoding tree."""
-        return self.node_class(mask=self.mask, N_min=N_min, AF=self.AF)
+        return self.node_class(mask=self.mask, AF=self.AF)

python_polar_coding/polar_codes/g_fast_ssc/functions.py

@@ -1,132 +1,29 @@
-import numpy as np
+from typing import Dict
 
 from python_polar_coding.polar_codes.fast_ssc import FastSSCNode
-from python_polar_coding.polar_codes.utils import splits
 
-from .functions import compute_g_repetition, compute_rg_parity
+from ..base import NodeTypes
 
 
-class GeneralizedFastSSCNode(FastSSCNode):
+class GFastSSCNode(FastSSCNode):
     """Decoder for Generalized Fast SSC code.
 
     Based on: https://arxiv.org/pdf/1804.09508.pdf
 
     """
-    G_REPETITION = 'G-REPETITION'
-    RG_PARITY = 'RG-PARITY'
-
-    MIN_CHUNKS = 2
-
-    def __init__(self, AF=1, *args, **kwargs):
-        self.AF = AF
-        self.last_chunk_type = None
-        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:
-            return ntype
-        if self._check_is_g_repetition(self._mask):
-            return self.G_REPETITION
-        if self._check_is_rg_parity(self._mask):
-            return self.RG_PARITY
-        return self.OTHER
-
-    def _build_decoding_tree(self):
-        """Build Generalized Fast SSC decoding tree."""
-        if self.is_simplified_node:
-            return
-
-        if self._mask.size == self.M:
-            return
-
-        left_mask, right_mask = np.split(self._mask, 2)
-        cls = self.__class__
-        cls(mask=left_mask, name=self.LEFT, N_min=self.M, parent=self, AF=self.AF)
-        cls(mask=right_mask, name=self.RIGHT, N_min=self.M, parent=self, AF=self.AF)
-
-    def _check_is_g_repetition(self, mask):
-        """Check the node is Generalized Repetition node.
-
-        Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, A.
-
-        """
-        # 1. Split mask into T chunks, T in range [2, 4, ..., N/2]
-        for t in splits(self.MIN_CHUNKS, self.N // 2):
-            chunks = np.split(mask, t)
-
-            last = chunks[-1]
-            last_ok = (
-                (self._check_is_spc(last) and last.size >= self.REPETITION_MIN_SIZE)
-                or self._check_is_one(last)
-            )
-            if not last_ok:
-                continue
-
-            others_ok = all(self._check_is_zero(c) for c in chunks[:-1])
-            if not others_ok:
-                continue
-
-            self.last_chunk_type = 1 if self._check_is_one(last) else 0
-            self.mask_steps = t
-            return True
-
-        return False
-
-    def _check_is_rg_parity(self, mask):
-        """Check the node is Relaxed Generalized Parity Check node.
-
-        Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, B.
-
-        """
-        # 1. Split mask into T chunks, T in range [2, 4, ..., N/2]
-        for t in splits(self.MIN_CHUNKS, self.N // 2):
-            chunks = np.split(mask, t)
-
-            first = chunks[0]
-            if not self._check_is_zero(first):
-                continue
-
-            ones = 0
-            spcs = 0
-            for c in chunks[1:]:
-                if self._check_is_one(c):
-                    ones += 1
-                elif c.size >= self.SPC_MIN_SIZE and self._check_is_spc(c):
-                    spcs += 1
-
-            others_ok = (ones + spcs + 1) == t and spcs <= self.AF
-            if not others_ok:
-                continue
-
-            self.mask_steps = t
-            return True
-
-        return False
-
-    def compute_leaf_beta(self):
-        super().compute_leaf_beta()
-        klass = self.__class__
-
-        if self._node_type == klass.G_REPETITION:
-            self._beta = compute_g_repetition(
-                llr=self.alpha,
-                mask_steps=self.mask_steps,
-                last_chunk_type=self.last_chunk_type,
-                N=self.N,
-            )
-        if self._node_type == klass.RG_PARITY:
-            self._beta = compute_rg_parity(
-                llr=self.alpha,
-                mask_steps=self.mask_steps,
-                N=self.N,
-            )
+    supported_nodes = (
+        NodeTypes.ZERO,
+        NodeTypes.ONE,
+        NodeTypes.SINGLE_PARITY_CHECK,
+        NodeTypes.REPETITION,
+        NodeTypes.RG_PARITY,
+        NodeTypes.G_REPETITION,
+    )
+
+    def get_decoding_params(self) -> Dict:
+        return dict(
+            node_type=self.node_type,
+            llr=self.alpha,
+            mask_steps=self.mask_steps,
+            last_chunk_type=self.last_chunk_type,
+        )