Refactor G-Fast-SCAN codec

Author: grigory

python_polar_coding/polar_codes/fast_scan/node.py

@@ -10,11 +10,3 @@ class FastSCANNode(RCSCANNode):
         NodeTypes.REPETITION,
         NodeTypes.SINGLE_PARITY_CHECK,
     )
-
-    @property
-    def is_repetition(self):
-        return self.node_type == NodeTypes.REPETITION
-
-    @property
-    def is_parity(self):
-        return self.node_type == NodeTypes.SINGLE_PARITY_CHECK

python_polar_coding/polar_codes/g_fast_scan/__init__.py

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

python_polar_coding/polar_codes/g_fast_scan/codec.py

@@ -15,23 +15,28 @@ def __init__(
             design_snr: float = 0.0,
             mask: Union[str, None] = None,
             pcc_method: str = RCSCANPolarCodec.BHATTACHARYYA,
-            AF: int = 1,
+            AF: int = 0,
             I: int = 1,
             * args, **kwargs,
     ):
 
         self.AF = AF
-        super().__init__(N=N, K=K,
-                         design_snr=design_snr,
-                         mask=mask,
-                         pcc_method=pcc_method,
-                         I=I,)
+        super().__init__(
+            N=N,
+            K=K,
+            design_snr=design_snr,
+            mask=mask,
+            pcc_method=pcc_method,
+            I=I,
+        )
 
     def init_decoder(self):
-        return self.decoder_class(n=self.n,
-                                  mask=self.mask,
-                                  AF=self.AF,
-                                  I=self.I)
+        return self.decoder_class(
+            n=self.n,
+            mask=self.mask,
+            AF=self.AF,
+            I=self.I,
+        )
 
     def to_dict(self):
         d = super().to_dict()

python_polar_coding/polar_codes/g_fast_scan/decoder.py

@@ -18,6 +18,6 @@ def __init__(
         self.AF = AF
         super().__init__(n=n, mask=mask, I=I)
 
-    def setup_decoding_tree(self, N_min, **kwargs):
+    def _setup_decoding_tree(self):
         """Setup decoding tree."""
         return self.node_class(mask=self.mask, AF=self.AF)

python_polar_coding/polar_codes/g_fast_scan/functions.py

@@ -1,54 +1,24 @@
-from python_polar_coding.polar_codes.fast_scan import (
-    compute_repetition_beta,
-    compute_spc_beta,
-)
-from python_polar_coding.polar_codes.g_fast_ssc import GeneralizedFastSSCNode
-from python_polar_coding.polar_codes.rc_scan import (
-    compute_beta_one_node,
-    compute_beta_zero_node,
-)
-
-from .functions import compute_g_repetition, compute_rg_parity
-
-
-class GFastSCANNode(GeneralizedFastSSCNode):
-
-    def compute_leaf_beta(self):
-        """Do nothing for ZERO and ONE nodes.
-
-        Unlike SC-based decoders SCAN decoders does not make decisions
-        in ZERO and ONE leaves.
-
-        """
-        if self.is_repetition:
-            self.beta = compute_repetition_beta(self.alpha)
-        if self.is_parity:
-            self.beta = compute_spc_beta(self.alpha)
-        if self.is_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.is_rg_parity:
-            self.beta = compute_rg_parity(
-                llr=self.alpha,
-                mask_steps=self.mask_steps,
-                N=self.N,
-            )
-
-    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_zero:
-            self._beta = compute_beta_zero_node(self.alpha)
-        if self.is_one:
-            self._beta = compute_beta_one_node(self.alpha)
+from typing import Dict
+
+from python_polar_coding.polar_codes.fast_scan import FastSCANNode
+
+from ..base import NodeTypes
+
+
+class GFastSCANNode(FastSCANNode):
+    supported_nodes = (
+        NodeTypes.ZERO,
+        NodeTypes.ONE,
+        NodeTypes.REPETITION,
+        NodeTypes.SINGLE_PARITY_CHECK,
+        NodeTypes.G_REPETITION,
+        NodeTypes.RG_PARITY,
+    )
+
+    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,
+        )