Reorganize code of SC Polar codec

grigory · grigory · commit a7f74db4409d · 2020-04-14T10:08:36.000+01:00
diff --git a/python_polar_coding/polar_codes/base/__init__.py b/python_polar_coding/polar_codes/base/__init__.py
@@ -1,4 +1,4 @@
-from .codec import BasePolarCodec, BasePolarCodeWithCRC
+from .codec import BaseCRCPolarCodec, BasePolarCodec
 from .decoder import BaseDecoder
-from .functions import *
 from .decoding_path import DecodingPathMixin
+from .functions import *
diff --git a/python_polar_coding/polar_codes/base/codec.py b/python_polar_coding/polar_codes/base/codec.py
@@ -1,11 +1,12 @@
 import abc
-from operator import itemgetter
 from typing import Union
+from operator import itemgetter
 
 import numpy as np
 
+from python_polar_coding.polar_codes import crc, pcc, utils
+
 from . import encoder
-from python_polar_coding.polar_codes import utils, crc, pcc
 
 
 class BasePolarCodec(metaclass=abc.ABCMeta):
@@ -133,7 +134,7 @@ def _construct_polar_mask(self, K):
         return np.array([m[2] for m in mask])
 
 
-class BasePolarCodeWithCRC(BasePolarCodec):
+class BaseCRCPolarCodec(BasePolarCodec):
     """Basic Polar code class with CRC support.
 
     Provides the support of CRC 16 and CRC 32.
diff --git a/python_polar_coding/polar_codes/sc/__init__.py b/python_polar_coding/polar_codes/sc/__init__.py
@@ -0,0 +1,2 @@
+from .codec import SCPolarCodec
+from .decoder import SCDecoder
diff --git a/python_polar_coding/polar_codes/sc/codec.py b/python_polar_coding/polar_codes/sc/codec.py
@@ -1,12 +1,12 @@
-from .base.polar_code import BasicPolarCode
-from .decoders import SCDecoder
+from ..base import BasePolarCodec
+from .decoder import SCDecoder
 
 
-class SCPolarCode(BasicPolarCode):
+class SCPolarCodec(BasePolarCodec):
     """Polar code with SC decoding algorithm."""
     decoder_class = SCDecoder
 
-    def get_decoder(self):
+    def init_decoder(self):
         return self.decoder_class(
             n=self.n, mask=self.mask, is_systematic=self.is_systematic
         )
diff --git a/python_polar_coding/polar_codes/sc/decoder.py b/python_polar_coding/polar_codes/sc/decoder.py
@@ -1,11 +1,10 @@
 import numba
 import numpy as np
 
-from ..base.decoder import BaseDecoder
-from ..base.functions import compute_encoding_step
+from ..base import decoder, functions
 
 
-class SCDecoder(BaseDecoder):
+class SCDecoder(decoder.BaseDecoder):
     """Implements SC decoding algorithm.
 
     Stores initial and intermediate LLR values, intermediate bit values and
@@ -134,7 +133,7 @@ def _compute_intermediate_beta(self, position):
             source = self.intermediate_bits[i + 1]
             result = self.intermediate_bits[i]
 
-            self.intermediate_bits[i] = compute_encoding_step(
+            self.intermediate_bits[i] = functions.compute_encoding_step(
                 i, self.n, source, result
             )
 
diff --git a/python_polar_coding/tests/mixins.py b/python_polar_coding/tests/mixins.py
diff --git a/python_polar_coding/tests/test_base/test_codec.py b/python_polar_coding/tests/test_base/test_codec.py
@@ -2,11 +2,11 @@
 
 import numpy as np
 
-from python_polar_coding.polar_codes.base.decoder import BaseDecoder
-from python_polar_coding.polar_codes.base.functions import make_hard_decision
-from python_polar_coding.polar_codes.base.polar_code import (
-    BasicPolarCode,
-    BasicPolarCodeWithCRC,
+from python_polar_coding.polar_codes.base import (
+    BaseCRCPolarCodec,
+    BaseDecoder,
+    BasePolarCodec,
+    make_hard_decision,
 )
 
 
@@ -16,21 +16,21 @@ def decode_internal(self, received_llr: np.array):
         return make_hard_decision(received_llr)
 
 
-class SimplePC(BasicPolarCode):
+class SimplePC(BasePolarCodec):
     """Simple polar code for testing."""
     decoder_class = SimpleDecoder
 
-    def get_decoder(self):
+    def init_decoder(self):
         return self.decoder_class(
             n=self.n, mask=self.mask, is_systematic=self.is_systematic
         )
 
 
-class SimplePCCRC(BasicPolarCodeWithCRC):
+class SimplePCCRC(BaseCRCPolarCodec):
     """Simple polar code with CRC support for testing."""
     decoder_class = SimpleDecoder
 
-    def get_decoder(self):
+    def init_decoder(self):
         return self.decoder_class(
             n=self.n, mask=self.mask, is_systematic=self.is_systematic
         )
diff --git a/python_polar_coding/tests/test_basic_polar_code.py b/python_polar_coding/tests/test_basic_polar_code.py
diff --git a/python_polar_coding/tests/test_sc/__init__.py b/python_polar_coding/tests/test_sc/__init__.py
diff --git a/python_polar_coding/tests/test_sc/test_decoder.py b/python_polar_coding/tests/test_sc/test_decoder.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from python_polar_coding.polar_codes.decoders import SCDecoder
+from python_polar_coding.polar_codes.sc import SCDecoder
 
 
 class TestSCDecoder(TestCase):
diff --git a/python_polar_coding/tests/test_sc/test_sc.py b/python_polar_coding/tests/test_sc/test_sc.py
@@ -1,52 +1,51 @@
 from unittest import TestCase
 
-from python_polar_coding.polar_codes.sc import SCPolarCode
-
-from .base import BasicVerifyPolarCode
+from python_polar_coding.polar_codes.sc import SCPolarCodec
+from python_polar_coding.tests.base import BasicVerifyPolarCode
 
 
 class TestSCCode_1024_512(BasicVerifyPolarCode, TestCase):
-    polar_code_class = SCPolarCode
+    polar_code_class = SCPolarCodec
     code_parameters = {
         'N': 1024,
         'K': 512,
     }
 
 
 class TestSCCode_1024_256(BasicVerifyPolarCode, TestCase):
-    polar_code_class = SCPolarCode
+    polar_code_class = SCPolarCodec
     code_parameters = {
         'N': 1024,
         'K': 256,
     }
 
 
 class TestSCCode_1024_768(BasicVerifyPolarCode, TestCase):
-    polar_code_class = SCPolarCode
+    polar_code_class = SCPolarCodec
     code_parameters = {
         'N': 1024,
         'K': 768,
     }
 
 
 class TestSCCode_2048_512(BasicVerifyPolarCode, TestCase):
-    polar_code_class = SCPolarCode
+    polar_code_class = SCPolarCodec
     code_parameters = {
         'N': 2048,
         'K': 512,
     }
 
 
 class TestSCCode_2048_1024(BasicVerifyPolarCode, TestCase):
-    polar_code_class = SCPolarCode
+    polar_code_class = SCPolarCodec
     code_parameters = {
         'N': 2048,
         'K': 1024,
     }
 
 
 class TestSCCode_2048_1536(BasicVerifyPolarCode, TestCase):
-    polar_code_class = SCPolarCode
+    polar_code_class = SCPolarCodec
     code_parameters = {
         'N': 2048,
         'K': 1536,
diff --git a/python_polar_coding/tests/test_utils.py b/python_polar_coding/tests/test_utils.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+from .codec import SCPolarCodec`
	`2`	`+from .decoder import SCDecoder`