Merge pull request #30 from YairMZ/bp_decoder

Fix numerical issue with BP decoder.
Add MS decoder.

Author: YairMZ

examples/ieee802_11_example.py

@@ -19,8 +19,8 @@
 
 # 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(spec=WiFiSpecCode.N648_R12, max_iter=20, channel_model=bsc_llr(p=p))
+p = 0.03
+decoder = DecoderWiFi(spec=WiFiSpecCode.N648_R12, max_iter=20, channel_model=bsc_llr(p=p), decoder_type="MS")
 
 # create a corrupted version of encoded codeword with error rate p
 corrupted = BitArray(encoded)

ldpc/decoder/graph.py

@@ -30,12 +30,13 @@ def add_v_node(self, ordering_key: int, name: str = "", channel_model: Optional[
         self.v_nodes[node.uid] = node
         return node
 
-    def add_c_node(self, name: str = "", ordering_key: Optional[int] = None) -> CNode:
+    def add_c_node(self, name: str = "", ordering_key: Optional[int] = None, decoder: Optional[str] = "BP") -> CNode:
         """
         :param ordering_key: use only for debug purposes
         :param name: name of node
+        :param decoder: must be either "BP" or "MS" for min-sum decoder
         """
-        node = CNode(name, ordering_key)
+        node = CNode(name, ordering_key, decoder=decoder)
         self.c_nodes[node.uid] = node
         return node
 
@@ -68,10 +69,10 @@ def add_edges_by_name(self, edges_set: set[tuple[str, str]]) -> None:
         for v_name, c_name in edges_set:
             v_uid = [node.uid for node in self.v_nodes.values() if node.name == v_name]
             if not v_uid:
-                raise ValueError("No v-node with name " + v_name + " in graph")
+                raise ValueError(f"No v-node with name {v_name} in graph")
             c_uid = [node.uid for node in self.c_nodes.values() if node.name == c_name]
             if not c_uid:
-                raise ValueError("No c-node with name " + c_name + " in graph")
+                raise ValueError(f"No c-node with name {c_name} in graph")
             self.add_edge(v_uid[0], c_uid[0])
 
     def get_edges(self, by_name: bool = False) -> Union[set[tuple[str, str]], EdgesSet]:
@@ -94,12 +95,14 @@ def to_nx(self) -> nx.Graph:
         return g
 
     @classmethod
-    def from_biadjacency_matrix(cls, h: npt.ArrayLike, channel_model: Optional[ChannelModel] = None) -> TannerGraph:
+    def from_biadjacency_matrix(cls, h: npt.ArrayLike, channel_model: Optional[ChannelModel] = None,
+                                decoder: Optional[str] = "BP") -> TannerGraph:
         """
         Creates a Tanner Graph from a biadjacency matrix, nodes are ordered according to matrix indices.
 
         :param channel_model: channel model to compute channel symbols llr within v nodes
         :param h: parity check matrix, shape MXN with M check nodes and N variable nodes. assumed binary matrix.
+        :param decoder: must be either "BP" or "MS" for min-sum decoder
         """
         g = TannerGraph()
         h = np.array(h)
@@ -109,7 +112,7 @@ def from_biadjacency_matrix(cls, h: npt.ArrayLike, channel_model: Optional[Chann
             v_uid = g.add_v_node(name=f"v{i}", channel_model=channel_model, ordering_key=i).uid
             ordered_vnode_uid[i] = v_uid
         for j in range(m):
-            c_uid = g.add_c_node(name=f"c{j}", ordering_key=j).uid
+            c_uid = g.add_c_node(name=f"c{j}", ordering_key=j, decoder=decoder).uid
             for i in range(n):
                 if h[j, i] == 1:
                     g.add_edge(ordered_vnode_uid[i], c_uid)

ldpc/decoder/ieee802_11_decoder.py

@@ -11,16 +11,19 @@ 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, spec: WiFiSpecCode, max_iter: int, channel_model: Optional[ChannelModel] = None):
+    def __init__(self, spec: WiFiSpecCode, max_iter: int, channel_model: Optional[ChannelModel] = None,
+                 decoder_type: Optional[str] = "BP"):
         """
 
         :param channel_model: a callable which receives a channel input, and returns the channel llr
         :param spec: specify which code from the spec we use
         :param max_iter: The maximal number of iterations for belief propagation algorithm
+        :param decoder_type: must be either "BP" or "MS" for min-sum decoder
         """
         self.spec = spec
         qc_file = QCFile.from_file(os.path.join(self._spec_base_path, spec.name + ".qc"))
         h = qc_file.to_array()
         m, n = h.shape
         k = n - m
-        super().__init__(h, max_iter, info_idx=np.array([True] * k + [False] * m), channel_model=channel_model)
+        super().__init__(h, max_iter, info_idx=np.array([True] * k + [False] * m), channel_model=channel_model,
+                         decoder_type=decoder_type)

ldpc/decoder/log_spa_decoder.py

@@ -20,18 +20,20 @@ class InfoBitsNotSpecified(Exception):
 class LogSpaDecoder:
     """Decode codewords according to Log-SPA version of the belief propagation algorithm"""
     def __init__(self, h: ArrayLike, max_iter: int, info_idx: Optional[NDArray[np.bool_]] = None,
-                 channel_model: Optional[ChannelModel] = None):
+                 channel_model: Optional[ChannelModel] = None, decoder_type: Optional[str] = "BP"):
         """
 
         :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.
+        :param decoder_type: must be either "BP" or "MS" for min-sum decoder
         """
+        self.decoder_type = decoder_type
         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.graph = TannerGraph.from_biadjacency_matrix(h=self.h, channel_model=channel_model, decoder=decoder_type)
         self.n = len(self.graph.v_nodes)
         self.max_iter = max_iter
         ordered_cnodes = sorted(self.graph.c_nodes.values())
@@ -80,14 +82,16 @@ def decode(self, channel_word: Sequence[np.float_], max_iter: Optional[int] = No
             if not syndrome.any():
                 break
 
+        # for each vnode how many equations are failed
+        vnode_validity: npt.NDArray[np.int_] = np.dot(syndrome, self.h)
         # for each vnode how many equations are fulfilled
-        vnode_validity: npt.NDArray[np.int_] = np.zeros(self.n, dtype=np.int_)
-        syndrome_compliance = {cnode: int(val == 0) for cnode, val in zip(self.ordered_cnodes_uids, syndrome)}
-
-        for idx, vnode in enumerate(self._ordered_vnodes):
-            neighbors = vnode.get_neighbors()
-            for neighbor in neighbors:
-                vnode_validity[idx] += 2*syndrome_compliance[neighbor] - 1
+        # vnode_validity: npt.NDArray[np.int_] = np.zeros(self.n, dtype=np.int_)
+        # syndrome_compliance = {cnode: int(val == 0) for cnode, val in zip(self.ordered_cnodes_uids, syndrome)}
+        #
+        # for idx, vnode in enumerate(self._ordered_vnodes):
+        #     neighbors = vnode.get_neighbors()
+        #     for neighbor in neighbors:
+        #         vnode_validity[idx] += 2*syndrome_compliance[neighbor] - 1
         return estimate, llr, not syndrome.any(), iteration+1, syndrome, vnode_validity
 
     def info_bits(self, estimate: NDArray[np.int_]) -> Bits:

ldpc/decoder/node.py

@@ -73,6 +73,16 @@ def __lt__(self, other: Any) -> bool:
 
 class CNode(Node):
     """Check nodes in Tanner graph"""
+
+    def __init__(self, name: str = "", ordering_key: Optional[int] = None,  decoder: Optional[str] = "BP") -> None:
+        """
+        :param ordering_key: used to order nodes per their order in the parity check matrix
+        :param name: optional name of node
+        :param decoder: must be either "BP" or "MS" for min-sum decoder
+        """
+        self.decoder_type = decoder
+        super().__init__(name, ordering_key)
+
     def initialize(self) -> None:
         """
         clear received messages
@@ -84,11 +94,16 @@ def message(self, requester_uid: int) -> np.float_:
         pass messages from c-nodes to v-nodes
         :param requester_uid: uid of requesting v-node
         """
+        q: npt.NDArray[np.float_] = np.array([msg for uid, msg in self.received_messages.items() if uid != requester_uid])
+        if self.decoder_type == "MS":
+            return np.prod(np.sign(q)) * np.absolute(q).min()   # type: ignore
+
+
+        # full BP
         def phi(x: npt.NDArray[np.float_]) -> Any:
             """see sources for definition and reasons for use of this function"""
-            return -np.log(np.tanh(x/2))
-        q: npt.NDArray[np.float_] = np.array([msg for uid, msg in self.received_messages.items() if uid != requester_uid])
-        return np.prod(np.sign(q))*phi(sum(phi(np.absolute(q)+np.finfo(np.float_).eps)))  # type: ignore
+            return -np.log(np.maximum(np.tanh(x/2), 1e3*np.finfo(np.float_).eps))
+        return np.prod(np.sign(q))*phi(sum(phi(np.absolute(q))))  # type: ignore
 
 
 class VNode(Node):

tests/test_ieee802_11.py

@@ -102,6 +102,31 @@ def test_decoder_648_r12(self) -> None:
         assert sum(encoded_ref ^ decoded) == 0
         assert sum(info_bits ^ decoded_info) == 0
 
+    def test_ms_decoder_1296_r23(self) -> None:
+        info_bits = Bits(auto=np.genfromtxt(
+            'tests/test_data/ieee_802_11/info_bits_N1296_R23.csv', delimiter=',', dtype=np.int_))
+        encoded_ref = Bits(auto=np.genfromtxt(
+            'tests/test_data/ieee_802_11/encoded_N1296_R23.csv', delimiter=',', dtype=np.int_))
+        p = 0.01
+
+        corrupted = BitArray(encoded_ref)
+        no_errors = int(len(corrupted) * p)
+        rng = np.random.default_rng()
+        error_idx = rng.choice(len(corrupted), size=no_errors, replace=False)
+        for idx in error_idx:
+            corrupted[idx] = not corrupted[idx]
+
+        decoder = DecoderWiFi(spec=WiFiSpecCode.N1296_R23, max_iter=20, channel_model=bsc_llr(p=p),decoder_type="MS")
+        decoded = Bits()
+        decoded_info = Bits()
+        for frame_idx in range(len(corrupted) // decoder.n):
+            decoder_output = decoder.decode(corrupted[frame_idx * decoder.n: (frame_idx + 1) * decoder.n])
+            decoded += decoder_output[0]
+            decoded_info += decoder.info_bits(decoder_output[0])
+            assert decoder_output[2] is True
+        assert sum(encoded_ref ^ decoded) == 0
+        assert sum(info_bits ^ decoded_info) == 0
+
     def test_decoder_no_info(self) -> None:
         p = 0.1
         h = QCFile.from_file("ldpc/code_specs/ieee802.11/N648_R12.qc").to_array()