<filename>test/test_scrambling.py
#
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
try:
    import sionna
except ImportError as e:
    import sys
    sys.path.append("../")

import unittest
import numpy as np
import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
print('Number of GPUs available :', len(gpus))
if gpus:
    gpu_num = 0 # Number of the GPU to be used
    try:
        tf.config.set_visible_devices(gpus[gpu_num], 'GPU')
        print('Only GPU number', gpu_num, 'used.')
        tf.config.experimental.set_memory_growth(gpus[gpu_num], True)
    except Runtime as e:
        print(e)
from sionna.fec.scrambling import Descrambler, Scrambler
from sionna.utils import BinarySource

class TestScrambler(unittest.TestCase):

    def test_sequence_dimension(self):
        """Test against correct dimensions of the sequence"""
        seq_lengths = [1, 100, 256, 1000, 1e4]
        batch_sizes = [1, 100, 256, 1000, 1e4]

        # keep_State=True
        for seq_length in seq_lengths:
            # init new scrambler for new sequence size;
            # only different batch_sizes are allowed in this mode
            s = Scrambler(binary=False)
            for batch_size in batch_sizes:
                llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                         tf.cast(seq_length, dtype=tf.int32)])
                # build scrambler
                x = s(llr).numpy()
                self.assertTrue(np.array_equal(np.array(x.shape),
                                 [int(batch_size), int(seq_length)]))
        # keep_State=False
        s = Scrambler(binary=False, keep_state=False)
        for seq_length in seq_lengths:
            for batch_size in batch_sizes:
                llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                         tf.cast(seq_length, dtype=tf.int32)])
                # build scrambler
                x = s(llr).numpy()
                self.assertTrue(np.array_equal(np.array(x.shape),
                                 [int(batch_size), int(seq_length)]))

    def test_sequence_offset(self):
        """Test that scrambling sequence has no offset, i.e., equal likely 0s
        and 1s"""
        seq_length = int(1e4)
        batch_size = int(1e2)
        for seed in (None, 1337, 1234, 1003): # test some initial seeds
            for keep_state in (False, True):
                s = Scrambler(seed=seed, keep_state=keep_state, binary=True)
                llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                         tf.cast(seq_length, dtype=tf.int32)])
                # build scrambler
                s(llr)
                # generate a random sequence
                x = s(tf.zeros_like(llr))
                self.assertAlmostEqual(np.mean(x),
                                        0.5,
                                        places=2)

    def test_sequence_batch(self):
        """Test that scrambling sequence is random per batch sample iff
        keep_batch_dims=True."""

        seq_length = int(1e6)
        batch_size = int(1e1)
        llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                 tf.cast(seq_length, dtype=tf.int32)])

        for keep_state in (False, True):
            s = Scrambler(keep_batch_constant=False,
                          keep_state=keep_state,
                          binary=True)
            # generate a random sequence
            x = s(tf.zeros_like(llr))
            for i in range(batch_size-1):
                for j in range(i+1,batch_size):
                    # each batch sample must be different
                    self.assertAlmostEqual(np.mean(np.abs(x[i,:]-x[j,:])),
                                           0.5,
                                           places=2)
        # test that the pattern is the same of option keep_batch_constant==True
        for keep_state in (False, True):
            s = Scrambler(keep_batch_constant=True,
                          keep_state=keep_state,
                          binary=True)
            # generate a random sequence
            x = s(tf.zeros_like(llr))
            for i in range(batch_size-1):
                for j in range(i+1,batch_size):
                    # each batch sample is the same
                    self.assertTrue(np.sum(np.abs(x[i,:]-x[j,:]))==0)

    def test_sequence_realization(self):
        """Test that scrambling sequences are random for each new realization.
        """

        seq_length = int(1e5)
        batch_size = int(1e2)
        s = Scrambler(keep_state=False, binary=True)
        llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                 tf.cast(seq_length, dtype=tf.int32)])
        # generate a random sequence
        x1 = s(tf.zeros_like(llr))
        x2 = s(tf.zeros_like(llr))
        self.assertAlmostEqual(np.mean(np.abs(x1-x2)), 0.5, places=3)

    def test_inverse(self):
        """Test that scrambling can be inverted/removed.
        2x scrambling must result in the original sequence (for binary and
         LLRs).
        """
        seq_length = int(1e5)
        batch_size = int(1e2)

        #check binary scrambling
        b = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                               tf.cast(seq_length, dtype=tf.int32)],
                               minval=0,
                               maxval=1)
        for keep_batch in (False, True):
            s = Scrambler(binary=True,
                          keep_batch_constant=keep_batch,
                          keep_state=True)
            # only works if keep_state=True
            b = tf.cast(tf.greater(0.5, b), dtype=tf.float32)
            x = s(b)
            x = s(x)
            self.assertIsNone(np.testing.assert_array_equal(x.numpy(),
                                                            b.numpy()))
            #check soft-value scrambling (flip sign)
            s = Scrambler(binary=False,
                          keep_batch_constant=keep_batch,
                          keep_state=True)
            llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                     tf.cast(seq_length, dtype=tf.int32)])
            x = s(llr)
            x = s(x)
            self.assertIsNone(np.testing.assert_array_equal(x.numpy(),
                                                            llr.numpy()))

    def test_llr(self):
        """Test that scrambling works for soft-values (sign flip)."""
        s = Scrambler(binary=False, seed=12345)
        b = tf.ones([100,200])
        x = s(b)
        s2 = Scrambler(binary=True, seed=12345)
        res = -2. * s2(tf.zeros_like(x)) + 1
        self.assertIsNone(np.testing.assert_array_equal(x.numpy(), res.numpy()))


    def test_keep_state(self):
        """Test that keep_state works as expected.
        Iff keep_state==True, the scrambled sequences must be constant."""
        seq_length = int(1e5)
        batch_size = int(1e2)
        llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                 tf.cast(seq_length, dtype=tf.int32)],
                                minval=-100,
                                maxval=100)
        S = Scrambler(binary=True, keep_state=True)
        res1 = S(tf.zeros_like(llr))
        res2 = S(tf.zeros_like(llr))

        self.assertTrue(np.array_equal(res1.numpy(), res2.numpy()))

        # also check that the sequence is unique with keep_state=False
        S = Scrambler(binary=True, keep_state=False)
        _ = S(llr)
        res1 = S(tf.zeros_like(llr))
        _ = S(llr)
        res2 = S(tf.zeros_like(llr))

        self.assertFalse(np.array_equal(res1.numpy(), res2.numpy()))


    def test_keras(self):
        """Test that Keras model can be compiled (supports dynamic shapes)."""
        bs = 10
        k = 100
        source = BinarySource()

        inputs = tf.keras.Input(shape=(k), dtype=tf.float32)
        x = Scrambler()(inputs)
        model = tf.keras.Model(inputs=inputs, outputs=x)
        # test that output batch dim is none
        self.assertTrue(model.output_shape[0] is None)

        # test that model can be called
        b = source([bs,k])
        model(b)
        # call twice to see that bs can change
        b2 = source([bs+1,k])
        model(b2)
        model.summary()

    def test_tf_fun(self):
        """Test that graph mode and XLA works as expected"""

        @tf.function()
        def run_graph(llr):
            return s(llr)

        @tf.function(jit_compile=True)
        def run_graph_xla(llr):
            return s(llr)

        for keep_state in (False, True):
            s = Scrambler(keep_state=keep_state)
            b = tf.ones([100,200])
            x1 = run_graph(b)
            x2 = run_graph_xla(b)
            # again with different batch_size
            b = tf.ones([101,200])
            x1 = run_graph(b)
            x2 = run_graph_xla(b)
            # and different sequence length
            b = tf.ones([101,201])
            x1 = run_graph(b)
            x2 = run_graph_xla(b)

            self.assertTrue(np.any(np.not_equal(x1.numpy(),b.numpy())))
            self.assertTrue(np.any(np.not_equal(x2.numpy(),b.numpy())))

    def test_seed(self):
        """Test that seed generates reproducible results."""
        seq_length = int(1e5)
        batch_size = int(1e2)
        b = tf.zeros([batch_size, seq_length])

        s1 = Scrambler(seed=1337, binary=True, keep_state=False)
        res_s1_1 = s1(b)
        res_s1_2 = s1(b)
        # new realization per call
        self.assertFalse(np.array_equal(res_s1_1.numpy(), res_s1_2.numpy()))

        # if keep_state=True, the same seed should lead to the same sequence
        s2 = Scrambler(seed=1337, binary=True, keep_state=True)
        res_s2_1 = s2(b)
        s3 = Scrambler(seed=1337)
        res_s3_1 = s3(b)
        # same seed lead to same sequence
        self.assertTrue(np.array_equal(res_s2_1.numpy(), res_s3_1.numpy()))

         # but with random seed it gives a new sequence for each init
        s4 = Scrambler(seed=None, binary=True, keep_state=True)
        res_s4_1 = s2(b)
        s5 = Scrambler(seed=None)
        res_s5_1 = s5(b)
        # same seed lead to same sequence
        self.assertFalse(np.array_equal(res_s4_1.numpy(), res_s5_1.numpy()))

        # for keep_State=False, even the same seed leads to new results
        s6 = Scrambler(seed=1337, binary=True, keep_state=False)
        res_s6_1 = s6(b)
        # different seed generates new sequence
        self.assertFalse(np.array_equal(res_s6_1.numpy(), res_s2_1.numpy()))

        # init with same seed as previous random seed
        s7 = Scrambler(seed=None, binary=True, keep_state=True)
        res_s7_1 = s7(b)
        s8 = Scrambler(seed=s7.seed, binary=True, keep_state=True)
        res_s8_1 = s8(b)
        # same seed lead to same sequence
        self.assertTrue(np.array_equal(res_s7_1.numpy(), res_s8_1.numpy()))

        # test that seed can be also provided to call
        seed = 987654
        s9 = Scrambler(seed=45234, keep_state=False)
        s10 = Scrambler(seed=76543, keep_state=True)
        x1 = s9([b, seed]).numpy()
        x2 = s9([b, seed+1]).numpy()
        x3 = s9([b, seed]).numpy()
        x4 = s10([b, seed]).numpy()
        self.assertFalse(np.array_equal(x1, x2)) # different seed
        self.assertTrue(np.array_equal(x1, x3)) # same seed
        self.assertTrue(np.array_equal(x1, x4)) # same seed (keep_state=f)

        # test that random seed allows inverse
        x5 = s9([b, seed])
        x6 = s9([b, seed]).numpy()
        # same seed
        self.assertTrue(np.array_equal(x5, x6)) # identity
        # different seed
        x7 = s9([b, seed+1])
        self.assertFalse(np.array_equal(x5, x7)) # identity
        # same seed again
        x8 = s9([b, seed+1])
        self.assertTrue(np.array_equal(x7, x8)) # identity

    def test_dtype(self):
        """Test that variable dtypes are supported."""
        seq_length = int(1e1)
        batch_size = int(1e2)

        dt_supported = [tf.float16, tf.float32, tf.float64]
        for dt in dt_supported:
            for dt_in in dt_supported:
                for dt_out in dt_supported:
                    b = tf.zeros([batch_size, seq_length], dtype=dt_in)
                    s1 = Scrambler(dtype=dt)
                    s2 = Descrambler(s1, dtype=dt_out)
                    x = s1(b)
                    y = s2(x)
                    assert (x.dtype==dt)
                    assert (y.dtype==dt_out)

    def test_descrambler(self):
        """"Test that descrambler works as expected."""
        seq_length = int(1e2)
        batch_size = int(1e1)

        b = tf.zeros([batch_size, seq_length])
        s1 = Scrambler()
        s2 = Descrambler(s1)
        x = s1(b)
        y = s2(x)
        assert (np.array_equal(b.numpy(), y.numpy()))

        x = s1([b, 1234])
        y = s2(x)
        assert (not np.array_equal(b.numpy(), y.numpy()))

        # check if seed is correctly retrieved from scrambler
        s3 = Scrambler(seed=12345)
        s4 = Descrambler(s3)
        x = s3(b)
        y = s4(x)
        assert (np.array_equal(b.numpy(), y.numpy()))

    def test_descrambler_nonbin(self):
        """"Test that descrambler works with non-binary."""
        seq_length = int(1e2)
        batch_size = int(1e1)

        b = tf.zeros([batch_size, seq_length])

        # scrambler binary, but descrambler non-binary
        scrambler = Scrambler(seed=1235456, binary=True)
        descrambler = Descrambler(scrambler, binary=False)
        # with explicit seed
        s = 8764
        y = scrambler([b, s])
        z = descrambler([2*y-1, s]) # bspk
        z = 1 + z # remove bpsk
        assert (np.array_equal(b.numpy(), z.numpy()))
        #without explicit seed
        y = scrambler(b)
        z = descrambler(2*y-1) # bspk
        z = 1 + z # remove bpsk
        assert (np.array_equal(b.numpy(), z.numpy()))

        # scrambler non-binary, but descrambler
        scrambler = Scrambler(seed=1235456, binary=False)
        descrambler = Descrambler(scrambler, binary=True)
        s = 546342
        y = scrambler([2*b-1, s]) # bspk
        y = 0.5*(1 + y) # remove bpsk
        z = descrambler([y, s])
        assert (np.array_equal(b.numpy(), z.numpy()))
        #without explicit seed
        y = scrambler(2*b-1) # bspk
        y = 0.5*(1 + y) # remove bpsk
        z = descrambler(y)
        y = 1 + y # remove bpsk
        assert (np.array_equal(b.numpy(), z.numpy()))

    def test_scrambler_binary(self):
        """test that binary flag can be used as input"""
        seq_length = int(1e2)
        batch_size = int(1e1)

        b = tf.ones([batch_size, seq_length])

        # scrambler binary, but descrambler non-binary

        scrambler = Scrambler(seed=1245, binary=True)

        s = 1234
        x1 = scrambler(b) # binary scrambling
        x2 = scrambler([b, s]) # binary scrambling different seed
        x3 = scrambler([b, s, True]) # binary scrambling different seed
        x4 = scrambler([b, s, False]) # non-binary scrambling different seed

        assert (not np.array_equal(x1.numpy(), x2.numpy())) # different seed
        assert (np.array_equal(x2.numpy(), x3.numpy())) # same seed
        # same but "bpsk modulated"
        assert (not np.array_equal(x1.numpy(), 0.5*(1+x4.numpy())))

<filename>test/test_scrambling.py
#
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
try:
    import sionna
except ImportError as e:
    import sys
    sys.path.append("../")

import unittest
import numpy as np
import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
print('Number of GPUs available :', len(gpus))
if gpus:
    gpu_num = 0 # Number of the GPU to be used
    try:
        tf.config.set_visible_devices(gpus[gpu_num], 'GPU')
        print('Only GPU number', gpu_num, 'used.')
        tf.config.experimental.set_memory_growth(gpus[gpu_num], True)
    except Runtime as e:
        print(e)
from sionna.fec.scrambling import Descrambler, Scrambler
from sionna.utils import BinarySource

class TestScrambler(unittest.TestCase):

    def test_sequence_dimension(self):
        """Test against correct dimensions of the sequence"""
        seq_lengths = [1, 100, 256, 1000, 1e4]
        batch_sizes = [1, 100, 256, 1000, 1e4]

        # keep_State=True
        for seq_length in seq_lengths:
            # init new scrambler for new sequence size;
            # only different batch_sizes are allowed in this mode
            s = Scrambler(binary=False)
            for batch_size in batch_sizes:
                llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                         tf.cast(seq_length, dtype=tf.int32)])
                # build scrambler
                x = s(llr).numpy()
                self.assertTrue(np.array_equal(np.array(x.shape),
                                 [int(batch_size), int(seq_length)]))
        # keep_State=False
        s = Scrambler(binary=False, keep_state=False)
        for seq_length in seq_lengths:
            for batch_size in batch_sizes:
                llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                         tf.cast(seq_length, dtype=tf.int32)])
                # build scrambler
                x = s(llr).numpy()
                self.assertTrue(np.array_equal(np.array(x.shape),
                                 [int(batch_size), int(seq_length)]))

    def test_sequence_offset(self):
        """Test that scrambling sequence has no offset, i.e., equal likely 0s
        and 1s"""
        seq_length = int(1e4)
        batch_size = int(1e2)
        for seed in (None, 1337, 1234, 1003): # test some initial seeds
            for keep_state in (False, True):
                s = Scrambler(seed=seed, keep_state=keep_state, binary=True)
                llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                         tf.cast(seq_length, dtype=tf.int32)])
                # build scrambler
                s(llr)
                # generate a random sequence
                x = s(tf.zeros_like(llr))
                self.assertAlmostEqual(np.mean(x),
                                        0.5,
                                        places=2)

    def test_sequence_batch(self):
        """Test that scrambling sequence is random per batch sample iff
        keep_batch_dims=True."""

        seq_length = int(1e6)
        batch_size = int(1e1)
        llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                 tf.cast(seq_length, dtype=tf.int32)])

        for keep_state in (False, True):
            s = Scrambler(keep_batch_constant=False,
                          keep_state=keep_state,
                          binary=True)
            # generate a random sequence
            x = s(tf.zeros_like(llr))
            for i in range(batch_size-1):
                for j in range(i+1,batch_size):
                    # each batch sample must be different
                    self.assertAlmostEqual(np.mean(np.abs(x[i,:]-x[j,:])),
                                           0.5,
                                           places=2)
        # test that the pattern is the same of option keep_batch_constant==True
        for keep_state in (False, True):
            s = Scrambler(keep_batch_constant=True,
                          keep_state=keep_state,
                          binary=True)
            # generate a random sequence
            x = s(tf.zeros_like(llr))
            for i in range(batch_size-1):
                for j in range(i+1,batch_size):
                    # each batch sample is the same
                    self.assertTrue(np.sum(np.abs(x[i,:]-x[j,:]))==0)

    def test_sequence_realization(self):
        """Test that scrambling sequences are random for each new realization.
        """

        seq_length = int(1e5)
        batch_size = int(1e2)
        s = Scrambler(keep_state=False, binary=True)
        llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                 tf.cast(seq_length, dtype=tf.int32)])
        # generate a random sequence
        x1 = s(tf.zeros_like(llr))
        x2 = s(tf.zeros_like(llr))
        self.assertAlmostEqual(np.mean(np.abs(x1-x2)), 0.5, places=3)

    def test_inverse(self):
        """Test that scrambling can be inverted/removed.
        2x scrambling must result in the original sequence (for binary and
         LLRs).
        """
        seq_length = int(1e5)
        batch_size = int(1e2)

        #check binary scrambling
        b = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                               tf.cast(seq_length, dtype=tf.int32)],
                               minval=0,
                               maxval=1)
        for keep_batch in (False, True):
            s = Scrambler(binary=True,
                          keep_batch_constant=keep_batch,
                          keep_state=True)
            # only works if keep_state=True
            b = tf.cast(tf.greater(0.5, b), dtype=tf.float32)
            x = s(b)
            x = s(x)
            self.assertIsNone(np.testing.assert_array_equal(x.numpy(),
                                                            b.numpy()))
            #check soft-value scrambling (flip sign)
            s = Scrambler(binary=False,
                          keep_batch_constant=keep_batch,
                          keep_state=True)
            llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                     tf.cast(seq_length, dtype=tf.int32)])
            x = s(llr)
            x = s(x)
            self.assertIsNone(np.testing.assert_array_equal(x.numpy(),
                                                            llr.numpy()))

    def test_llr(self):
        """Test that scrambling works for soft-values (sign flip)."""
        s = Scrambler(binary=False, seed=12345)
        b = tf.ones([100,200])
        x = s(b)
        s2 = Scrambler(binary=True, seed=12345)
        res = -2. * s2(tf.zeros_like(x)) + 1
        self.assertIsNone(np.testing.assert_array_equal(x.numpy(), res.numpy()))


    def test_keep_state(self):
        """Test that keep_state works as expected.
        Iff keep_state==True, the scrambled sequences must be constant."""
        seq_length = int(1e5)
        batch_size = int(1e2)
        llr = tf.random.uniform([tf.cast(batch_size, dtype=tf.int32),
                                 tf.cast(seq_length, dtype=tf.int32)],
                                minval=-100,
                                maxval=100)
        S = Scrambler(binary=True, keep_state=True)
        res1 = S(tf.zeros_like(llr))
        res2 = S(tf.zeros_like(llr))

        self.assertTrue(np.array_equal(res1.numpy(), res2.numpy()))

        # also check that the sequence is unique with keep_state=False
        S = Scrambler(binary=True, keep_state=False)
        _ = S(llr)
        res1 = S(tf.zeros_like(llr))
        _ = S(llr)
        res2 = S(tf.zeros_like(llr))

        self.assertFalse(np.array_equal(res1.numpy(), res2.numpy()))


    def test_keras(self):
        """Test that Keras model can be compiled (supports dynamic shapes)."""
        bs = 10
        k = 100
        source = BinarySource()

        inputs = tf.keras.Input(shape=(k), dtype=tf.float32)
        x = Scrambler()(inputs)
        model = tf.keras.Model(inputs=inputs, outputs=x)
        # test that output batch dim is none
        self.assertTrue(model.output_shape[0] is None)

        # test that model can be called
        b = source([bs,k])
        model(b)
        # call twice to see that bs can change
        b2 = source([bs+1,k])
        model(b2)
        model.summary()

    def test_tf_fun(self):
        """Test that graph mode and XLA works as expected"""

        @tf.function()
        def run_graph(llr):
            return s(llr)

        @tf.function(jit_compile=True)
        def run_graph_xla(llr):
            return s(llr)

        for keep_state in (False, True):
            s = Scrambler(keep_state=keep_state)
            b = tf.ones([100,200])
            x1 = run_graph(b)
            x2 = run_graph_xla(b)
            # again with different batch_size
            b = tf.ones([101,200])
            x1 = run_graph(b)
            x2 = run_graph_xla(b)
            # and different sequence length
            b = tf.ones([101,201])
            x1 = run_graph(b)
            x2 = run_graph_xla(b)

            self.assertTrue(np.any(np.not_equal(x1.numpy(),b.numpy())))
            self.assertTrue(np.any(np.not_equal(x2.numpy(),b.numpy())))

    def test_seed(self):
        """Test that seed generates reproducible results."""
        seq_length = int(1e5)
        batch_size = int(1e2)
        b = tf.zeros([batch_size, seq_length])

        s1 = Scrambler(seed=1337, binary=True, keep_state=False)
        res_s1_1 = s1(b)
        res_s1_2 = s1(b)
        # new realization per call
        self.assertFalse(np.array_equal(res_s1_1.numpy(), res_s1_2.numpy()))

        # if keep_state=True, the same seed should lead to the same sequence
        s2 = Scrambler(seed=1337, binary=True, keep_state=True)
        res_s2_1 = s2(b)
        s3 = Scrambler(seed=1337)
        res_s3_1 = s3(b)
        # same seed lead to same sequence
        self.assertTrue(np.array_equal(res_s2_1.numpy(), res_s3_1.numpy()))

         # but with random seed it gives a new sequence for each init
        s4 = Scrambler(seed=None, binary=True, keep_state=True)
        res_s4_1 = s2(b)
        s5 = Scrambler(seed=None)
        res_s5_1 = s5(b)
        # same seed lead to same sequence
        self.assertFalse(np.array_equal(res_s4_1.numpy(), res_s5_1.numpy()))

        # for keep_State=False, even the same seed leads to new results
        s6 = Scrambler(seed=1337, binary=True, keep_state=False)
        res_s6_1 = s6(b)
        # different seed generates new sequence
        self.assertFalse(np.array_equal(res_s6_1.numpy(), res_s2_1.numpy()))

        # init with same seed as previous random seed
        s7 = Scrambler(seed=None, binary=True, keep_state=True)
        res_s7_1 = s7(b)
        s8 = Scrambler(seed=s7.seed, binary=True, keep_state=True)
        res_s8_1 = s8(b)
        # same seed lead to same sequence
        self.assertTrue(np.array_equal(res_s7_1.numpy(), res_s8_1.numpy()))

        # test that seed can be also provided to call
        seed = 987654
        s9 = Scrambler(seed=45234, keep_state=False)
        s10 = Scrambler(seed=76543, keep_state=True)
        x1 = s9([b, seed]).numpy()
        x2 = s9([b, seed+1]).numpy()
        x3 = s9([b, seed]).numpy()
        x4 = s10([b, seed]).numpy()
        self.assertFalse(np.array_equal(x1, x2)) # different seed
        self.assertTrue(np.array_equal(x1, x3)) # same seed
        self.assertTrue(np.array_equal(x1, x4)) # same seed (keep_state=f)

        # test that random seed allows inverse
        x5 = s9([b, seed])
        x6 = s9([b, seed]).numpy()
        # same seed
        self.assertTrue(np.array_equal(x5, x6)) # identity
        # different seed
        x7 = s9([b, seed+1])
        self.assertFalse(np.array_equal(x5, x7)) # identity
        # same seed again
        x8 = s9([b, seed+1])
        self.assertTrue(np.array_equal(x7, x8)) # identity

    def test_dtype(self):
        """Test that variable dtypes are supported."""
        seq_length = int(1e1)
        batch_size = int(1e2)

        dt_supported = [tf.float16, tf.float32, tf.float64]
        for dt in dt_supported:
            for dt_in in dt_supported:
                for dt_out in dt_supported:
                    b = tf.zeros([batch_size, seq_length], dtype=dt_in)
                    s1 = Scrambler(dtype=dt)
                    s2 = Descrambler(s1, dtype=dt_out)
                    x = s1(b)
                    y = s2(x)
                    assert (x.dtype==dt)
                    assert (y.dtype==dt_out)

    def test_descrambler(self):
        """"Test that descrambler works as expected."""
        seq_length = int(1e2)
        batch_size = int(1e1)

        b = tf.zeros([batch_size, seq_length])
        s1 = Scrambler()
        s2 = Descrambler(s1)
        x = s1(b)
        y = s2(x)
        assert (np.array_equal(b.numpy(), y.numpy()))

        x = s1([b, 1234])
        y = s2(x)
        assert (not np.array_equal(b.numpy(), y.numpy()))

        # check if seed is correctly retrieved from scrambler
        s3 = Scrambler(seed=12345)
        s4 = Descrambler(s3)
        x = s3(b)
        y = s4(x)
        assert (np.array_equal(b.numpy(), y.numpy()))

    def test_descrambler_nonbin(self):
        """"Test that descrambler works with non-binary."""
        seq_length = int(1e2)
        batch_size = int(1e1)

        b = tf.zeros([batch_size, seq_length])

        # scrambler binary, but descrambler non-binary
        scrambler = Scrambler(seed=1235456, binary=True)
        descrambler = Descrambler(scrambler, binary=False)
        # with explicit seed
        s = 8764
        y = scrambler([b, s])
        z = descrambler([2*y-1, s]) # bspk
        z = 1 + z # remove bpsk
        assert (np.array_equal(b.numpy(), z.numpy()))
        #without explicit seed
        y = scrambler(b)
        z = descrambler(2*y-1) # bspk
        z = 1 + z # remove bpsk
        assert (np.array_equal(b.numpy(), z.numpy()))

        # scrambler non-binary, but descrambler
        scrambler = Scrambler(seed=1235456, binary=False)
        descrambler = Descrambler(scrambler, binary=True)
        s = 546342
        y = scrambler([2*b-1, s]) # bspk
        y = 0.5*(1 + y) # remove bpsk
        z = descrambler([y, s])
        assert (np.array_equal(b.numpy(), z.numpy()))
        #without explicit seed
        y = scrambler(2*b-1) # bspk
        y = 0.5*(1 + y) # remove bpsk
        z = descrambler(y)
        y = 1 + y # remove bpsk
        assert (np.array_equal(b.numpy(), z.numpy()))

    def test_scrambler_binary(self):
        """test that binary flag can be used as input"""
        seq_length = int(1e2)
        batch_size = int(1e1)

        b = tf.ones([batch_size, seq_length])

        # scrambler binary, but descrambler non-binary

        scrambler = Scrambler(seed=1245, binary=True)

        s = 1234
        x1 = scrambler(b) # binary scrambling
        x2 = scrambler([b, s]) # binary scrambling different seed
        x3 = scrambler([b, s, True]) # binary scrambling different seed
        x4 = scrambler([b, s, False]) # non-binary scrambling different seed

        assert (not np.array_equal(x1.numpy(), x2.numpy())) # different seed
        assert (np.array_equal(x2.numpy(), x3.numpy())) # same seed
        # same but "bpsk modulated"
        assert (not np.array_equal(x1.numpy(), 0.5*(1+x4.numpy())))

# # SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Number of the GPU to be used Test against correct dimensions of the sequence # keep_State=True # init new scrambler for new sequence size; # only different batch_sizes are allowed in this mode # build scrambler # keep_State=False # build scrambler Test that scrambling sequence has no offset, i.e., equal likely 0s         and 1s # test some initial seeds # build scrambler # generate a random sequence Test that scrambling sequence is random per batch sample iff         keep_batch_dims=True. # generate a random sequence # each batch sample must be different # test that the pattern is the same of option keep_batch_constant==True # generate a random sequence # each batch sample is the same Test that scrambling sequences are random for each new realization. # generate a random sequence Test that scrambling can be inverted/removed.         2x scrambling must result in the original sequence (for binary and          LLRs). #check binary scrambling # only works if keep_state=True #check soft-value scrambling (flip sign) Test that scrambling works for soft-values (sign flip). Test that keep_state works as expected.         Iff keep_state==True, the scrambled sequences must be constant. # also check that the sequence is unique with keep_state=False Test that Keras model can be compiled (supports dynamic shapes). # test that output batch dim is none # test that model can be called # call twice to see that bs can change Test that graph mode and XLA works as expected # again with different batch_size # and different sequence length Test that seed generates reproducible results. # new realization per call # if keep_state=True, the same seed should lead to the same sequence # same seed lead to same sequence # but with random seed it gives a new sequence for each init # same seed lead to same sequence # for keep_State=False, even the same seed leads to new results # different seed generates new sequence # init with same seed as previous random seed # same seed lead to same sequence # test that seed can be also provided to call # different seed # same seed # same seed (keep_state=f) # test that random seed allows inverse # same seed # identity # different seed # identity # same seed again # identity Test that variable dtypes are supported. "Test that descrambler works as expected. # check if seed is correctly retrieved from scrambler "Test that descrambler works with non-binary. # scrambler binary, but descrambler non-binary # with explicit seed # bspk # remove bpsk #without explicit seed # bspk # remove bpsk # scrambler non-binary, but descrambler # bspk # remove bpsk #without explicit seed # bspk # remove bpsk # remove bpsk test that binary flag can be used as input # scrambler binary, but descrambler non-binary # binary scrambling # binary scrambling different seed # binary scrambling different seed # non-binary scrambling different seed # different seed # same seed # same but "bpsk modulated"