Merge pull request #24 from YairMZ/bp_decoder

allow decoder to work by receiving llr's, and without specifying a channel model

Author: YairMZ

README.md

@@ -23,6 +23,11 @@ python -m pytest  -n auto --cov-report=html
 ```
 to run tests in parallel (with number of CPU's dictated by machine) to speed up tests.
 
+Verify static typing with
+```shell
+mypy --strict --config-file .mypy.ini .
+```
+
 -----
 ## Included modules
  - [Utilities](ldpc/utils/README.md): implementing various utility operations to assist with encoding, decoding and 
@@ -36,7 +41,7 @@ simulations.
 ```python
 import numpy as np
 from bitstring import BitArray, Bits
-from ldpc.decoder import LogSpaDecoder, bsc_llr
+from ldpc.decoder import DecoderWiFi, bsc_llr
 from ldpc.encoder import EncoderWiFi
 from ldpc.wifi_spec_codes import WiFiSpecCode
 from ldpc.utils import QCFile
@@ -56,7 +61,7 @@ np.dot(h, np.array(encoded)) % 2  # creates an all zero vector as required.
 # create a decoder which assumes a probability of p=0.05 for bit flips by the channel
 # allow up to 20 iterations for the bp decoder.
 p = 0.05
-decoder = LogSpaDecoder(bsc_llr(p=p), h=h, max_iter=20, info_idx=np.array([True]*324 + [False]*324))
+decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20, channel_model=bsc_llr(p=p))
 
 # create a corrupted version of encoded codeword with error rate p
 corrupted = BitArray(encoded)
@@ -68,6 +73,15 @@ decoded, llr, decode_success, num_of_iterations = decoder.decode(corrupted)
 # Verify correct decoding
 print(Bits(decoded) == encoded)  # true
 info = decoder.info_bits(decoded)
+
+# a decoder can also be instantiated without a channel model, in which case llr is expected to be sent for decoding instead of
+# hard channel outputs.
+channel = bsc_llr(p=p)
+channel_llr = channel(np.array(corrupted, dtype=np.int_))
+decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20)
+decoded, llr2, decode_success, num_of_iterations = decoder.decode(channel_llr)
+print(Bits(decoded) == encoded)  # true
+info = decoder.info_bits(decoded)
 ```
 The example is also included as a jupyter notebook. Note however, that you need to launch the notebook from the correct 
 path for it to be able to access installed packages. To run the notebook:

examples/ieee802_11_example.py

@@ -20,7 +20,7 @@
 # create a decoder which assumes a probability of p=0.05 for bit flips by the channel
 # allow up to 20 iterations for the bp decoder.
 p = 0.05
-decoder = DecoderWiFi(bsc_llr(p=p), spec=WiFiSpecCode.N648_R12, max_iter=20)
+decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20, channel_model=bsc_llr(p=p))
 
 # create a corrupted version of encoded codeword with error rate p
 corrupted = BitArray(encoded)
@@ -32,3 +32,12 @@
 # Verify correct decoding
 print(Bits(decoded) == encoded)  # true
 info = decoder.info_bits(decoded)
+
+# a decoder can also be instantiated without a channel model, in which case llr is expected to be sent for decoding instead of
+# hard channel outputs.
+channel = bsc_llr(p=p)
+channel_llr = channel(np.array(corrupted, dtype=np.int_))
+decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20)
+decoded, llr2, decode_success, num_of_iterations = decoder.decode(channel_llr)
+print(Bits(decoded) == encoded)  # true
+info = decoder.info_bits(decoded)

ldpc/decoder/channel_models.py

@@ -3,7 +3,7 @@
 
 __all__ = ["ChannelModel", "bsc_llr"]
 
-ChannelModel = Callable[[int], np.float_]
+ChannelModel = Callable[[np.float_], np.float_]
 
 
 def bsc_llr(p: float) -> ChannelModel:

ldpc/decoder/graph.py

@@ -20,13 +20,13 @@ def __init__(self) -> None:
         self.c_nodes: dict[int, CNode] = {}
         self.edges: EdgesSet = set()
 
-    def add_v_node(self, channel_model: ChannelModel, ordering_key: int, name: str = "") -> VNode:
+    def add_v_node(self, ordering_key: int, name: str = "", channel_model: Optional[ChannelModel] = None) -> VNode:
         """
         :param ordering_key: should reflect order according to parity check matrix, channel symbols in order
         :param name: name of node.
-        :param channel_model: add an exiting node to graph. If not used a new node is created.
+        :param channel_model: optional channel model to compute llr out of hard channel outputs. Of not used llr are expected.
         """
-        node = VNode(channel_model, ordering_key, name)
+        node = VNode(ordering_key, name=name, channel_model=channel_model)
         self.v_nodes[node.uid] = node
         return node
 
@@ -94,7 +94,7 @@ def to_nx(self) -> nx.Graph:
         return g
 
     @classmethod
-    def from_biadjacency_matrix(cls, h: npt.ArrayLike, channel_model: ChannelModel) -> TannerGraph:
+    def from_biadjacency_matrix(cls, h: npt.ArrayLike, channel_model: Optional[ChannelModel] = None) -> TannerGraph:
         """
         Creates a Tanner Graph from a biadjacency matrix, nodes are ordered according to matrix indices.
 

ldpc/decoder/ieee802_11_decoder.py

@@ -4,13 +4,14 @@
 from ldpc.wifi_spec_codes import WiFiSpecCode
 import os
 from ldpc.utils.qc_format import QCFile
+from typing import Optional
 
 
 class DecoderWiFi(LogSpaDecoder):
     """Decode messages according to the codes in the IEEE802.11n standard using Log SPA decoder"""
     _spec_base_path: str = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'code_specs', 'ieee802.11')
 
-    def __init__(self, channel_model: ChannelModel, spec: WiFiSpecCode, max_iter: int):
+    def __init__(self, spec: WiFiSpecCode, max_iter: int, channel_model: Optional[ChannelModel] = None):
         """
 
         :param channel_model: a callable which receives a channel input, and returns the channel llr
@@ -22,4 +23,4 @@ def __init__(self, channel_model: ChannelModel, spec: WiFiSpecCode, max_iter: in
         h = qc_file.to_array()
         m, n = h.shape
         k = n - m
-        super().__init__(channel_model, h, max_iter, info_idx=np.array([True] * k + [False] * m))
+        super().__init__(h, max_iter, info_idx=np.array([True] * k + [False] * m), channel_model=channel_model)

ldpc/decoder/log_spa_decoder.py

@@ -19,22 +19,23 @@ class InfoBitsNotSpecified(Exception):
 
 class LogSpaDecoder:
     """Decode codewords according to Log-SPA version of the belief propagation algorithm"""
-    def __init__(self, channel_model: ChannelModel, h: ArrayLike, max_iter: int,
-                 info_idx: Optional[NDArray[np.bool_]] = None):
+    def __init__(self, h: ArrayLike, max_iter: int, info_idx: Optional[NDArray[np.bool_]] = None,
+                 channel_model: Optional[ChannelModel] = None):
         """
 
-        :param channel_model: a callable which receives a channel input, and returns the channel llr
         :param h:the parity check code matrix of the code
         :param max_iter: The maximal number of iterations for belief propagation algorithm
         :param info_idx: a boolean array representing the indices of information bits in the code
+        :param channel_model: optional, a callable which receives a channel input, and returns the channel llr. If not
+        specified, llr is expected to be fed into the decoder.
         """
         self.info_idx = info_idx
         self.h: npt.NDArray[np.int_] = np.array(h)
         self.graph = TannerGraph.from_biadjacency_matrix(h=self.h, channel_model=channel_model)
         self.n = len(self.graph.v_nodes)
         self.max_iter = max_iter
 
-    def decode(self, channel_word: Sequence[int], max_iter: Optional[int] = None) \
+    def decode(self, channel_word: Sequence[np.float_], max_iter: Optional[int] = None) \
             -> tuple[NDArray[np.int_], NDArray[np.float_], bool, int]:
         """
         decode a sequence received from the channel
@@ -69,7 +70,8 @@ def decode(self, channel_word: Sequence[int], max_iter: Optional[int] = None) \
 
             # Check stop condition
             llr: npt.NDArray[np.float_] = np.array([node.estimate() for node in self.graph.ordered_v_nodes()])
-            estimate: npt.NDArray[np.int_] = np.array([1 if node_llr < 0 else 0 for node_llr in llr], dtype=np.int_)
+            estimate: npt.NDArray[np.int_] = np.array(llr < 0, dtype=np.int_)
+            # np.array([1 if node_llr < 0 else 0 for node_llr in llr], dtype=np.int_)
             syndrome = self.h.dot(estimate) % 2
             if not syndrome.any():
                 break

ldpc/decoder/node.py

@@ -25,7 +25,7 @@ def __init__(self, name: str = "", ordering_key: Optional[int] = None) -> None:
         :param name: name of node
         """
         self.uid = next(Node._uid_generator)
-        self.name = name if name else str(self.uid)
+        self.name = name or str(self.uid)
         self.ordering_key = ordering_key if ordering_key is not None else self.uid
         self.neighbors: dict[int, Node] = {}  # keys as senders uid
         self.received_messages: dict[int, Any] = {}  # keys as senders uid, values as messages
@@ -34,10 +34,7 @@ def register_neighbor(self, neighbor: Node) -> None:
         self.neighbors[neighbor.uid] = neighbor
 
     def __str__(self) -> str:
-        if self.name:
-            return self.name
-        else:
-            return str(self.uid)
+        return self.name or str(self.uid)
 
     def get_neighbors(self) -> list[int]:
         """
@@ -96,24 +93,27 @@ def phi(x: npt.NDArray[np.float_]) -> Any:
 
 class VNode(Node):
     """Variable nodes in Tanner graph"""
-    def __init__(self, channel_model: ChannelModel, ordering_key: int, name: str = ""):
+    def __init__(self, ordering_key: int, name: str = "", channel_model: Optional[ChannelModel] = None):
         """
         :param channel_model: a function which receives channel outputs anr returns relevant message
         :param ordering_key: used to order nodes per their order in the parity check matrix
         :param name: optional name of node
         """
-        self.channel_model: ChannelModel = channel_model
-        self.channel_symbol: int = None  # type: ignore # currently assuming hard channel symbols
+        self.channel_model: Optional[ChannelModel] = channel_model
+        self.channel_symbol: np.float_ = None  # type: ignore
         self.channel_llr: np.float_ = None  # type: ignore
         super().__init__(name, ordering_key)
 
-    def initialize(self, channel_symbol: int) -> None:  # type: ignore
+    def initialize(self, channel_symbol: np.float_) -> None:  # type: ignore
         """
         clear received messages and initialize channel llr with channel bit
         :param channel_symbol: bit received from channel, currently assumes hard inputs.
         """
         self.channel_symbol = channel_symbol
-        self.channel_llr = self.channel_model(channel_symbol)
+        if self.channel_model is not None:
+            self.channel_llr = self.channel_model(channel_symbol)
+        else:
+            self.channel_llr = channel_symbol
         self.received_messages = {node_uid: 0 for node_uid in self.neighbors}
 
     def message(self, requester_uid: int) -> np.float_:
@@ -127,4 +127,4 @@ def message(self, requester_uid: int) -> np.float_:
 
     def estimate(self) -> np.float_:
         """provide a soft bit estimate"""
-        return self.channel_llr + np.sum(list(self.received_messages.values()))  # type: ignore
+        return self.channel_llr + np.sum(list(self.received_messages.values()), dtype=np.float_)  # type: ignore

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup
 import pathlib
 
-VERSION = '0.1.3'
+VERSION = '0.2.0'
 
 # The directory containing this file
 HERE = pathlib.Path(__file__).resolve().parent

tests/test_ieee802_11.py

@@ -72,7 +72,7 @@ class TestDecoder802_11:
     def test_incorrect_length(self) -> None:
         p = 0.1
         h = QCFile.from_file("ldpc/code_specs/ieee802.11/N648_R12.qc").to_array()
-        decoder = LogSpaDecoder(bsc_llr(p=p), h=h, max_iter=20)
+        decoder = LogSpaDecoder(h=h, max_iter=20, channel_model=bsc_llr(p=p))
         bits: npt.NDArray[np.int_] = np.array([1, 1, 0], dtype=np.int_)
         with pytest.raises(IncorrectLength):
             decoder.decode(bits)  # type: ignore
@@ -91,7 +91,7 @@ def test_decoder_648_r12(self) -> None:
         for idx in error_idx:
             corrupted[idx] = not corrupted[idx]
 
-        decoder = DecoderWiFi(bsc_llr(p=p), spec=WiFiSpecCode.N648_R12, max_iter=20)
+        decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20, channel_model=bsc_llr(p=p))
         decoded = Bits()
         decoded_info = Bits()
         for frame_idx in range(len(corrupted) // decoder.n):
@@ -105,7 +105,7 @@ def test_decoder_648_r12(self) -> None:
     def test_decoder_no_info(self) -> None:
         p = 0.1
         h = QCFile.from_file("ldpc/code_specs/ieee802.11/N648_R12.qc").to_array()
-        decoder = LogSpaDecoder(bsc_llr(p=p), h=h, max_iter=20)
+        decoder = LogSpaDecoder(h=h, max_iter=20, channel_model=bsc_llr(p=p))
         estimate: npt.NDArray[np.int_] = np.array(Bits(bytes=bytes(list(range(81)))), np.int_)
         with pytest.raises(InfoBitsNotSpecified):
             decoder.info_bits(estimate)

wifi_example.ipynb

@@ -11,16 +11,16 @@
    "source": [
     "import numpy as np\n",
     "from bitstring import BitArray, Bits\n",
-    "from ldpc.decoder import LogSpaDecoder, bsc_llr\n",
+    "from ldpc.decoder import DecoderWiFi, bsc_llr\n",
     "from ldpc.encoder import EncoderWiFi\n",
     "from ldpc.wifi_spec_codes import WiFiSpecCode"
    ],
-   "execution_count": 4,
+   "execution_count": 1,
    "outputs": []
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 2,
    "outputs": [],
    "source": [
     "# create information bearing bits\n",
@@ -38,7 +38,7 @@
     "# create a decoder which assumes a probability of p=0.05 for bit flips by the channel\n",
     "# allow up to 20 iterations for the bp decoder.\n",
     "p = 0.05\n",
-    "decoder = LogSpaDecoder(bsc_llr(p=p), h=h, max_iter=20, info_idx=np.array([True]*324 + [False]*324))"
+    "decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20, channel_model=bsc_llr(p=p))"
    ],
    "metadata": {
     "collapsed": false,
@@ -49,7 +49,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "outputs": [
     {
      "name": "stdout",
@@ -77,6 +77,35 @@
      "name": "#%%\n"
     }
    }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "True\n"
+     ]
+    }
+   ],
+   "source": [
+    "# a decoder can also be instantiated without a channel model, in which case llr is expected to be sent for decoding instead of\n",
+    "# hard channel outputs.\n",
+    "channel = bsc_llr(p=p)\n",
+    "channel_llr = channel(np.array(corrupted, dtype=np.int_))\n",
+    "decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20)\n",
+    "decoded, llr2, decode_success, num_of_iterations = decoder.decode(channel_llr)\n",
+    "print(Bits(decoded) == encoded)  # true\n",
+    "info = decoder.info_bits(decoded)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
   }
  ],
  "metadata": {