Add numba compilation of vnode messages, suppress mypy issues

Author: YairMZ

.mypy.ini

@@ -30,4 +30,7 @@ allow_untyped_calls = True
 ignore_missing_imports = True
 
 [mypy-setuptools.*]
+ignore_missing_imports = True
+
+[mypy-numba.*]
 ignore_missing_imports = True

ldpc/decoder/log_spa_decoder.py

@@ -83,7 +83,7 @@ def decode(self, channel_word: Sequence[np.float_], max_iter: Optional[int] = No
                 break
 
         # for each vnode how many equations are failed
-        vnode_validity: npt.NDArray[np.int_] = np.dot(syndrome, self.h)
+        vnode_validity: npt.NDArray[np.int_] = np.dot(syndrome, self.h)  # type: ignore
         # 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)}

ldpc/decoder/node.py

@@ -1,16 +1,27 @@
 from __future__ import annotations
 import numpy as np
 import itertools
-from typing import Any, Optional
+from typing import Any, Optional, Tuple
 from functools import total_ordering
 from abc import ABC, abstractmethod
 import numpy.typing as npt
 from ldpc.decoder.channel_models import ChannelModel
+from numba import jit
 
 
 __all__ = ["Node", "CNode", "VNode"]
 
 
+@jit(nopython=True, cache=True)  # type: ignore
+def c_message(requester_uid: int, senders: Tuple[int], received_messages: Tuple[float]) -> np.float_:
+    q: npt.NDArray[np.float_] = np.array([msg for uid, msg in zip(senders, received_messages) if uid != requester_uid])
+    return -np.prod(np.sign(q)) * np.log(  # type: ignore
+        np.maximum(np.tanh(
+            sum(
+                -np.log(np.maximum(np.tanh(np.absolute(q) / 2), 1e3 * np.finfo(np.float_).eps))
+            ) / 2), 1e3 * np.finfo(np.float_).eps))
+
+
 @total_ordering  # type: ignore
 class Node(ABC):
     """Base class VNodes anc CNodes.
@@ -28,7 +39,7 @@ def __init__(self, name: str = "", ordering_key: Optional[int] = None) -> None:
         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
+        self.received_messages: dict[int, np.float_] = {}  # keys as senders uid, values as messages
 
     def register_neighbor(self, neighbor: Node) -> None:
         self.neighbors[neighbor.uid] = neighbor
@@ -48,7 +59,7 @@ def receive_messages(self) -> None:
             self.received_messages[node_id] = node.message(self.uid)
 
     @abstractmethod
-    def message(self, requester_uid: int) -> Any:
+    def message(self, requester_uid: int) -> np.float_:
         """Used to return a message to the requesting node"""
         pass
 
@@ -87,23 +98,30 @@ def initialize(self) -> None:
         """
         clear received messages
         """
-        self.received_messages = {node_uid: 0 for node_uid in self.neighbors}
+        self.received_messages = {node_uid: 0 for node_uid in self.neighbors}  # type: ignore
 
     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":
+            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)) * 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.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
+        return c_message(requester_uid, tuple(self.received_messages.keys()), tuple(self.received_messages.values()))  # type: ignore
+
+        # def phi(x: npt.NDArray[np.float_]) -> npt.NDArray[np.float_]:
+        #     """see sources for definition and reasons for use of this function"""
+        #     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
+
+        # return -np.prod(np.sign(q)) * np.log(
+        #     np.maximum(np.tanh(
+        #         sum(
+        #             -np.log(np.maximum(np.tanh(np.absolute(q) / 2), 1e3 * np.finfo(np.float_).eps))
+        #         ) / 2), 1e3 * np.finfo(np.float_).eps))
 
 
 class VNode(Node):
@@ -129,7 +147,7 @@ def initialize(self, channel_symbol: np.float_) -> None:  # type: ignore
             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}
+        self.received_messages = {node_uid: 0 for node_uid in self.neighbors}  # type: ignore
 
     def message(self, requester_uid: int) -> np.float_:
         """

ldpc/encoder/h_based_encoder.py

@@ -34,6 +34,6 @@ def encode(self, information_bits: Bits) -> Bits:
         p: npt.NDArray[np.int_] = np.mod(np.matmul(self.h[:, :self.k], info), 2)
         if not self.identity_p:
             for l in range(1, self.m):
-                p[l] += np.mod(np.dot(self.h[l, self.k:self.k+l], p[:l]), 2)
+                p[l] += np.mod(np.dot(self.h[l, self.k:self.k+l], p[:l]), 2)  # type: ignore
         encoded[self.k:] = p
         return Bits(encoded)

requirements.txt

@@ -1,4 +1,5 @@
-numpy~=1.22.0
-bitstring~=3.1.9
-scipy~=1.7.2
-networkx~=2.6.3
+numpy~=1.21.6
+bitstring>=3.1.9
+scipy>=1.7.2
+networkx>=2.6.3
+numba

setup.cfg

@@ -1,2 +1,2 @@
 [metadata]
-description-file = README.md
+description_file = README.md

setup.py

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

tests/test_a_list.py

@@ -44,7 +44,7 @@ def test_array_io(self) -> None:
                                                  [0, 1, 3],
                                                  [0, 2, 3]]
         b = a.to_array()
-        np.testing.assert_array_equal(arr, b)
+        np.testing.assert_array_equal(arr, b)  # type: ignore
 
     def test_verify_alist(self) -> None:
         original_file = "ldpc/code_specs/Mackay_96.3.963.alist"

tests/test_encoders.py

@@ -17,7 +17,7 @@ def test_params(self) -> None:
         enc = EncoderG(g)
         assert enc.n == 7
         assert enc.k == 4
-        np.testing.assert_array_equal(g, enc.generator)
+        np.testing.assert_array_equal(g, enc.generator)  # type: ignore
 
     def test_encoding(self) -> None:
         g = AList.from_file("tests/test_data/Hamming_7_4_g.alist").to_array()
@@ -47,7 +47,7 @@ def test_params(self) -> None:
         assert enc.n == 4098
         assert enc.m == 3095
         assert enc.k == 4098-3095
-        np.testing.assert_array_equal(h, enc.h)
+        np.testing.assert_array_equal(h, enc.h)  # type: ignore
 
     def test_encoding(self) -> None:
         h = AList.from_file("tests/test_data/systematic_4098_3095.alist").to_array()

tests/test_ieee802_11.py

@@ -18,9 +18,9 @@ def test_incorrect_length(self) -> None:
     def test_encoding_648_r12(self) -> None:
         # comparing encodings with reference implementation by: https://github.com/tavildar/LDPC
         info_bits = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/info_bits_N648_R12.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/info_bits_N648_R12.csv', delimiter=',', dtype=np.int_))  # type: ignore
         encoded_ref = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/encoded_N648_R12.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/encoded_N648_R12.csv', delimiter=',', dtype=np.int_))  # type: ignore
         enc = EncoderWiFi(WiFiSpecCode.N648_R12)
 
         encoded = Bits()
@@ -31,9 +31,9 @@ def test_encoding_648_r12(self) -> None:
     def test_encoding_648_r56(self) -> None:
         # comparing encodings with reference implementation by: https://github.com/tavildar/LDPC
         info_bits = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/info_bits_N648_R56.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/info_bits_N648_R56.csv', delimiter=',', dtype=np.int_))  # type: ignore
         encoded_ref = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/encoded_N648_R56.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/encoded_N648_R56.csv', delimiter=',', dtype=np.int_))  # type: ignore
         enc = EncoderWiFi(WiFiSpecCode.N648_R56)
 
         encoded = Bits()
@@ -44,9 +44,9 @@ def test_encoding_648_r56(self) -> None:
     def test_encoding_1296_r23(self) -> None:
         # comparing encodings with reference implementation by: https://github.com/tavildar/LDPC
         info_bits = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/info_bits_N1296_R23.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/info_bits_N1296_R23.csv', delimiter=',', dtype=np.int_))  # type: ignore
         encoded_ref = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/encoded_N1296_R23.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/encoded_N1296_R23.csv', delimiter=',', dtype=np.int_))  # type: ignore
         enc = EncoderWiFi(WiFiSpecCode.N1296_R23)
 
         encoded = Bits()
@@ -57,9 +57,9 @@ def test_encoding_1296_r23(self) -> None:
     def test_encoding_1944_r34(self) -> None:
         # comparing encodings with reference implementation by: https://github.com/tavildar/LDPC
         info_bits = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/info_bits_N1944_R34.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/info_bits_N1944_R34.csv', delimiter=',', dtype=np.int_))  # type: ignore
         encoded_ref = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/encoded_N1944_R34.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/encoded_N1944_R34.csv', delimiter=',', dtype=np.int_))  # type: ignore
         enc = EncoderWiFi(WiFiSpecCode.N1944_R34)
 
         encoded = Bits()
@@ -79,9 +79,9 @@ def test_incorrect_length(self) -> None:
 
     def test_decoder_648_r12(self) -> None:
         info_bits = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/info_bits_N648_R12.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/info_bits_N648_R12.csv', delimiter=',', dtype=np.int_))  # type: ignore
         encoded_ref = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/encoded_N648_R12.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/encoded_N648_R12.csv', delimiter=',', dtype=np.int_))  # type: ignore
         p = 0.01
 
         corrupted = BitArray(encoded_ref)
@@ -104,9 +104,9 @@ def test_decoder_648_r12(self) -> None:
 
     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_))
+            'tests/test_data/ieee_802_11/info_bits_N1296_R23.csv', delimiter=',', dtype=np.int_))  # type: ignore
         encoded_ref = Bits(auto=np.genfromtxt(
-            'tests/test_data/ieee_802_11/encoded_N1296_R23.csv', delimiter=',', dtype=np.int_))
+            'tests/test_data/ieee_802_11/encoded_N1296_R23.csv', delimiter=',', dtype=np.int_))  # type: ignore
         p = 0.01
 
         corrupted = BitArray(encoded_ref)