Refactors, adds doc and tests for modulation.py and utilities.py

Author: BastienTr

commpy/modulation.py

@@ -30,13 +30,46 @@
 from numpy.fft import fft, ifft
 from numpy.linalg import qr, norm
 
-from commpy.utilities import bitarray2dec, dec2bitarray
+from commpy.utilities import bitarray2dec, dec2bitarray, signal_power
 
 __all__ = ['PSKModem', 'QAMModem', 'ofdm_tx', 'ofdm_rx', 'mimo_ml', 'kbest', 'best_first_detector',
            'bit_lvl_repr', 'max_log_approx']
 
 
 class Modem:
+
+    """ Creates a custom Modem object.
+
+        Parameters
+        ----------
+        constellation : array-like with a length which is a power of 2
+                        Constellation of the custom modem
+
+        Attributes
+        ----------
+        constellation : 1D-ndarray of complex
+                        Modem constellation. If changed, the length of the new constellation must be a power of 2.
+
+        Es            : float
+                        Average energy per symbols.
+
+        m             : integer
+                        Constellation length.
+
+        num_bits_symb : integer
+                        Number of bits per symbol.
+
+        Raises
+        ------
+        ValueError
+                        If the constellation is changed to an array-like with length that is not a power of 2.
+        """
+
+    def __init__(self, constellation):
+        """ Creates a custom Modem object. """
+
+        self.constellation = constellation
+
     def modulate(self, input_bits):
         """ Modulate (map) an array of bits to constellation symbols.
 
@@ -52,7 +85,7 @@ def modulate(self, input_bits):
 
         """
         mapfunc = vectorize(lambda i:
-                            self.constellation[bitarray2dec(input_bits[i:i + self.num_bits_symbol])])
+                            self._constellation[bitarray2dec(input_bits[i:i + self.num_bits_symbol])])
 
         baseband_symbols = mapfunc(arange(0, len(input_bits), self.num_bits_symbol))
 
@@ -80,9 +113,8 @@ def demodulate(self, input_symbols, demod_type, noise_var=0):
 
         """
         if demod_type == 'hard':
-            index_list = map(lambda i: argmin(abs(input_symbols[i] - self.constellation)),
-                             range(0, len(input_symbols)))
-            demod_bits = array([dec2bitarray(i, self.num_bits_symbol) for i in index_list]).reshape(-1)
+            index_list = abs(input_symbols - self._constellation[:, None]).argmin(0)
+            demod_bits = dec2bitarray(index_list, self.num_bits_symbol)
 
         elif demod_type == 'soft':
             demod_bits = zeros(len(input_symbols) * self.num_bits_symbol)
@@ -94,10 +126,10 @@ def demodulate(self, input_symbols, demod_type, noise_var=0):
                     for const_index in self.symbol_mapping:
                         if (const_index >> bit_index) & 1:
                             llr_num = llr_num + exp(
-                                (-abs(current_symbol - self.constellation[const_index]) ** 2) / noise_var)
+                                (-abs(current_symbol - self._constellation[const_index]) ** 2) / noise_var)
                         else:
                             llr_den = llr_den + exp(
-                                (-abs(current_symbol - self.constellation[const_index]) ** 2) / noise_var)
+                                (-abs(current_symbol - self._constellation[const_index]) ** 2) / noise_var)
                     demod_bits[i * self.num_bits_symbol + self.num_bits_symbol - 1 - bit_index] = log(llr_num / llr_den)
         else:
             raise ValueError('demod_type must be "hard" or "soft"')
@@ -142,36 +174,92 @@ def plot_constellation(self):
         plt.grid()
         plt.show()
 
+    @property
+    def constellation(self):
+        """ Constellation of the modem. """
+        return self._constellation
 
-class PSKModem(Modem):
-    """ Creates a Phase Shift Keying (PSK) Modem object. """
+    @constellation.setter
+    def constellation(self, value):
+        # Check value input
+        num_bits_symbol = log2(len(value))
+        if num_bits_symbol != int(num_bits_symbol):
+            raise ValueError('Constellation length must be a power of 2.')
 
-    Es = 1
+        # Set constellation as an array
+        self._constellation = array(value)
 
-    def _constellation_symbol(self, i):
-        return cos(2 * pi * (i - 1) / self.m) + sin(2 * pi * (i - 1) / self.m) * (0 + 1j)
+        # Update other attributes
+        self.Es = signal_power(self.constellation)
+        self.m = self._constellation.size
+        self.num_bits_symbol = int(num_bits_symbol)
+        self.constellation_mapping = arange(self.m)
 
-    def __init__(self, m):
-        """ Creates a Phase Shift Keying (PSK) Modem object.
+
+class PSKModem(Modem):
+    """ Creates a Phase Shift Keying (PSK) Modem object.
 
         Parameters
         ----------
         m : int
             Size of the PSK constellation.
 
+        Attributes
+        ----------
+        constellation : 1D-ndarray of complex
+                        Modem constellation. If changed, the length of the new constellation must be a power of 2.
+
+        Es            : float
+                        Average energy per symbols.
+
+        m             : integer
+                        Constellation length.
+
+        num_bits_symb : integer
+                        Number of bits per symbol.
+
+        Raises
+        ------
+        ValueError
+                        If the constellation is changed to an array-like with length that is not a power of 2.
         """
-        self.m = m
-        self.num_bits_symbol = int(log2(self.m))
-        self.symbol_mapping = arange(self.m)
-        self.constellation = list(map(self._constellation_symbol,
-                                      self.symbol_mapping))
+
+    def __init__(self, m):
+        """ Creates a Phase Shift Keying (PSK) Modem object. """
+
+        def _constellation_symbol(i):
+            return cos(2 * pi * (i - 1) / m) + sin(2 * pi * (i - 1) / m) * (0 + 1j)
+
+        self.constellation = list(map(_constellation_symbol, arange(m)))
 
 
 class QAMModem(Modem):
-    """ Creates a Quadrature Amplitude Modulation (QAM) Modem object."""
+    """ Creates a Quadrature Amplitude Modulation (QAM) Modem object.
 
-    def _constellation_symbol(self, i):
-        return (2 * i[0] - 1) + (2 * i[1] - 1) * (1j)
+        Parameters
+        ----------
+        m : int
+            Size of the PSK constellation.
+
+        Attributes
+        ----------
+        constellation : 1D-ndarray of complex
+                        Modem constellation. If changed, the length of the new constellation must be a power of 2.
+
+        Es            : float
+                        Average energy per symbols.
+
+        m             : integer
+                        Constellation length.
+
+        num_bits_symb : integer
+                        Number of bits per symbol.
+
+        Raises
+        ------
+        ValueError
+                        If the constellation is changed to an array-like with length that is not a power of 2.
+    """
 
     def __init__(self, m):
         """ Creates a Quadrature Amplitude Modulation (QAM) Modem object.
@@ -183,13 +271,12 @@ def __init__(self, m):
 
         """
 
-        self.m = m
-        self.num_bits_symbol = int(log2(self.m))
-        self.symbol_mapping = arange(self.m)
-        mapping_array = arange(1, sqrt(self.m) + 1) - (sqrt(self.m) / 2)
-        self.constellation = list(map(self._constellation_symbol,
+        def _constellation_symbol(i):
+            return (2 * i[0] - 1) + (2 * i[1] - 1) * (1j)
+
+        mapping_array = arange(1, sqrt(m) + 1) - (sqrt(m) / 2)
+        self.constellation = list(map(_constellation_symbol,
                                       list(product(mapping_array, repeat=2))))
-        self.Es = 2 * (self.m - 1) / 3
 
 
 def ofdm_tx(x, nfft, nsc, cp_length):
@@ -510,6 +597,11 @@ def bit_lvl_repr(H, w):
     ------
     A : 2D nbarray (shape : nb_rx, nb_tx*beta)
         Channel matrix adapted to the bit-level representation.
+
+    raises
+    ------
+    ValueError
+                    If beta (the length of w) is not even)
     """
     beta = len(w)
     if beta % 2 == 0:
@@ -518,7 +610,7 @@ def bit_lvl_repr(H, w):
         kr = kron(In, w)
         return dot(H, kr)
     else:
-        raise ValueError('Beta must be even.')
+        raise ValueError('Beta (length of w) must be even.')
 
 
 def max_log_approx(y, h, noise_var, pts_list, demode):

commpy/tests/test_modulation.py

@@ -1,13 +1,16 @@
 # Authors: CommPy contributors
 # License: BSD 3-Clause
 
+from itertools import product
+
 from numpy import zeros, identity, arange, concatenate, log2, array, inf
 from numpy.random import seed
-from numpy.testing import run_module_suite, assert_allclose, dec
+from numpy.testing import run_module_suite, assert_allclose, dec, assert_raises, assert_array_equal
 
 from commpy.channels import MIMOFlatChannel
 from commpy.links import *
-from commpy.modulation import QAMModem, mimo_ml, bit_lvl_repr, max_log_approx
+from commpy.modulation import QAMModem, mimo_ml, bit_lvl_repr, max_log_approx, PSKModem, Modem
+from commpy.utilities import signal_power
 
 
 @dec.slow
@@ -49,6 +52,10 @@ def receiver_without_blr(y, H, cons, noise_var):
     assert_allclose(ber_without_blr, ber_with_blr, rtol=0.5,
                     err_msg='bit_lvl_repr changes the performance')
 
+    # Test error raising
+    with assert_raises(ValueError):
+        bit_lvl_repr(RayleighChannel.channel_gains[0], array((2, 4, 6)))
+
 
 def test_max_log_approx():
     x = array((-1, -1, 1))
@@ -71,8 +78,68 @@ def decode(pt):
                     err_msg='Wrong LLRs without noise')
 
 
-def test_kbest():
-    pass  # Tested in test_links
+class ModemTestcase:
+    qam_modems = [QAMModem(4), QAMModem(16), QAMModem(64)]
+    psk_modems = [PSKModem(4), PSKModem(16), PSKModem(64)]
+    modems = qam_modems + psk_modems
+
+    def __init__(self):
+        # Create a custom Modem
+        custom_constellation = [re + im * 1j for re, im in product((-3.5, -0.5, 0.5, 3.5), repeat=2)]
+        self.custom_modems = [Modem(custom_constellation)]
+
+        # Add to custom modems a QAM modem with modified constellation
+        QAM_custom = QAMModem(16)
+        QAM_custom.constellation = custom_constellation
+        self.custom_modems.append(QAM_custom)
+        self.modems += self.custom_modems
+
+        # Assert that error is raised when the contellation length is not a power of 2
+        with assert_raises(ValueError):
+            QAM_custom.constellation = (0, 0, 0)
+
+    def test(self):
+        for modem in self.modems:
+            self.do(modem)
+        for modem in self.qam_modems:
+            self.do_qam(modem)
+        for modem in self.psk_modems:
+            self.do_psk(modem)
+        for modem in self.custom_modems:
+            self.do_custom(modem)
+
+    # Default methods for TestClasses that not implement a specific test
+    def do(self, modem):
+        pass
+
+    def do_qam(self, modem):
+        pass
+
+    def do_psk(self, modem):
+        pass
+
+    def do_custom(self, modem):
+        pass
+
+
+class TestModulateHardDemodulate(ModemTestcase):
+
+    def do(self, modem):
+        for bits in product(*((0, 1),) * modem.num_bits_symbol):
+            assert_array_equal(bits, modem.demodulate(modem.modulate(bits), 'hard'),
+                               err_msg='Bits are not equal after modulation and hard demodulation')
+
+
+class TestEs(ModemTestcase):
+
+    def do_qam(self, modem):
+        assert_allclose(signal_power(modem.constellation), 2 * (modem.m - 1) / 3)
+
+    def do_psk(self, modem):
+        assert_allclose(modem.Es, 1)
+
+    def do_custom(self, modem):
+        assert_allclose(modem.Es, 12.5)
 
 
 if __name__ == "__main__":

commpy/tests/test_utilities.py

@@ -0,0 +1,19 @@
+# Authors: CommPy contributors
+# License: BSD 3-Clause
+
+from numpy import array
+from numpy.random import seed
+from numpy.testing import run_module_suite, assert_array_equal
+
+from commpy.utilities import dec2bitarray
+
+
+def test_dec2bitarray():
+    # Assert result
+    assert_array_equal(dec2bitarray(17, 8), array((0, 0, 0, 1, 0, 0, 0, 1)))
+    assert_array_equal(dec2bitarray((17, 12), 5), array((1, 0, 0, 0, 1, 0, 1, 1, 0, 0)))
+
+
+if __name__ == "__main__":
+    seed(17121996)
+    run_module_suite()

commpy/utilities.py

@@ -16,21 +16,25 @@
    upsample             -- Upsample by an integral factor (zero insertion).
    signal_power         -- Compute the power of a discrete time signal.
 """
-from __future__ import division  # Python 2 compatibility
+
+import itertools as it
 
 import numpy as np
 
 __all__ = ['dec2bitarray', 'bitarray2dec', 'hamming_dist', 'euclid_dist', 'upsample',
            'signal_power']
 
+vectorized_binary_repr = np.vectorize(np.binary_repr)
+
+
 def dec2bitarray(in_number, bit_width):
     """
     Converts a positive integer to NumPy array of the specified size containing
     bits (0 and 1).
 
     Parameters
     ----------
-    in_number : int
+    in_number : int or array-like of int
         Positive integer to be converted to a bit array.
 
     bit_width : int
@@ -39,17 +43,12 @@ def dec2bitarray(in_number, bit_width):
     Returns
     -------
     bitarray : 1D ndarray of ints
-        Array containing the binary representation of the input decimal.
+        Array containing the binary representation of all the input decimal(s).
 
     """
 
-    binary_string = bin(in_number)
-    length = len(binary_string)
-    bitarray = np.zeros(bit_width, 'int')
-    for i in range(length - 2):
-        bitarray[bit_width - i - 1] = int(binary_string[length - i - 1])
-
-    return bitarray
+    binary_words = vectorized_binary_repr(np.array(in_number, ndmin=1), bit_width)
+    return np.fromiter(it.chain.from_iterable(binary_words), dtype=np.int8)
 
 
 def bitarray2dec(in_bitarray):