gen_data.py (1761B)
1 #!/usr/bin/python 2 3 import numpy as np 4 from numpy.random import random, normal 5 from sys import argv 6 from os import path, makedirs 7 8 if len(argv) != 4: 9 print("Usage: python gen_data.py <NUM_RX> <NUM_TX> <NUM_SC>") 10 exit(1) 11 NUM_RX = int(argv[1]) # number of receive antennas 12 NUM_TX = int(argv[2]) # number of transmit antennas 13 NUM_SC = int(argv[3]) # number of subcarriers 14 15 NOISE_STD_DEVIATION = np.sqrt(.5) / 100 # noise standard deviation 16 17 # Transmitter signal 18 x = random((NUM_TX, NUM_SC)) \ 19 + 1.j * random((NUM_TX, NUM_SC)) 20 x = (x - 0.5 - 0.5j) * 2 # scale it from [0, 1] to [-1, 1] 21 # Channel 22 H = random((NUM_RX, NUM_TX, NUM_SC)) \ 23 + 1.j * random((NUM_RX, NUM_TX, NUM_SC)) 24 H = (H - 0.5 - 0.5j) * 2 # scale it from [0, 1] to [-1, 1] 25 # Noise 26 n = normal(0, NOISE_STD_DEVIATION, (NUM_RX, NUM_SC)) \ 27 + 1.j * normal(0, NOISE_STD_DEVIATION, (NUM_RX, NUM_SC)) 28 # Received signal 29 y = np.einsum("ijk,jk->ik", H, x) + n 30 # Noise covariance matrix 31 R = np.eye(NUM_TX, NUM_TX, dtype=np.complex64) * NOISE_STD_DEVIATION**2 32 33 data_tuple = ( 34 x.real.astype(np.float32), x.imag.astype(np.float32), 35 H.real.astype(np.float32), H.imag.astype(np.float32), 36 R.real.astype(np.float32), R.imag.astype(np.float32), 37 y.real.astype(np.float32), y.imag.astype(np.float32) 38 ) 39 data_filenames = ( 40 "data/x_re.bin", "data/x_im.bin", 41 "data/H_re.bin", "data/H_im.bin", 42 "data/R_re.bin", "data/R_im.bin", 43 "data/y_re.bin", "data/y_im.bin" 44 ) 45 46 # Create "data" directory if it does not exist 47 data_dir = path.join(path.dirname(__file__), "..", "data") 48 if not path.exists(data_dir): 49 makedirs(data_dir) 50 51 # Write data to bin files 52 for data, filename in zip(data_tuple, data_filenames): 53 with open(filename, "wb") as f: 54 f.write(data.tobytes())