stream_management.py (9637B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 "Classeds and functions related to stream management in MIMO systems" 6 7 import numpy as np 8 9 10 class StreamManagement(): 11 """Class for management of streams in multi-cell MIMO networks. 12 13 Parameters 14 ---------- 15 rx_tx_association : [num_rx, num_tx], np.int 16 A binary NumPy array where ``rx_tx_association[i,j]=1`` means 17 that receiver `i` gets one or multiple streams from 18 transmitter `j`. 19 20 num_streams_per_tx : int 21 Indicates the number of streams that are transmitted by each 22 transmitter. 23 24 Note 25 ---- 26 Several symmetry constraints on ``rx_tx_association`` are imposed 27 to ensure efficient processing. All row sums and all column sums 28 must be equal, i.e., all receivers have the same number of associated 29 transmitters and all transmitters have the same number of associated 30 receivers. It is also assumed that all transmitters send the same 31 number of streams ``num_streams_per_tx``. 32 """ 33 def __init__(self, 34 rx_tx_association, 35 num_streams_per_tx): 36 37 super().__init__() 38 self._num_streams_per_tx = int(num_streams_per_tx) 39 self.rx_tx_association = rx_tx_association 40 41 @property 42 def rx_tx_association(self): 43 """Association between receivers and transmitters. 44 45 A binary NumPy array of shape `[num_rx, num_tx]`, 46 where ``rx_tx_association[i,j]=1`` means that receiver `i` 47 gets one ore multiple streams from transmitter `j`. 48 """ 49 return self._rx_tx_association 50 51 @property 52 def num_rx(self): 53 "Number of receivers." 54 return self._num_rx 55 56 @property 57 def num_tx(self): 58 "Number of transmitters." 59 return self._num_tx 60 61 @property 62 def num_streams_per_tx(self): 63 "Number of streams per transmitter." 64 return self._num_streams_per_tx 65 66 @property 67 def num_streams_per_rx(self): 68 "Number of streams transmitted to each receiver." 69 return int(self.num_tx*self.num_streams_per_tx/self.num_rx) 70 71 @property 72 def num_interfering_streams_per_rx(self): 73 "Number of interfering streams received at each eceiver." 74 return int(self.num_tx*self.num_streams_per_tx 75 - self.num_streams_per_rx) 76 77 @property 78 def num_tx_per_rx(self): 79 "Number of transmitters communicating with a receiver." 80 return self._num_tx_per_rx 81 82 @property 83 def num_rx_per_tx(self): 84 "Number of receivers communicating with a transmitter." 85 return self._num_rx_per_tx 86 87 @property 88 def precoding_ind(self): 89 """Indices needed to gather channels for precoding. 90 91 A NumPy array of shape `[num_tx, num_rx_per_tx]`, 92 where ``precoding_ind[i,:]`` contains the indices of the 93 receivers to which transmitter `i` is sending streams. 94 """ 95 return self._precoding_ind 96 97 @property 98 def stream_association(self): 99 """Association between receivers, transmitters, and streams. 100 101 A binary NumPy array of shape 102 `[num_rx, num_tx, num_streams_per_tx]`, where 103 ``stream_association[i,j,k]=1`` means that receiver `i` gets 104 the `k` th stream from transmitter `j`. 105 """ 106 return self._stream_association 107 108 @property 109 def detection_desired_ind(self): 110 """Indices needed to gather desired channels for receive processing. 111 112 A NumPy array of shape `[num_rx*num_streams_per_rx]` that 113 can be used to gather desired channels from the flattened 114 channel tensor of shape 115 `[...,num_rx, num_tx, num_streams_per_tx,...]`. 116 The result of the gather operation can be reshaped to 117 `[...,num_rx, num_streams_per_rx,...]`. 118 """ 119 return self._detection_desired_ind 120 121 @property 122 def detection_undesired_ind(self): 123 """Indices needed to gather undesired channels for receive processing. 124 125 A NumPy array of shape `[num_rx*num_streams_per_rx]` that 126 can be used to gather undesired channels from the flattened 127 channel tensor of shape `[...,num_rx, num_tx, num_streams_per_tx,...]`. 128 The result of the gather operation can be reshaped to 129 `[...,num_rx, num_interfering_streams_per_rx,...]`. 130 """ 131 return self._detection_undesired_ind 132 133 @property 134 def tx_stream_ids(self): 135 """Mapping of streams to transmitters. 136 137 A NumPy array of shape `[num_tx, num_streams_per_tx]`. 138 Streams are numbered from 0,1,... and assiged to transmitters in 139 increasing order, i.e., transmitter 0 gets the first 140 `num_streams_per_tx` and so on. 141 """ 142 return self._tx_stream_ids 143 144 @property 145 def rx_stream_ids(self): 146 """Mapping of streams to receivers. 147 148 A Numpy array of shape `[num_rx, num_streams_per_rx]`. 149 This array is obtained from ``tx_stream_ids`` together with 150 the ``rx_tx_association``. ``rx_stream_ids[i,:]`` contains 151 the indices of streams that are supposed to be decoded by receiver `i`. 152 """ 153 return self._rx_stream_ids 154 155 @property 156 def stream_ind(self): 157 """Indices needed to gather received streams in the correct order. 158 159 A NumPy array of shape `[num_rx*num_streams_per_rx]` that can be 160 used to gather streams from the flattened tensor of received streams 161 of shape `[...,num_rx, num_streams_per_rx,...]`. The result of the 162 gather operation is then reshaped to 163 `[...,num_tx, num_streams_per_tx,...]`. 164 """ 165 return self._stream_ind 166 167 @rx_tx_association.setter 168 def rx_tx_association(self, rx_tx_association): 169 """Sets the rx_tx_association and derives related properties. """ 170 171 # Make sure that rx_tx_association is a binary NumPy array 172 rx_tx_association = np.array(rx_tx_association, np.int32) 173 assert all(x in [0,1] for x in np.nditer(rx_tx_association)), \ 174 "All elements of `stream_association` must be 0 or 1." 175 176 # Obtain num_rx, num_tx from stream_association shape 177 self._num_rx, self._num_tx = np.shape(rx_tx_association) 178 179 # Each receiver must be associated with the same number of transmitters 180 num_tx_per_rx = np.sum(rx_tx_association, 1) 181 assert np.min(num_tx_per_rx) == np.max(num_tx_per_rx), \ 182 """Each receiver needs to be associated with the same number 183 of transmitters.""" 184 self._num_tx_per_rx = num_tx_per_rx[0] 185 186 # Each transmitter must be associated with the same number of receivers 187 num_rx_per_tx = np.sum(rx_tx_association, 0) 188 assert np.min(num_rx_per_tx) == np.max(num_rx_per_tx), \ 189 """Each transmitter needs to be associated with the same number 190 of receivers.""" 191 self._num_rx_per_tx = num_rx_per_tx[0] 192 193 self._rx_tx_association = rx_tx_association 194 195 # Compute indices for precoding 196 self._precoding_ind = np.zeros([self.num_tx, self.num_rx_per_tx], 197 np.int32) 198 for i in range(self.num_tx): 199 self._precoding_ind[i,:] = np.where(self.rx_tx_association[:,i])[0] 200 201 # Construct the stream association matrix 202 # The element [i,j,k]=1 indicates that receiver i, get the kth stream 203 # from transmitter j. 204 stream_association = np.zeros( 205 [self.num_rx, self.num_tx, self.num_streams_per_tx], np.int32) 206 n_streams = np.min([self.num_streams_per_rx, self.num_streams_per_tx]) 207 tmp = np.ones([n_streams]) 208 for j in range(self.num_tx): 209 c = 0 210 for i in range(self.num_rx): 211 # If receiver i gets anything from transmitter j 212 if rx_tx_association[i,j]: 213 stream_association[i,j,c:c+self.num_streams_per_rx] = tmp 214 c += self.num_streams_per_rx 215 self._stream_association = stream_association 216 217 # Get indices of desired and undesired channel coefficients from 218 # the flattened stream_association. These indices can be used by 219 # a receiver to gather channels of desired and undesired streams. 220 self._detection_desired_ind = \ 221 np.where(np.reshape(stream_association, [-1])==1)[0] 222 223 self._detection_undesired_ind = \ 224 np.where(np.reshape(stream_association, [-1])==0)[0] 225 226 # We number streams from 0,1,... and assign them to the TX 227 # TX 0 gets the first num_streams_per_tx and so on: 228 self._tx_stream_ids = np.reshape( 229 np.arange(0, self.num_tx*self.num_streams_per_tx), 230 [self.num_tx, self.num_streams_per_tx]) 231 232 # We now compute the stream_ids for each receiver 233 self._rx_stream_ids = np.zeros([self.num_rx, self.num_streams_per_rx], 234 np.int32) 235 for i in range(self.num_rx): 236 c = [] 237 for j in range(self.num_tx): 238 # If receiver i gets anything from transmitter j 239 if rx_tx_association[i,j]: 240 tmp = np.where(stream_association[i,j])[0] 241 tmp += j*self.num_streams_per_tx 242 c += list(tmp) 243 self._rx_stream_ids[i,:] = c 244 245 # Get indices to bring received streams back to the right order in 246 # which they were transmitted. 247 self._stream_ind = np.argsort(np.reshape(self._rx_stream_ids, [-1]))