best_first_detector seems to perform as expected

Author: BastienTr

README.md

@@ -54,6 +54,7 @@ Available Features
 - OFDM Tx/Rx signal processing
 - MIMO Maximum Likelihood (ML) Detection.
 - MIMO K-best Schnorr-Euchner Detection.
+- MIMO Best-First Detection.
 - Convert channel matrix to Bit-level representation.
 - Computation of LogLikelihood ratio using max-log approximation.
 

commpy/modulation.py

@@ -15,8 +15,9 @@
    ofdm_rx              -- OFDM Receive Signal Processing
    mimo_ml              -- MIMO Maximum Likelihood (ML) Detection.
    kbest                -- MIMO K-best Schnorr-Euchner Detection.
-   bit_lvl_repr         -- Bit level representation.
-   max_log_approx       -- Max-log approximation.
+   best_first_detector  -- MIMO Best-First Detection.
+   bit_lvl_repr         -- Bit Level Representation.
+   max_log_approx       -- Max-Log Approximation.
 
 """
 from bisect import insort
@@ -31,7 +32,8 @@
 
 from commpy.utilities import bitarray2dec, dec2bitarray
 
-__all__ = ['PSKModem', 'QAMModem', 'ofdm_tx', 'ofdm_rx', 'mimo_ml', 'kbest', 'bit_lvl_repr', 'max_log_approx']
+__all__ = ['PSKModem', 'QAMModem', 'ofdm_tx', 'ofdm_rx', 'mimo_ml', 'kbest', 'best_first_detector',
+           'bit_lvl_repr', 'max_log_approx']
 
 
 class Modem:
@@ -347,48 +349,80 @@ def kbest(y, h, constellation, K, noise_var=0, output_type='hard', demode=None):
         raise ValueError('output_type must be "hard" or "soft"')
 
 
-class _Node:
-    """ Helper data model for best_first_detector. Implement __lt__ aka '<' as required to use bisect.insort. """
-    def __init__(self, symb_vectors, partial_metrics):
-        """
-        Recursive initializer that build a sequence of siblings.
-        Inputs are assumed to be ordered based on metric
-        """
-        if len(partial_metrics) == 1:
-            # There is one node to build
-            self.symb_vector = symb_vectors.reshape(-1)  # Insure that self.symb_vector is a 1d-ndarray
-            self.partial_metric = partial_metrics
-            self.best_sibling = None
-        else:
-            # Recursive call to build several nodes
-            self.symb_vector = symb_vectors[:, 0].reshape(-1)  # Insure that self.symb_vector is a 1d-ndarray
-            self.partial_metric = partial_metrics[0]
-            self.best_sibling = _Node(symb_vectors[:, 1:], partial_metrics[1:])
+def best_first_detector(y, h, constellation, stack_size, noise_var, demode, llr_max):
+    """ MIMO Best-First Detection.
 
-    def __lt__(self, other):
-        return self.partial_metric < other.partial_metric
+    Reference: G. He, X. Zhang, et Z. Liang, "Algorithm and Architecture of an Efficient MIMO Detector With Cross-Level
+     Parallel Tree-Search", IEEE Transactions on Very Large Scale Integration (VLSI) Systems, 2019
 
-    def expand(self, yt, r, constellation):
-        """ Build all children and return the best one. constellation must be a numpy ndarray."""
-        # Construct children's symbol vector
-        child_size = self.symb_vector.size + 1
-        children_symb_vectors = empty((child_size, constellation.size), constellation.dtype)
-        children_symb_vectors[:-1] = self.symb_vector[:, newaxis]
-        children_symb_vectors[-1] = constellation
 
-        # Compute children's partial metric and sort
-        children_metric = abs(yt[-child_size] - r[-child_size, -child_size:].dot(children_symb_vectors)) ** 2
-        children_metric += self.partial_metric
-        ordering = children_metric.argsort()
+    Parameters
+    ----------
+    y : 1D ndarray
+        Received complex symbols (length: num_receive_antennas)
+
+    h : 2D ndarray
+        Channel Matrix (shape: num_receive_antennas x num_transmit_antennas)
 
-        # Build children and return the best one
-        return _Node(children_symb_vectors[:, ordering], children_metric[ordering])
+    constellation : 1D ndarray of floats
+        Constellation used to modulate the symbols
 
+    stack_size : tuple of integers
+        Size of each stack (length: num_transmit_antennas - 1)
 
-def best_first_detector(y, h, constellation, stack_size, noise_var, demode, llr_max):
-    # TODO doc, __all__, readme
-    # TODO doc details : ref sans a priori, ordre de modulation entier
-    # TODO TESTS!!!
+    noise_var : positive float
+        Noise variance.
+        *Default* value is 0.
+
+    demode : function with prototype binary_word = demode(point)
+        Function that provide the binary word corresponding to a symbol vector.
+
+    llr_max : float
+        Max value for LLR clipping
+
+    Returns
+    -------
+    x : 1D ndarray of Log-Likelihood Ratios.
+        Detected vector (length: num_receive_antennas).
+    """
+
+    class _Node:
+        """ Helper data model that implements __lt__ (aka '<') as required to use bisect.insort. """
+
+        def __init__(self, symb_vectors, partial_metrics):
+            """
+            Recursive initializer that build a sequence of siblings.
+            Inputs are assumed to be ordered based on metric
+            """
+            if len(partial_metrics) == 1:
+                # There is one node to build
+                self.symb_vector = symb_vectors.reshape(-1)  # Insure that self.symb_vector is a 1d-ndarray
+                self.partial_metric = partial_metrics[0]
+                self.best_sibling = None
+            else:
+                # Recursive call to build several nodes
+                self.symb_vector = symb_vectors[:, 0].reshape(-1)  # Insure that self.symb_vector is a 1d-ndarray
+                self.partial_metric = partial_metrics[0]
+                self.best_sibling = _Node(symb_vectors[:, 1:], partial_metrics[1:])
+
+        def __lt__(self, other):
+            return self.partial_metric < other.partial_metric
+
+        def expand(self, yt, r, constellation):
+            """ Build all children and return the best one. constellation must be a numpy ndarray."""
+            # Construct children's symbol vector
+            child_size = self.symb_vector.size + 1
+            children_symb_vectors = empty((child_size, constellation.size), constellation.dtype)
+            children_symb_vectors[1:] = self.symb_vector[:, newaxis]
+            children_symb_vectors[0] = constellation
+
+            # Compute children's partial metric and sort
+            children_metric = abs(yt[-child_size] - r[-child_size, -child_size:].dot(children_symb_vectors)) ** 2
+            children_metric += self.partial_metric
+            ordering = children_metric.argsort()
+
+            # Build children and return the best one
+            return _Node(children_symb_vectors[:, ordering], children_metric[ordering])
 
     # Extract information from arguments
     nb_tx, nb_rx = h.shape
@@ -409,7 +443,7 @@ def best_first_detector(y, h, constellation, stack_size, noise_var, demode, llr_
     # Start process by adding the best root's child in the last stack
     stacks[-1].append(_Node(empty(0, constellation.dtype), array(0, float, ndmin=1)).expand(yt, r, constellation))
 
-    # While there is at least one non-empty stack (exempt first one)
+    # While there is at least one non-empty stack (exempt the first one)
     while any(stacks[1:]):
         # Node processing
         for idx_next_stack in range(len(stacks) - 1):
@@ -428,7 +462,7 @@ def best_first_detector(y, h, constellation, stack_size, noise_var, demode, llr_
                     try:
                         a2 = counter_hyp_metric[idx_this_stack:][map_bit_vector[idx_this_stack:] != bit_vector].max()
                     except ValueError:
-                        a2 = inf  # NumPy can compute max on empty an matrix
+                        a2 = inf  # NumPy cannot compute max on an empty matrix
                     radius = max(counter_hyp_metric[:idx_this_stack].max(), a2)
 
                 # Process best sibling
@@ -458,7 +492,7 @@ def best_first_detector(y, h, constellation, stack_size, noise_var, demode, llr_
         for idx_next_stack in range(len(stacks) - 1):
             del stacks[idx_next_stack + 1][stack_size[idx_next_stack]:]
 
-    return (counter_hyp_metric - map_metric) * map_bit_vector / 2 / noise_var
+    return ((map_metric - counter_hyp_metric) * map_bit_vector).reshape(-1) / 2 / noise_var
 
 
 def bit_lvl_repr(H, w):