Refactor Fast-SSC codec

Author: grigory

python_polar_coding/polar_codes/base/__init__.py

@@ -1,6 +1,6 @@
 from .codec import BaseCRCPolarCodec, BasePolarCodec
 from .constants import *
-from .decoder import BaseDecoder
+from .decoder import BaseDecoder, BaseTreeDecoder
 from .decoding_path import DecodingPathMixin
 from .encoder import *
 from .functions import *

python_polar_coding/polar_codes/base/decoder.py

@@ -1,6 +1,7 @@
 import abc
 
 import numpy as np
+from anytree import PreOrderIter
 
 
 class BaseDecoder(metaclass=abc.ABCMeta):
@@ -14,13 +15,13 @@ def __init__(self, n, mask: np.array, is_systematic: bool = True):
 
     def decode(self, received_llr: np.array) -> np.array:
         decoded = self.decode_internal(received_llr)
-        return self.get_result(decoded)
+        return self.extract_result(decoded)
 
     @abc.abstractmethod
     def decode_internal(self, received_llr: np.array) -> np.array:
         """Implementation of particular decoding method."""
 
-    def get_result(self, decoded: np.array) -> np.array:
+    def extract_result(self, decoded: np.array) -> np.array:
         """Get decoding result.
 
         Extract info bits from decoded message due to polar code mask.
@@ -32,3 +33,97 @@ def get_result(self, decoded: np.array) -> np.array:
             if self.mask[i] == 1:
                 decoded_info = np.append(decoded_info, decoded[i])
         return np.array(decoded_info, dtype=np.int)
+
+
+class BaseTreeDecoder(metaclass=abc.ABCMeta):
+    """Basic class for polar decoder that use tree for decoding."""
+
+    node_class: 'BaseDecodingNode'
+
+    def __init__(self, n, mask: np.array):
+        self.N = mask.shape[0]
+        self.n = n
+        self.mask = mask
+
+        self._decoding_tree = self._setup_decoding_tree()
+        self._position = 0
+
+    def __call__(self, received_llr: np.array) -> np.array:
+        decoded = self.decode(received_llr)
+        return self.extract_result(decoded)
+
+    @property
+    def leaves(self):
+        return self._decoding_tree.leaves
+
+    @property
+    def root(self):
+        """Returns root node of decoding tree."""
+        return self._decoding_tree.root
+
+    @property
+    def result(self):
+        return self.root.beta
+
+    def decode(self, received_llr: np.array) -> np.array:
+        """Implementation of decoding using tree."""
+        self._set_initial_state(received_llr)
+        self._reset_tree_computed_state()
+
+        for leaf in self.leaves:
+            self._set_decoder_state(self._position)
+            self._compute_intermediate_alpha(leaf)
+            leaf()
+            self._compute_intermediate_beta(leaf)
+            self._set_next_state(leaf.N)
+
+        return self.result
+
+    def extract_result(self, decoded: np.array) -> np.array:
+        """Get decoding result.
+
+        Extract info bits from decoded message due to polar code mask.
+
+        """
+        decoded_info = list()
+
+        for i in range(self.N):
+            if self.mask[i] == 1:
+                decoded_info = np.append(decoded_info, decoded[i])
+        return np.array(decoded_info, dtype=np.int)
+
+    def _setup_decoding_tree(self, ):
+        """Setup decoding tree."""
+        return self.node_class(mask=self.mask)
+
+    def _set_initial_state(self, received_llr):
+        """Initialize decoder with received message."""
+        self.current_state = np.zeros(self.n, dtype=np.int8)
+        self.previous_state = np.ones(self.n, dtype=np.int8)
+
+        # LLR values at intermediate steps
+        self._position = 0
+        self._decoding_tree.root.alpha = received_llr
+
+    def _reset_tree_computed_state(self):
+        """Reset the state of the tree before decoding"""
+        for node in PreOrderIter(self._decoding_tree):
+            node.is_computed = False
+
+    def _set_decoder_state(self, position):
+        """Set current state of the decoder."""
+        bits = np.unpackbits(
+            np.array([position], dtype=np.uint32).byteswap().view(np.uint8)
+        )
+        self.current_state = bits[-self.n:]
+
+    @abc.abstractmethod
+    def _compute_intermediate_alpha(self, leaf):
+        """Compute intermediate Alpha values (LLR)."""
+
+    @abc.abstractmethod
+    def _compute_intermediate_beta(self, node):
+        """Compute intermediate Beta values (Bits or LLR)."""
+
+    def _set_next_state(self, leaf_size):
+        self._position += leaf_size

python_polar_coding/polar_codes/base/functions/__init__.py

@@ -5,7 +5,11 @@
     function_1,
     function_2,
 )
-from .beta_hard import compute_beta_hard, compute_parent_beta_hard
+from .beta_hard import (
+    compute_beta_hard,
+    compute_parent_beta_hard,
+    make_hard_decision,
+)
 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/fast_ssc/codec.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from python_polar_coding.polar_codes.base import BasePolarCodec
 
 from .decoder import FastSSCDecoder
@@ -12,9 +14,8 @@ class FastSSCPolarCodec(BasePolarCodec):
     decoder_class = FastSSCDecoder
 
     def init_decoder(self):
-        return self.decoder_class(n=self.n, mask=self.mask,
-                                  is_systematic=self.is_systematic)
+        return self.decoder_class(n=self.n, mask=self.mask)
 
-    @property
-    def tree(self):
-        return self.decoder._decoding_tree
+    def decode(self, received_message: np.array) -> np.array:
+        """Decode received message presented as LLR values."""
+        return self.decoder(received_message)

python_polar_coding/polar_codes/fast_ssc/decoder.py

@@ -1,72 +1,18 @@
-import numpy as np
-from anytree import PreOrderIter
-
+from python_polar_coding.polar_codes.base import BaseTreeDecoder
 from python_polar_coding.polar_codes.base.functions import (
     compute_left_alpha,
-    compute_right_alpha,
     compute_parent_beta_hard,
+    compute_right_alpha,
 )
-from python_polar_coding.polar_codes.sc import SCDecoder
 
 from .node import FastSSCNode
 
 
-class FastSSCDecoder(SCDecoder):
+class FastSSCDecoder(BaseTreeDecoder):
     """Implements Fast SSC decoding algorithm."""
     node_class = FastSSCNode
 
-    def __init__(
-            self,
-            n: int,
-            mask: np.array,
-            is_systematic: bool = True,
-            code_min_size: int = 0,
-    ):
-        super().__init__(n=n, mask=mask, is_systematic=is_systematic)
-        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)
-        self.previous_state = np.ones(self.n, dtype=np.int8)
-
-        # LLR values at intermediate steps
-        self._position = 0
-        self._decoding_tree.root.alpha = received_llr
-
-    def decode_internal(self, received_llr: np.array) -> np.array:
-        """Implementation of SC decoding method."""
-        self._set_initial_state(received_llr)
-
-        # Reset the state of the tree before decoding
-        for node in PreOrderIter(self._decoding_tree):
-            node.is_computed = False
-
-        for leaf in self._decoding_tree.leaves:
-            self._set_decoder_state(self._position)
-            self.compute_intermediate_alpha(leaf)
-            leaf()
-            self.compute_intermediate_beta(leaf)
-            self.set_next_state(leaf.N)
-
-        return self.result
-
-    @property
-    def root(self):
-        """Returns root node of decoding tree."""
-        return self._decoding_tree.root
-
-    @property
-    def result(self):
-        if self.is_systematic:
-            return self.root.beta
-
-    def compute_intermediate_alpha(self, leaf):
+    def _compute_intermediate_alpha(self, leaf):
         """Compute intermediate Alpha values (LLR)."""
         for node in leaf.path[1:]:
             if node.is_computed:
@@ -87,7 +33,7 @@ def compute_intermediate_alpha(self, leaf):
             node.alpha = compute_right_alpha(parent_alpha, left_beta)
             node.is_computed = True
 
-    def compute_intermediate_beta(self, node):
+    def _compute_intermediate_beta(self, node):
         """Compute intermediate Beta values (BIT)."""
         if node.is_left:
             return
@@ -98,7 +44,4 @@ def compute_intermediate_beta(self, node):
         parent = node.parent
         left = node.siblings[0]
         parent.beta = compute_parent_beta_hard(left.beta, node.beta)
-        return self.compute_intermediate_beta(parent)
-
-    def set_next_state(self, leaf_size):
-        self._position += leaf_size
+        return self._compute_intermediate_beta(parent)

python_polar_coding/tests/test_base/test_codec.py

@@ -128,7 +128,7 @@ def test_precode_and_extract(self):
         precoded = self.systematic_code.encoder._precode(self.message)
         self.assertEqual(precoded.size, self.codeword_length)
 
-        extracted = self.systematic_code.decoder.get_result(precoded)
+        extracted = self.systematic_code.decoder.extract_result(precoded)
         self.assertEqual(extracted.size, self.info_length)
 
         self.assertTrue(all(extracted == self.message))
@@ -143,13 +143,13 @@ def test_systematic_encode(self):
         encoded = self.systematic_code.encode(self.message)
         self.assertTrue(all(encoded == self.sys_enc_msg))
 
-        extracted = self.systematic_code.decoder.get_result(encoded)
+        extracted = self.systematic_code.decoder.extract_result(encoded)
         self.assertTrue(all(extracted == self.message))
 
     def test_systematic_encode_with_crc(self):
         """Test for systematic encoding with CRC support"""
         encoded = self.systematic_crc_code.encode(self.message)
         self.assertTrue(all(encoded == self.sys_crc_enc_msg))
 
-        extracted = self.systematic_crc_code.decoder.get_result(encoded)
+        extracted = self.systematic_crc_code.decoder.extract_result(encoded)
         self.assertTrue(all(extracted[:self.info_length] == self.message))

python_polar_coding/tests/test_fast_ssc/test_decoder.py

@@ -23,7 +23,7 @@ def setUpClass(cls):
 
     def test_zero_node_decoder(self):
         mask = np.zeros(self.length, dtype=np.int8)
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=self.n)
+        decoder = FastSSCDecoder(mask=mask, n=self.n)
         self.assertEqual(len(decoder._decoding_tree.leaves), 1)
 
         decoder.decode(self.received_llr)
@@ -34,7 +34,7 @@ def test_zero_node_decoder(self):
 
     def test_one_node_decoder(self):
         mask = np.ones(self.length, dtype=np.int8)
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=self.n)
+        decoder = FastSSCDecoder(mask=mask, n=self.n)
         self.assertEqual(len(decoder._decoding_tree.leaves), 1)
 
         decoder.decode(self.received_llr)
@@ -45,7 +45,7 @@ def test_one_node_decoder(self):
 
     def test_spc_node_decoder(self):
         mask = np.array([0, 1, 1, 1, 1, 1, 1, 1], dtype=np.int8)
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=self.n)
+        decoder = FastSSCDecoder(mask=mask, n=self.n)
         self.assertEqual(len(decoder._decoding_tree.leaves), 1)
 
         decoder.decode(self.received_llr)
@@ -56,7 +56,7 @@ def test_spc_node_decoder(self):
 
     def test_repetition_node_decoder(self):
         mask = np.array([0, 0, 0, 0, 0, 0, 0, 1], dtype=np.int8)
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=self.n)
+        decoder = FastSSCDecoder(mask=mask, n=self.n)
         self.assertEqual(len(decoder._decoding_tree.leaves), 1)
 
         decoder.decode(self.received_llr)
@@ -67,7 +67,7 @@ def test_repetition_node_decoder(self):
 
     def test_repetition_spc_node_decoder(self):
         mask = np.array([0, 0, 0, 1, 0, 1, 1, 1], dtype=np.int8)
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=self.n)
+        decoder = FastSSCDecoder(mask=mask, n=self.n)
         self.assertEqual(len(decoder._decoding_tree.leaves), 2)
 
         decoder.decode(self.received_llr)
@@ -103,7 +103,7 @@ def test_repetition_spc_node_decoder(self):
 
     def test_spc_repetition_node_decoder(self):
         mask = np.array([0, 1, 1, 1, 0, 0, 0, 1], dtype=np.int8)
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=self.n)
+        decoder = FastSSCDecoder(mask=mask, n=self.n)
         self.assertEqual(len(decoder._decoding_tree.leaves), 2)
 
         decoder.decode(self.received_llr)
@@ -154,7 +154,7 @@ def test_complex(self):
             np.array([0, ], dtype=np.int8),
             np.array([0, 1, 1, 1, ], dtype=np.int8),
         ]
-        decoder = FastSSCDecoder(mask=mask, is_systematic=True, n=4)
+        decoder = FastSSCDecoder(mask=mask, n=4)
 
         # Check tree structure
         self.assertEqual(len(decoder._decoding_tree.leaves), len(sub_codes))