tests: Improve coverage for mapping submodule

Adds missing tests to the tests/ subdirectory in the test_mapping.py file.

Author: uwimt

poetry.lock

@@ -38,6 +38,7 @@ jaxtyping = "^0.2.36"
 pre-commit = "*"
 pytest = "*"
 pytest-coverage = "*"
+pytest-datafiles = "*"
 black = "*"
 flake8 = "*"
 pylint = "*"

pyproject.toml

@@ -1,6 +1,14 @@
 from mokka import mapping
+import mokka.utils.bitops.torch
+import mokka.utils.bitops.numpy
 import torch
 import numpy as np
+import pytest
+
+from pathlib import Path
+
+FIXTURE_DIR = Path(__file__).parent.resolve() / "data"
+
 
 # PyTorch tests
 
@@ -43,6 +51,56 @@ def test_qam_constellation_mapper():
     assert np.allclose(symbols, reference_symbols)
 
 
+def test_custom_constellation_mapper():
+    m = 4
+    bits = mokka.utils.bitops.numpy.generate_all_input_bits(m)
+    symbols = mokka.utils.bitops.torch.bits2idx(torch.from_numpy(bits))
+    mapper = mapping.torch.CustomConstellationMapper(m, symbols)
+    reference_symbols = mapper.get_constellation().detach().numpy().flatten()
+    assert np.allclose(symbols, reference_symbols)
+
+
+def test_separated_constellation_mapper():
+    m = 4
+    for i in range(0, 3):
+        if i == 0:
+            mapper = mapping.torch.SeparatedConstellationMapper(m, qam_init=True)
+            symbols = mapper.get_constellation().detach().numpy().flatten()
+            reference_symbols = mapping.numpy.QAM(m).get_constellation().flatten()
+            assert np.allclose(symbols, reference_symbols)
+        elif i == 1:
+            mapper = mapping.torch.SeparatedConstellationMapper(m, m_real=1)
+            symbols = mapper.get_constellation()
+            assert symbols.shape[0] == 16
+        elif i == 2:
+            mapper = mapping.torch.SeparatedConstellationMapper(m, m_imag=1)
+            symbols = mapper.get_constellation()
+            assert symbols.shape[0] == 16
+
+
+@pytest.mark.datafiles(
+    FIXTURE_DIR / "AWGN_16QAM_mapper.bin", FIXTURE_DIR / "AWGN_16QAM_demapper.bin"
+)
+def test_constellation_demapper(datafiles):
+    m = 4
+    bits = torch.from_numpy(
+        np.array([[0, 1, 0, 0], [1, 0, 0, 0], [1, 1, 1, 1], [0, 1, 0, 1]], dtype=int)
+    )
+    for df in datafiles.iterdir():
+        if "demapper" in df.name:
+            demapper_dict = torch.load(df, map_location=torch.device("cpu"))
+        elif "mapper" in df.name:
+            mapper_dict = torch.load(df, map_location=torch.device("cpu"))
+        else:
+            raise ValueError("Neither mapper nor demapper in filename")
+    mapper = mapping.torch.ConstellationMapper(m).load_model(mapper_dict)
+    symbols = mapper(bits).flatten()
+    demapper = mapping.torch.ConstellationDemapper(m).load_model(demapper_dict)
+    llrs = demapper(symbols)
+    rx_bits = (llrs.detach().numpy() < 0).astype(int)
+    assert np.allclose(bits, rx_bits)
+
+
 def test_classical_demapper():
     m = 4
     sigma = torch.tensor(0.1)
@@ -59,6 +117,24 @@ def test_classical_demapper():
     assert np.allclose(bits, rx_bits)
 
 
+def test_classical_demapper_symbolwise():
+    m = 4
+    sigma = torch.tensor(0.1)
+    bits = torch.from_numpy(
+        np.array([[0, 1, 0, 0], [1, 0, 0, 0], [1, 1, 1, 1], [0, 1, 0, 1]], dtype=int)
+    )
+    onehot = mokka.utils.bitops.torch.bits_to_onehot(bits)
+    mapper = mapping.torch.QAMConstellationMapper(m)
+    symbols = mapper(bits).flatten()
+    demapper = mapping.torch.ClassicalDemapper(
+        sigma, mapper.get_constellation().flatten(), bitwise=False
+    )
+    q_value = demapper(symbols)
+    print(q_value)
+    rx_onehot = (q_value.detach().numpy() >= 1.0000e-5).astype(float)
+    assert np.allclose(onehot, rx_onehot)
+
+
 def test_gaussian_demapper():
     m = 4
     bits = torch.from_numpy(
@@ -73,3 +149,25 @@ def test_gaussian_demapper():
     print(bits)
     print(rx_bits)
     assert np.allclose(bits, rx_bits)
+
+
+def test_separated_simple_demapper():
+    m = 4
+    bits = torch.from_numpy(
+        np.array([[0, 1, 0, 0], [1, 0, 0, 0], [1, 1, 1, 1], [0, 1, 0, 1]], dtype=int)
+    )
+    mapper = mapping.torch.QAMConstellationMapper(m)
+    symbols = mapper(bits).flatten()
+    demapper = mapping.torch.SeparatedSimpleDemapper(m, demapper_width=128)
+    llrs = demapper(symbols)
+    rx_bits = (llrs.detach().numpy() < 0).astype(int)
+    assert not np.allclose(bits, rx_bits)
+
+
+def test_pcssampler():
+    m = 4
+    batchsize = 10
+    sampler = mapping.torch.PCSSampler(m)
+    indices = sampler(batchsize).detach().numpy()
+    print(indices)
+    assert len(indices) == batchsize and max(indices) << 2 ^ m and min(indices) >= 0