Add Generalized Fast SCAN decoder

Author: grigory

.gitignore

@@ -5,6 +5,7 @@
 .idea
 .python-version
 .pytest_cache
+apak/*
 firestore/*credentials*
 src*
 tests/_trial_temp*

check.ipynb

@@ -1,5 +1,6 @@
 from .fast_scan import FastSCANCodec
 from .fast_ssc import FastSSCPolarCodec
+from .g_fast_scan import GFastSCANCodec
 from .g_fast_ssc import GeneralizedFastSSCPolarCodec
 from .rc_scan import RCSCANPolarCodec
 from .sc import SCPolarCodec

e_g_fast_ssc_analisys.ipynb

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

python_polar_coding/polar_codes/__init__.py

@@ -0,0 +1,38 @@
+from typing import Union
+
+from python_polar_coding.polar_codes.rc_scan import RCSCANPolarCodec
+
+from .decoder import GFastSCANDecoder
+
+
+class GFastSCANCodec(RCSCANPolarCodec):
+    decoder_class = GFastSCANDecoder
+
+    def __init__(
+            self,
+            N: int,
+            K: int,
+            design_snr: float = 0.0,
+            mask: Union[str, None] = None,
+            pcc_method: str = RCSCANPolarCodec.BHATTACHARYYA,
+            AF: int = 1,
+            I: int = 1,
+    ):
+
+        self.AF = AF
+        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)
+
+    def to_dict(self):
+        d = super().to_dict()
+        d.update({'AF': self.AF})
+        return d

python_polar_coding/polar_codes/g_fast_scan/__init__.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+from ..rc_scan import RCSCANDecoder
+from .node import GFastSCANNode
+
+
+class GFastSCANDecoder(RCSCANDecoder):
+
+    node_class = GFastSCANNode
+
+    def __init__(
+            self,
+            n: int,
+            mask: np.array,
+            AF: int = 1,
+            I: int = 1,
+    ):
+        self.AF = AF
+        super().__init__(n=n, mask=mask, I=I)
+
+    def setup_decoding_tree(self, N_min, **kwargs):
+        """Setup decoding tree."""
+        return self.node_class(mask=self.mask, AF=self.AF)

python_polar_coding/polar_codes/g_fast_scan/codec.py

@@ -0,0 +1,53 @@
+import numba
+import numpy as np
+
+from ..fast_scan import compute_spc_beta
+from ..rc_scan import compute_beta_one_node
+
+
+@numba.njit
+def compute_g_repetition(llr, mask_steps, last_chunk_type, N):
+    """Compute bits for Generalized Repetition node.
+
+    Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, A.
+
+    """
+    step = N // mask_steps  # step is equal to a chunk size
+
+    last_alpha = np.zeros(step)
+    for i in range(step):
+        last_alpha[i] = np.sum(np.array([
+            llr[i + j * step] for j in range(mask_steps)
+        ]))
+
+    last_beta = (
+        compute_beta_one_node(last_alpha) if last_chunk_type == 1
+        else compute_spc_beta(last_alpha)
+    )
+
+    result = np.zeros(N)
+    for i in range(0, N, step):
+        result[i: i + step] = last_beta
+
+    return result
+
+
+@numba.njit
+def compute_rg_parity(llr, mask_steps, N):
+    """Compute bits for Relaxed Generalized Parity Check node.
+
+    Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, B.
+
+    """
+    step = N // mask_steps  # step is equal to a chunk size
+    result = np.zeros(N)
+
+    for i in range(step):
+        alpha = np.zeros(mask_steps)
+        for j in range(mask_steps):
+            alpha[j] = llr[i + j * step]
+
+        beta = compute_spc_beta(alpha)
+        result[i:N:step] = beta
+
+    return result

python_polar_coding/polar_codes/g_fast_scan/decoder.py

@@ -0,0 +1,54 @@
+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)

python_polar_coding/polar_codes/g_fast_scan/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.