[fix] fix nrCodeBlockDesegmentLDPC()

catkira · Aug 23, 2023 · 69f7585 · 69f7585
1 parent bc27312
commit 69f7585
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/py3gpp/nrCodeBlockDesegmentLDPC.py b/py3gpp/nrCodeBlockDesegmentLDPC.py
@@ -7,20 +7,20 @@
 def nrCodeBlockDesegmentLDPC(cbs, bgn, blklen):
     blk = None
     err = False
-    if len(cbs.shape) == 1:
+    if len(cbs.shape) == 1 or cbs.shape[1] == 1:
         blk = cbs[:blklen]
         # there is no appended CRC when there is only one segment
     else:
         cbsInfo = getCBSInfo(blklen, bgn)
         idx = 0
-        blk = np.zeros(blklen)
+        blk = np.zeros((blklen, 1))
         for i in np.arange(cbsInfo['C']):
             if i < cbsInfo['C'] - 1:
-                blk[idx:][:cbsInfo['CBZ']] = cbs[0 : cbsInfo['CBZ'], i]
+                blk[idx:, 0][:cbsInfo['CBZ']] = cbs[0 : cbsInfo['CBZ'], i]
                 _, crc = nrCRCDecode(cbs[0 : cbsInfo['CBZ'] + cbsInfo['Lcb'], i], '24B')
                 idx += cbsInfo['CBZ']
             else:
-                blk[idx:][:blklen - idx] = cbs[0 : blklen - idx, i]
+                blk[idx:, 0][:blklen - idx] = cbs[0 : blklen - idx, i]
                 _, crc = nrCRCDecode(cbs[0 : cbsInfo['CBZ'] + cbsInfo['Lcb'], i], '24B')
 
             if crc != 0:

diff --git a/tests/matlab/test_nrCodeBlockDesegmentLDPC.py b/tests/matlab/test_nrCodeBlockDesegmentLDPC.py
@@ -12,20 +12,21 @@
 def run_nrCodeBlockDesegmentLDPC_fixed_data(eng):
     bgn = 2
     blklen = 640 + 16 # 16 bit transport block CRC
+    cbs = test_data.ldpc.decBits
+    cbs = np.expand_dims(cbs, axis=1)
     ref_data = eng.nrCodeBlockDesegmentLDPC(eng.transpose(eng.double(test_data.ldpc.decBits)), eng.double(bgn), eng.double(blklen))
-    ref_data = np.array(list(itertools.chain(*ref_data)))
-
-    data, _ = nrCodeBlockDesegmentLDPC(test_data.ldpc.decBits, bgn, blklen)
+    ref_data = np.asarray(ref_data)
 
+    data, _ = nrCodeBlockDesegmentLDPC(cbs, bgn, blklen)
     assert (ref_data == data).all()
 
 def run_nrCodeBlockDesegmentLDPC(blklen, eng):
     bgn = 2
-    in_blk = np.random.randint(2, size = blklen)  # this is a transport block
-    cbs = eng.nrCodeBlockSegmentLDPC(eng.transpose(eng.double(in_blk)), eng.double(bgn))  # outputs code block segments
+    in_blk = np.random.randint(2, size = (blklen, 1))  # this is a transport block
+    cbs = eng.nrCodeBlockSegmentLDPC(eng.double(in_blk), eng.double(bgn))  # outputs code block segments
     cbs = np.asarray(cbs)
     ref_blk = eng.nrCodeBlockDesegmentLDPC(eng.double(cbs), eng.double(bgn), eng.double(blklen))
-    ref_blk = np.ravel(np.asarray(ref_blk).astype(int))
+    ref_blk = np.asarray(ref_blk).astype(int)
     assert np.array_equal(ref_blk, in_blk)
 
     out_blk, _ = nrCodeBlockDesegmentLDPC(cbs, bgn, blklen)
@@ -40,11 +41,12 @@ def eng():
 def test_nrCodeBlockDesegmentLDPC_fixed_data(eng):
     run_nrCodeBlockDesegmentLDPC_fixed_data(eng)
 
-@pytest.mark.parametrize('blklen', [4000, 4001, 10000])
+@pytest.mark.parametrize('blklen', [100, 1000, 4000, 4001, 10000])
 def test_nrCodeBlockDesegmentLDPC(blklen, eng):
     run_nrCodeBlockDesegmentLDPC(blklen, eng)
 
 if __name__ == '__main__':
     _eng = matlab.engine.connect_matlab()
     run_nrCodeBlockDesegmentLDPC(4000, _eng)
+    # run_nrCodeBlockDesegmentLDPC_fixed_data(_eng)
     _eng.quit()