Refactor Fast SSC node

Author: grigory

.isort.cfg

@@ -3,3 +3,4 @@ line_length=79
 length_sort_stdlib=1
 multi_line_output=3
 include_trailing_comma=True
+skip=apak,venv

python_polar_coding/__init__.py

@@ -1,7 +0,0 @@
-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
-from .sc_list import SCListPolarCodec

python_polar_coding/polar_codes/__init__.py

@@ -4,3 +4,4 @@
 from .decoding_path import DecodingPathMixin
 from .encoder import *
 from .functions import *
+from .node import *

python_polar_coding/polar_codes/base/__init__.py

@@ -0,0 +1,11 @@
+from .alpha import (
+    compute_alpha,
+    compute_left_alpha,
+    compute_right_alpha,
+    function_1,
+    function_2,
+)
+from .beta_hard import compute_beta_hard, compute_parent_beta_hard
+from .beta_soft import compute_beta_soft
+from .encoding import compute_encoding_step
+from .node_types import NodeTypes, get_node_type

python_polar_coding/polar_codes/base/functions.py

@@ -0,0 +1,53 @@
+import numba
+import numpy as np
+
+
+@numba.njit
+def compute_alpha(a: np.array, b: np.array) -> np.array:
+    """Basic function to compute intermediate LLR values."""
+    c = np.zeros(a.shape[0])
+    for i in range(c.shape[0]):
+        c[i] = (
+            np.sign(a[i]) *
+            np.sign(b[i]) *
+            np.fabs(np.array([a[i], b[i]])).min()
+        )
+    return c
+
+
+@numba.njit
+def compute_left_alpha(llr: np.array) -> np.array:
+    """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: np.array, left_beta: np.array) -> np.array:
+    """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
+
+
+@numba.njit
+def function_1(a: np.array, b: np.array, c: np.array) -> np.array:
+    """Function 1.
+
+    Source: doi:10.1007/s12243-018-0634-7, formula 1.
+
+    """
+    return compute_alpha(a, b + c)
+
+
+@numba.njit
+def function_2(a: np.array, b: np.array, c: np.array) -> np.array:
+    """Function 2.
+
+    Source: doi:10.1007/s12243-018-0634-7, formula 2.
+
+    """
+    return compute_alpha(a, b) + c

python_polar_coding/polar_codes/base/functions/__init__.py

@@ -0,0 +1,156 @@
+"""Common functions for polar coding."""
+import numba
+import numpy as np
+
+from .node_types import NodeTypes
+
+# -----------------------------------------------------------------------------
+# Making hard decisions during the decoding
+# -----------------------------------------------------------------------------
+
+
+@numba.njit
+def zero(
+        llr: np.array,
+        mask_steps: int = 0,
+        last_chunk_type: int = 0,
+) -> np.array:
+    """Makes hard decision based on soft input values (LLR)."""
+    return np.zeros(llr.size, dtype=np.int8)
+
+
+@numba.njit
+def make_hard_decision(
+        llr: np.array,
+        mask_steps: int = 0,
+        last_chunk_type: int = 0,
+) -> np.array:
+    """Makes hard decision based on soft input values (LLR)."""
+    return np.array([s < 0 for s in llr], dtype=np.int8)
+
+
+@numba.njit
+def single_parity_check(
+        llr: np.array,
+        mask_steps: int = 0,
+        last_chunk_type: int = 0,
+) -> np.array:
+    """Compute bits for Single Parity Check node.
+
+    Based on: https://arxiv.org/pdf/1307.7154.pdf, Section IV, A.
+
+    """
+    bits = make_hard_decision(llr)
+    parity = np.sum(bits) % 2
+    arg_min = np.abs(llr).argmin()
+    bits[arg_min] = (bits[arg_min] + parity) % 2
+    return bits
+
+
+@numba.njit
+def repetition(
+        llr: np.array,
+        mask_steps: int = 0,
+        last_chunk_type: int = 0,
+) -> np.array:
+    """Compute bits for Repetition node.
+
+    Based on: https://arxiv.org/pdf/1307.7154.pdf, Section IV, B.
+
+    """
+    return (
+        np.zeros(llr.size, dtype=np.int8) if np.sum(llr) >= 0
+        else np.ones(llr.size, dtype=np.int8)
+    )
+
+
+@numba.njit
+def g_repetition(
+        llr: np.array,
+        mask_steps: int,
+        last_chunk_type: int,
+) -> np.array:
+    """Compute bits for Generalized Repetition node.
+
+    Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, A.
+
+    """
+    N = llr.size
+    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 = (
+        make_hard_decision(last_alpha) if last_chunk_type == 1
+        else single_parity_check(last_alpha)
+    )
+
+    result = np.zeros(N)
+    for i in range(0, N, step):
+        result[i: i + step] = last_beta
+
+    return result
+
+
+@numba.njit
+def rg_parity(
+        llr: np.array,
+        mask_steps: int,
+        last_chunk_type: int = 0,
+) -> np.array:
+    """Compute bits for Relaxed Generalized Parity Check node.
+
+    Based on: https://arxiv.org/pdf/1804.09508.pdf, Section III, B.
+
+    """
+    N = llr.size
+    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 = single_parity_check(alpha)
+        result[i:N:step] = beta
+
+    return result
+
+
+# Mapping between decoding node types and corresponding decoding methods
+_methods_map = {
+    NodeTypes.ZERO: zero,
+    NodeTypes.ONE: make_hard_decision,
+    NodeTypes.SINGLE_PARITY_CHECK: single_parity_check,
+    NodeTypes.REPETITION: repetition,
+    NodeTypes.RG_PARITY: rg_parity,
+    NodeTypes.G_REPETITION: g_repetition,
+}
+
+
+def compute_beta_hard(
+        node_type: str,
+        llr: np.array,
+        mask_steps: int = 0,
+        last_chunk_type: int = 0,
+        *args, **kwargs,
+) -> np.array:
+    """Unites functions for making hard decisions during decoding."""
+    method = _methods_map[node_type]
+    return method(llr, mask_steps, last_chunk_type, *args, **kwargs)
+
+
+@numba.njit
+def compute_parent_beta_hard(left: np.array, right: np.array) -> np.array:
+    """Compute Beta values for parent Node."""
+    N = left.size
+    result = np.zeros(N * 2, dtype=np.int8)
+    result[:N] = (left + right) % 2
+    result[N:] = right
+
+    return result