1#!/usr/bin/env python
2#
3# Copyright 2015 Free Software Foundation, Inc.
4#
5# This file is part of GNU Radio
6#
7# GNU Radio is free software; you can redistribute it and/or modify
8# it under the terms of the GNU General Public License as published by
9# the Free Software Foundation; either version 3, or (at your option)
10# any later version.
11#
12# GNU Radio is distributed in the hope that it will be useful,
13# but WITHOUT ANY WARRANTY; without even the implied warranty of
14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15# GNU General Public License for more details.
16#
17# You should have received a copy of the GNU General Public License
18# along with GNU Radio; see the file COPYING.  If not, write to
19# the Free Software Foundation, Inc., 51 Franklin Street,
20# Boston, MA 02110-1301, USA.
21#
22
23from __future__ import absolute_import
24from __future__ import division
25
26
27import numpy as np
28
29from gnuradio import gr, gr_unittest, blocks, fec
30from gnuradio.fec.extended_encoder import extended_encoder
31from gnuradio.fec.polar.encoder import PolarEncoder
32from gnuradio.fec.polar import channel_construction as cc
33
34# import os
35# print('PID:', os.getpid())
36# raw_input('tell me smth')
37
38
39class test_polar_encoder(gr_unittest.TestCase):
40
41    def setUp(self):
42        self.tb = gr.top_block()
43
44    def tearDown(self):
45        self.tb = None
46
47    def test_001_setup(self):
48        block_size = 16
49        num_info_bits = 8
50        frozen_bit_positions = cc.frozen_bit_positions(block_size, num_info_bits, 0.0)
51        frozen_bit_values = np.array([],)
52
53        polar_encoder = fec.polar_encoder.make(block_size, num_info_bits, frozen_bit_positions, frozen_bit_values)
54
55        self.assertEqual(block_size, polar_encoder.get_output_size())
56        self.assertEqual(num_info_bits, polar_encoder.get_input_size())
57        self.assertFloatTuplesAlmostEqual((float(num_info_bits) / block_size, ), (polar_encoder.rate(), ))
58        self.assertFalse(polar_encoder.set_frame_size(10))
59
60    def test_002_work_function_packed(self):
61        is_packed = True
62        block_size = 256
63        num_info_bits = block_size // 2
64
65        data, ref, polar_encoder = self.get_test_data(block_size, num_info_bits, 1, is_packed)
66        src = blocks.vector_source_b(data, False)
67        enc_block = extended_encoder(polar_encoder, None, '11')
68        snk = blocks.vector_sink_b(1)
69
70        self.tb.connect(src, enc_block, snk)
71        self.tb.run()
72
73        res = np.array(snk.data()).astype(dtype=int)
74        self.assertTupleEqual(tuple(res), tuple(ref))
75
76    def test_003_work_function_unpacked(self):
77        is_packed = False
78        block_size = 256
79        num_info_bits = block_size // 2
80
81        data, ref, polar_encoder = self.get_test_data(block_size, num_info_bits, 1, is_packed)
82        src = blocks.vector_source_b(data, False)
83        enc_block = extended_encoder(polar_encoder, None, '11')
84        snk = blocks.vector_sink_b(1)
85
86        self.tb.connect(src, enc_block, snk)
87        self.tb.run()
88
89        res = np.array(snk.data()).astype(dtype=int)
90        self.assertTupleEqual(tuple(res), tuple(ref))
91
92    def test_004_big_input(self):
93        is_packed = False
94        num_blocks = 30
95        block_size = 1024
96        num_info_bits = block_size // 8
97
98        data, ref, polar_encoder = self.get_test_data(block_size, num_info_bits, num_blocks, is_packed)
99        src = blocks.vector_source_b(data, False)
100        enc_block = extended_encoder(polar_encoder, None, '11')
101        snk = blocks.vector_sink_b(1)
102
103        self.tb.connect(src, enc_block, snk)
104        self.tb.run()
105
106        res = np.array(snk.data()).astype(dtype=int)
107        self.assertTupleEqual(tuple(res), tuple(ref))
108
109    def get_test_data(self, block_size, num_info_bits, num_blocks, is_packed):
110        # helper function to set up test data and together with encoder object.
111        num_frozen_bits = block_size - num_info_bits
112        frozen_bit_positions = cc.frozen_bit_positions(block_size, num_info_bits, 0.0)
113        frozen_bit_values = np.array([0] * num_frozen_bits,)
114        python_encoder = PolarEncoder(block_size, num_info_bits, frozen_bit_positions, frozen_bit_values)
115
116        data = np.array([], dtype=int)
117        ref = np.array([], dtype=int)
118        for i in range(num_blocks):
119            d = np.random.randint(2, size=num_info_bits)
120            data = np.append(data, d)
121            ref = np.append(ref, python_encoder.encode(d))
122        polar_encoder = fec.polar_encoder.make(block_size, num_info_bits, frozen_bit_positions, frozen_bit_values, is_packed)
123        return data, ref, polar_encoder
124
125
126if __name__ == '__main__':
127    gr_unittest.run(test_polar_encoder, "test_polar_encoder.xml")
128