anomaly-detection-material-parameters-calibration

Sionna param calibration (research proj)
git clone https://git.ea.contact/anomaly-detection-material-parameters-calibration
Log | Files | Refs | README

channel_coefficients.py (38496B)


      1 #
      2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
      3 # SPDX-License-Identifier: Apache-2.0
      4 #
      5 """
      6 Class for sampling channel impulse responses following 3GPP TR38.901
      7 specifications and giving LSPs and rays.
      8 """
      9 
     10 import tensorflow as tf
     11 from tensorflow import sin, cos, acos
     12 
     13 from sionna import PI, SPEED_OF_LIGHT
     14 import sionna
     15 
     16 class Topology:
     17     # pylint: disable=line-too-long
     18     r"""
     19     Class for conveniently storing the network topology information required
     20     for sampling channel impulse responses
     21 
     22     Parameters
     23     -----------
     24 
     25     velocities : [batch size, number of UTs], tf.float
     26         UT velocities
     27 
     28     moving_end : str
     29         Indicated which end of the channel (TX or RX) is moving. Either "tx" or
     30         "rx".
     31 
     32     los_aoa : [batch size, number of BSs, number of UTs], tf.float
     33         Azimuth angle of arrival of LoS path [radian]
     34 
     35     los_aod : [batch size, number of BSs, number of UTs], tf.float
     36         Azimuth angle of departure of LoS path [radian]
     37 
     38     los_zoa : [batch size, number of BSs, number of UTs], tf.float
     39         Zenith angle of arrival for of path [radian]
     40 
     41     los_zod : [batch size, number of BSs, number of UTs], tf.float
     42         Zenith angle of departure for of path [radian]
     43 
     44     los : [batch size, number of BSs, number of UTs], tf.bool
     45         Indicate for each BS-UT link if it is in LoS
     46 
     47     distance_3d : [batch size, number of UTs, number of UTs], tf.float
     48         Distance between the UTs in X-Y-Z space (not only X-Y plan).
     49 
     50     tx_orientations : [batch size, number of TXs, 3], tf.float
     51         Orientations of the transmitters, which are either BSs or UTs depending
     52         on the link direction [radian].
     53 
     54     rx_orientations : [batch size, number of RXs, 3], tf.float
     55         Orientations of the receivers, which are either BSs or UTs depending on
     56         the link direction [radian].
     57     """
     58 
     59     def __init__(self,  velocities,
     60                         moving_end,
     61                         los_aoa,
     62                         los_aod,
     63                         los_zoa,
     64                         los_zod,
     65                         los,
     66                         distance_3d,
     67                         tx_orientations,
     68                         rx_orientations):
     69         self.velocities = velocities
     70         self.moving_end = moving_end
     71         self.los_aoa = los_aoa
     72         self.los_aod = los_aod
     73         self.los_zoa = los_zoa
     74         self.los_zod = los_zod
     75         self.los = los
     76         self.tx_orientations = tx_orientations
     77         self.rx_orientations = rx_orientations
     78         self.distance_3d = distance_3d
     79 
     80 class ChannelCoefficientsGenerator:
     81     # pylint: disable=line-too-long
     82     r"""
     83     Sample channel impulse responses according to LSPs rays.
     84 
     85     This class implements steps 10 and 11 from the TR 38.901 specifications,
     86     (section 7.5).
     87 
     88     Parameters
     89     ----------
     90     carrier_frequency : float
     91         Carrier frequency [Hz]
     92 
     93     tx_array : PanelArray
     94         Panel array used by the transmitters.
     95         All transmitters share the same antenna array configuration.
     96 
     97     rx_array : PanalArray
     98         Panel array used by the receivers.
     99         All transmitters share the same antenna array configuration.
    100 
    101     subclustering : bool
    102         Use subclustering if set to `True` (see step 11 for section 7.5 in
    103         TR 38.901). CDL does not use subclustering. System level models (UMa,
    104         UMi, RMa) do.
    105 
    106     dtype : Complex tf.DType
    107         Defines the datatype for internal calculations and the output
    108         dtype. Defaults to `tf.complex64`.
    109 
    110     Input
    111     -----
    112     num_time_samples : int
    113         Number of samples
    114 
    115     sampling_frequency : float
    116         Sampling frequency [Hz]
    117 
    118     k_factor : [batch_size, number of TX, number of RX]
    119         K-factor
    120 
    121     rays : Rays
    122         Rays from which to compute thr CIR
    123 
    124     topology : Topology
    125         Topology of the network
    126 
    127     c_ds : [batch size, number of TX, number of RX]
    128         Cluster DS [ns]. Only needed when subclustering is used
    129         (``subclustering`` set to `True`), i.e., with system level models.
    130         Otherwise can be set to None.
    131         Defaults to None.
    132 
    133     debug : bool
    134         If set to `True`, additional information is returned in addition to
    135         paths coefficients and delays: The random phase shifts (see step 10 of
    136         section 7.5 in TR38.901 specification), and the time steps at which the
    137         channel is sampled.
    138 
    139     Output
    140     ------
    141     h : [batch size, num TX, num RX, num paths, num RX antenna, num TX antenna, num samples], tf.complex
    142         Paths coefficients
    143 
    144     delays : [batch size, num TX, num RX, num paths], tf.real
    145         Paths delays [s]
    146 
    147     phi : [batch size, number of BSs, number of UTs, 4], tf.real
    148         Initial phases (see step 10 of section 7.5 in TR 38.901 specification).
    149         Last dimension corresponds to the four polarization combinations.
    150 
    151     sample_times : [number of time steps], tf.float
    152         Sampling time steps
    153     """
    154 
    155     def __init__(self,  carrier_frequency,
    156                         tx_array, rx_array,
    157                         subclustering,
    158                         dtype=tf.complex64):
    159         assert dtype.is_complex, "dtype must be a complex datatype"
    160         self._dtype = dtype
    161 
    162         # Wavelength (m)
    163         self._lambda_0 = tf.constant(SPEED_OF_LIGHT/carrier_frequency,
    164             dtype.real_dtype)
    165         self._tx_array = tx_array
    166         self._rx_array = rx_array
    167         self._subclustering = subclustering
    168 
    169         # Sub-cluster information for intra cluster delay spread clusters
    170         # This is hardcoded from Table 7.5-5
    171         self._sub_cl_1_ind = tf.constant([0,1,2,3,4,5,6,7,18,19], tf.int32)
    172         self._sub_cl_2_ind = tf.constant([8,9,10,11,16,17], tf.int32)
    173         self._sub_cl_3_ind = tf.constant([12,13,14,15], tf.int32)
    174         self._sub_cl_delay_offsets = tf.constant([0, 1.28, 2.56],
    175                                                     dtype.real_dtype)
    176 
    177     def __call__(self, num_time_samples, sampling_frequency, k_factor, rays,
    178                  topology, c_ds=None, debug=False):
    179         # Sample times
    180         sample_times = (tf.range(num_time_samples,
    181                 dtype=self._dtype.real_dtype)/sampling_frequency)
    182 
    183         # Step 10
    184         phi = self._step_10(tf.shape(rays.aoa))
    185 
    186         # Step 11
    187         h, delays = self._step_11(phi, topology, k_factor, rays, sample_times,
    188                                                                         c_ds)
    189 
    190         # Return additional information if requested
    191         if debug:
    192             return h, delays, phi, sample_times
    193 
    194         return h, delays
    195 
    196     ###########################################
    197     # Utility functions
    198     ###########################################
    199 
    200     def _unit_sphere_vector(self, theta, phi):
    201         r"""
    202         Generate vector on unit sphere (7.1-6)
    203 
    204         Input
    205         -------
    206         theta : Arbitrary shape, tf.float
    207             Zenith [radian]
    208 
    209         phi : Same shape as ``theta``, tf.float
    210             Azimuth [radian]
    211 
    212         Output
    213         --------
    214         rho_hat : ``phi.shape`` + [3, 1]
    215             Vector on unit sphere
    216 
    217         """
    218         rho_hat = tf.stack([sin(theta)*cos(phi),
    219                             sin(theta)*sin(phi),
    220                             cos(theta)], axis=-1)
    221         return tf.expand_dims(rho_hat, axis=-1)
    222 
    223     def _forward_rotation_matrix(self, orientations):
    224         r"""
    225         Forward composite rotation matrix (7.1-4)
    226 
    227         Input
    228         ------
    229             orientations : [...,3], tf.float
    230                 Orientation to which to rotate [radian]
    231 
    232         Output
    233         -------
    234         R : [...,3,3], tf.float
    235             Rotation matrix
    236         """
    237         a, b, c = orientations[...,0], orientations[...,1], orientations[...,2]
    238 
    239         row_1 = tf.stack([cos(a)*cos(b),
    240             cos(a)*sin(b)*sin(c)-sin(a)*cos(c),
    241             cos(a)*sin(b)*cos(c)+sin(a)*sin(c)], axis=-1)
    242 
    243         row_2 = tf.stack([sin(a)*cos(b),
    244             sin(a)*sin(b)*sin(c)+cos(a)*cos(c),
    245             sin(a)*sin(b)*cos(c)-cos(a)*sin(c)], axis=-1)
    246 
    247         row_3 = tf.stack([-sin(b),
    248             cos(b)*sin(c),
    249             cos(b)*cos(c)], axis=-1)
    250 
    251         rot_mat = tf.stack([row_1, row_2, row_3], axis=-2)
    252         return rot_mat
    253 
    254     def _rot_pos(self, orientations, positions):
    255         r"""
    256         Rotate the ``positions`` according to the ``orientations``
    257 
    258         Input
    259         ------
    260         orientations : [...,3], tf.float
    261             Orientation to which to rotate [radian]
    262 
    263         positions : [...,3,1], tf.float
    264             Positions to rotate
    265 
    266         Output
    267         -------
    268         : [...,3,1], tf.float
    269             Rotated positions
    270         """
    271         rot_mat = self._forward_rotation_matrix(orientations)
    272         return tf.matmul(rot_mat, positions)
    273 
    274     def _reverse_rotation_matrix(self, orientations):
    275         r"""
    276         Reverse composite rotation matrix (7.1-4)
    277 
    278         Input
    279         ------
    280         orientations : [...,3], tf.float
    281             Orientations to rotate to  [radian]
    282 
    283         Output
    284         -------
    285         R_inv : [...,3,3], tf.float
    286             Inverse of the rotation matrix corresponding to ``orientations``
    287         """
    288         rot_mat = self._forward_rotation_matrix(orientations)
    289         rot_mat_inv = tf.linalg.matrix_transpose(rot_mat)
    290         return rot_mat_inv
    291 
    292     def _gcs_to_lcs(self, orientations, theta, phi):
    293         # pylint: disable=line-too-long
    294         r"""
    295         Compute the angles ``theta``, ``phi`` in LCS rotated according to
    296         ``orientations`` (7.1-7/8)
    297 
    298         Input
    299         ------
    300         orientations : [...,3] of rank K, tf.float
    301             Orientations to which to rotate to [radian]
    302 
    303         theta : Broadcastable to the first K-1 dimensions of ``orientations``, tf.float
    304             Zenith to rotate [radian]
    305 
    306         phi : Same dimension as ``theta``, tf.float
    307             Azimuth to rotate [radian]
    308 
    309         Output
    310         -------
    311         theta_prime : Same dimension as ``theta``, tf.float
    312             Rotated zenith
    313 
    314         phi_prime : Same dimensions as ``theta`` and ``phi``, tf.float
    315             Rotated azimuth
    316         """
    317 
    318         rho_hat = self._unit_sphere_vector(theta, phi)
    319         rot_inv = self._reverse_rotation_matrix(orientations)
    320         rot_rho = tf.matmul(rot_inv, rho_hat)
    321         v1 = tf.constant([0,0,1], self._dtype.real_dtype)
    322         v1 = tf.reshape(v1, [1]*(rot_rho.shape.rank-1)+[3])
    323         v2 = tf.constant([1+0j,1j,0], self._dtype)
    324         v2 = tf.reshape(v2, [1]*(rot_rho.shape.rank-1)+[3])
    325         z = tf.matmul(v1, rot_rho)
    326         z = tf.clip_by_value(z, tf.constant(-1., self._dtype.real_dtype),
    327                              tf.constant(1., self._dtype.real_dtype))
    328         theta_prime = acos(z)
    329         phi_prime = tf.math.angle((tf.matmul(v2, tf.cast(rot_rho,
    330             self._dtype))))
    331         theta_prime = tf.squeeze(theta_prime, axis=[phi.shape.rank,
    332             phi.shape.rank+1])
    333         phi_prime = tf.squeeze(phi_prime, axis=[phi.shape.rank,
    334             phi.shape.rank+1])
    335 
    336         return (theta_prime, phi_prime)
    337 
    338     def _compute_psi(self, orientations, theta, phi):
    339         # pylint: disable=line-too-long
    340         r"""
    341         Compute displacement angle :math:`Psi` for the transformation of LCS-GCS
    342         field components in (7.1-15) of TR38.901 specification
    343 
    344         Input
    345         ------
    346         orientations : [...,3], tf.float
    347             Orientations to which to rotate to [radian]
    348 
    349         theta :  Broadcastable to the first K-1 dimensions of ``orientations``, tf.float
    350             Spherical position zenith [radian]
    351 
    352         phi : Same dimensions as ``theta``, tf.float
    353             Spherical position azimuth [radian]
    354 
    355         Output
    356         -------
    357             Psi : Same shape as ``theta`` and ``phi``, tf.float
    358                 Displacement angle :math:`Psi`
    359         """
    360         a = orientations[...,0]
    361         b = orientations[...,1]
    362         c = orientations[...,2]
    363         real = sin(c)*cos(theta)*sin(phi-a)
    364         real += cos(c)*(cos(b)*sin(theta)-sin(b)*cos(theta)*cos(phi-a))
    365         imag = sin(c)*cos(phi-a) + sin(b)*cos(c)*sin(phi-a)
    366         psi = tf.math.angle(tf.complex(real, imag))
    367         return psi
    368 
    369     def _l2g_response(self, f_prime, orientations, theta, phi):
    370         # pylint: disable=line-too-long
    371         r"""
    372         Transform field components from LCS to GCS (7.1-11)
    373 
    374         Input
    375         ------
    376         f_prime : K-Dim Tensor of shape [...,2], tf.float
    377             Field components
    378 
    379         orientations : K-Dim Tensor of shape [...,3], tf.float
    380             Orientations of LCS-GCS [radian]
    381 
    382         theta : K-1-Dim Tensor with matching dimensions to ``f_prime`` and ``phi``, tf.float
    383             Spherical position zenith [radian]
    384 
    385         phi : Same dimensions as ``theta``, tf.float
    386             Spherical position azimuth [radian]
    387 
    388         Output
    389         ------
    390             F : K+1-Dim Tensor with shape [...,2,1], tf.float
    391                 The first K dimensions are identical to those of ``f_prime``
    392         """
    393         psi = self._compute_psi(orientations, theta, phi)
    394         row1 = tf.stack([cos(psi), -sin(psi)], axis=-1)
    395         row2 = tf.stack([sin(psi), cos(psi)], axis=-1)
    396         mat = tf.stack([row1, row2], axis=-2)
    397         f = tf.matmul(mat, tf.expand_dims(f_prime, -1))
    398         return f
    399 
    400     def _step_11_get_tx_antenna_positions(self, topology):
    401         r"""Compute d_bar_tx in (7.5-22), i.e., the positions in GCS of elements
    402         forming the transmit panel
    403 
    404         Input
    405         -----
    406         topology : Topology
    407             Topology of the network
    408 
    409         Output
    410         -------
    411         d_bar_tx : [batch_size, num TXs, num TX antenna, 3]
    412             Positions of the antenna elements in the GCS
    413         """
    414         # Get BS orientations got broadcasting
    415         tx_orientations = topology.tx_orientations
    416         tx_orientations = tf.expand_dims(tx_orientations, 2)
    417 
    418         # Get antenna element positions in LCS and reshape for broadcasting
    419         tx_ant_pos_lcs = self._tx_array.ant_pos
    420         tx_ant_pos_lcs = tf.reshape(tx_ant_pos_lcs,
    421             [1,1]+tx_ant_pos_lcs.shape+[1])
    422 
    423         # Compute antenna element positions in GCS
    424         tx_ant_pos_gcs = self._rot_pos(tx_orientations, tx_ant_pos_lcs)
    425         tx_ant_pos_gcs = tf.reshape(tx_ant_pos_gcs,
    426             tf.shape(tx_ant_pos_gcs)[:-1])
    427 
    428         d_bar_tx = tx_ant_pos_gcs
    429 
    430         return d_bar_tx
    431 
    432     def _step_11_get_rx_antenna_positions(self, topology):
    433         r"""Compute d_bar_rx in (7.5-22), i.e., the positions in GCS of elements
    434         forming the receive antenna panel
    435 
    436         Input
    437         -----
    438         topology : Topology
    439             Topology of the network
    440 
    441         Output
    442         -------
    443         d_bar_rx : [batch_size, num RXs, num RX antenna, 3]
    444             Positions of the antenna elements in the GCS
    445         """
    446         # Get UT orientations got broadcasting
    447         rx_orientations = topology.rx_orientations
    448         rx_orientations = tf.expand_dims(rx_orientations, 2)
    449 
    450         # Get antenna element positions in LCS and reshape for broadcasting
    451         rx_ant_pos_lcs = self._rx_array.ant_pos
    452         rx_ant_pos_lcs = tf.reshape(rx_ant_pos_lcs,
    453             [1,1]+rx_ant_pos_lcs.shape+[1])
    454 
    455         # Compute antenna element positions in GCS
    456         rx_ant_pos_gcs = self._rot_pos(rx_orientations, rx_ant_pos_lcs)
    457         rx_ant_pos_gcs = tf.reshape(rx_ant_pos_gcs,
    458             tf.shape(rx_ant_pos_gcs)[:-1])
    459 
    460         d_bar_rx = rx_ant_pos_gcs
    461 
    462         return d_bar_rx
    463 
    464     def _step_10(self, shape):
    465         r"""
    466         Generate random and uniformly distributed phases for all rays and
    467         polarization combinations
    468 
    469         Input
    470         -----
    471         shape : Shape tensor
    472             Shape of the leading dimensions for the tensor of phases to generate
    473 
    474         Output
    475         ------
    476         phi : [shape] + [4], tf.float
    477             Phases for all polarization combinations
    478         """
    479         phi = sionna.config.tf_rng.uniform(
    480                              tf.concat([shape, [4]], axis=0),
    481                              minval=-PI,
    482                              maxval=PI,
    483                              dtype=self._dtype.real_dtype)
    484 
    485         return phi
    486 
    487     def _step_11_phase_matrix(self, phi, rays):
    488         # pylint: disable=line-too-long
    489         r"""
    490         Compute matrix with random phases in (7.5-22)
    491 
    492         Input
    493         -----
    494         phi : [batch size, num TXs, num RXs, num clusters, num rays, 4], tf.float
    495             Initial phases for all combinations of polarization
    496 
    497         rays : Rays
    498             Rays
    499 
    500         Output
    501         ------
    502         h_phase : [batch size, num TXs, num RXs, num clusters, num rays, 2, 2], tf.complex
    503             Matrix with random phases in (7.5-22)
    504         """
    505         xpr = rays.xpr
    506 
    507         xpr_scaling = tf.complex(tf.sqrt(1/xpr),
    508             tf.constant(0., self._dtype.real_dtype))
    509         e0 = tf.exp(tf.complex(tf.constant(0., self._dtype.real_dtype),
    510             phi[...,0]))
    511         e3 = tf.exp(tf.complex(tf.constant(0., self._dtype.real_dtype),
    512             phi[...,3]))
    513         e1 = xpr_scaling*tf.exp(tf.complex(tf.constant(0.,
    514                                 self._dtype.real_dtype), phi[...,1]))
    515         e2 = xpr_scaling*tf.exp(tf.complex(tf.constant(0.,
    516                                 self._dtype.real_dtype), phi[...,2]))
    517         shape = tf.concat([tf.shape(e0), [2,2]], axis=-1)
    518         h_phase = tf.reshape(tf.stack([e0, e1, e2, e3], axis=-1), shape)
    519 
    520         return h_phase
    521 
    522     def _step_11_doppler_matrix(self, topology, aoa, zoa, t):
    523         # pylint: disable=line-too-long
    524         r"""
    525         Compute matrix with phase shifts due to mobility in (7.5-22)
    526 
    527         Input
    528         -----
    529         topology : Topology
    530             Topology of the network
    531 
    532         aoa : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    533             Azimuth angles of arrivals [radian]
    534 
    535         zoa : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    536             Zenith angles of arrivals [radian]
    537 
    538         t : [number of time steps]
    539             Time steps at which the channel is sampled
    540 
    541         Output
    542         ------
    543         h_doppler : [batch size, num_tx, num rx, num clusters, num rays, num time steps], tf.complex
    544             Matrix with phase shifts due to mobility in (7.5-22)
    545         """
    546         lambda_0 = self._lambda_0
    547         velocities = topology.velocities
    548 
    549         # Add an extra dimension to make v_bar broadcastable with the time
    550         # dimension
    551         # v_bar [batch size, num tx or num rx, 3, 1]
    552         v_bar = velocities
    553         v_bar = tf.expand_dims(v_bar, axis=-1)
    554 
    555         # Depending on which end of the channel is moving, tx or rx, we add an
    556         # extra dimension to make this tensor broadcastable with the other end
    557         if topology.moving_end == 'rx':
    558             # v_bar [batch size, 1, num rx, num tx, 1]
    559             v_bar = tf.expand_dims(v_bar, 1)
    560         elif topology.moving_end == 'tx':
    561             # v_bar [batch size, num tx, 1, num tx, 1]
    562             v_bar = tf.expand_dims(v_bar, 2)
    563 
    564         # v_bar [batch size, 1, num rx, 1, 1, 3, 1]
    565         # or    [batch size, num tx, 1, 1, 1, 3, 1]
    566         v_bar = tf.expand_dims(tf.expand_dims(v_bar, -3), -3)
    567 
    568         # v_bar [batch size, num_tx, num rx, num clusters, num rays, 3, 1]
    569         r_hat_rx = self._unit_sphere_vector(zoa, aoa)
    570 
    571         # Compute phase shift due to doppler
    572         # [batch size, num_tx, num rx, num clusters, num rays, num time steps]
    573         exponent = 2*PI/lambda_0*tf.reduce_sum(r_hat_rx*v_bar, -2)*t
    574         h_doppler = tf.exp(tf.complex(tf.constant(0.,
    575                                     self._dtype.real_dtype), exponent))
    576 
    577         # [batch size, num_tx, num rx, num clusters, num rays, num time steps]
    578         return h_doppler
    579 
    580     def _step_11_array_offsets(self, topology, aoa, aod, zoa, zod):
    581         # pylint: disable=line-too-long
    582         r"""
    583         Compute matrix accounting for phases offsets between antenna elements
    584 
    585         Input
    586         -----
    587         topology : Topology
    588             Topology of the network
    589 
    590         aoa : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    591             Azimuth angles of arrivals [radian]
    592 
    593         aod : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    594             Azimuth angles of departure [radian]
    595 
    596         zoa : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    597             Zenith angles of arrivals [radian]
    598 
    599         zod : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    600             Zenith angles of departure [radian]
    601         Output
    602         ------
    603         h_array : [batch size, num_tx, num rx, num clusters, num rays, num rx antennas, num tx antennas], tf.complex
    604             Matrix accounting for phases offsets between antenna elements
    605         """
    606 
    607         lambda_0 = self._lambda_0
    608 
    609         r_hat_rx = self._unit_sphere_vector(zoa, aoa)
    610         r_hat_rx = tf.squeeze(r_hat_rx, axis=r_hat_rx.shape.rank-1)
    611         r_hat_tx = self._unit_sphere_vector(zod, aod)
    612         r_hat_tx = tf.squeeze(r_hat_tx, axis=r_hat_tx.shape.rank-1)
    613         d_bar_rx = self._step_11_get_rx_antenna_positions(topology)
    614         d_bar_tx = self._step_11_get_tx_antenna_positions(topology)
    615 
    616         # Reshape tensors for broadcasting
    617         # r_hat_rx/tx have
    618         # shape [batch_size, num_tx, num_rx, num_clusters, num_rays,    3]
    619         # and will be reshaoed to
    620         # [batch_size, num_tx, num_rx, num_clusters, num_rays, 1, 3]
    621         r_hat_tx = tf.expand_dims(r_hat_tx, -2)
    622         r_hat_rx = tf.expand_dims(r_hat_rx, -2)
    623 
    624         # d_bar_tx has shape [batch_size, num_tx,          num_tx_antennas, 3]
    625         # and will be reshaped to
    626         # [batch_size, num_tx, 1, 1, 1, num_tx_antennas, 3]
    627         s = tf.shape(d_bar_tx)
    628         shape = tf.concat([s[:2], [1,1,1], s[2:]], 0)
    629         d_bar_tx = tf.reshape(d_bar_tx, shape)
    630 
    631         # d_bar_rx has shape [batch_size,    num_rx,       num_rx_antennas, 3]
    632         # and will be reshaped to
    633         # [batch_size, 1, num_rx, 1, 1, num_rx_antennas, 3]
    634         s = tf.shape(d_bar_rx)
    635         shape = tf.concat([[s[0]], [1, s[1], 1,1], s[2:]], 0)
    636         d_bar_rx = tf.reshape(d_bar_rx, shape)
    637 
    638         # Compute all tensor elements
    639 
    640         # As broadcasting of such high-rank tensors is not fully supported
    641         # in all cases, we need to do a hack here by explicitly
    642         # broadcasting one dimension:
    643         s = tf.shape(d_bar_rx)
    644         shape = tf.concat([ [s[0]], [tf.shape(r_hat_rx)[1]], s[2:]], 0)
    645         d_bar_rx = tf.broadcast_to(d_bar_rx, shape)
    646         exp_rx = 2*PI/lambda_0*tf.reduce_sum(r_hat_rx*d_bar_rx,
    647             axis=-1, keepdims=True)
    648         exp_rx = tf.exp(tf.complex(tf.constant(0.,
    649                                     self._dtype.real_dtype), exp_rx))
    650 
    651         # The hack is for some reason not needed for this term
    652         # exp_tx = 2*PI/lambda_0*tf.reduce_sum(r_hat_tx*d_bar_tx,
    653         #     axis=-1, keepdims=True)
    654         exp_tx = 2*PI/lambda_0*tf.reduce_sum(r_hat_tx*d_bar_tx,
    655             axis=-1)
    656         exp_tx = tf.exp(tf.complex(tf.constant(0.,
    657                                     self._dtype.real_dtype), exp_tx))
    658         exp_tx = tf.expand_dims(exp_tx, -2)
    659 
    660         h_array = exp_rx*exp_tx
    661 
    662         return h_array
    663 
    664     def _step_11_field_matrix(self, topology, aoa, aod, zoa, zod, h_phase):
    665         # pylint: disable=line-too-long
    666         r"""
    667         Compute matrix accounting for the element responses, random phases
    668         and xpr
    669 
    670         Input
    671         -----
    672         topology : Topology
    673             Topology of the network
    674 
    675         aoa : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    676             Azimuth angles of arrivals [radian]
    677 
    678         aod : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    679             Azimuth angles of departure [radian]
    680 
    681         zoa : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    682             Zenith angles of arrivals [radian]
    683 
    684         zod : [batch size, num TXs, num RXs, num clusters, num rays], tf.float
    685             Zenith angles of departure [radian]
    686 
    687         h_phase : [batch size, num_tx, num rx, num clusters, num rays, num time steps], tf.complex
    688             Matrix with phase shifts due to mobility in (7.5-22)
    689 
    690         Output
    691         ------
    692         h_field : [batch size, num_tx, num rx, num clusters, num rays, num rx antennas, num tx antennas], tf.complex
    693             Matrix accounting for element responses, random phases and xpr
    694         """
    695 
    696         tx_orientations = topology.tx_orientations
    697         rx_orientations = topology.rx_orientations
    698 
    699         # Transform departure angles to the LCS
    700         s = tf.shape(tx_orientations)
    701         shape = tf.concat([s[:2], [1,1,1,s[-1]]], 0)
    702         tx_orientations = tf.reshape(tx_orientations, shape)
    703         zod_prime, aod_prime = self._gcs_to_lcs(tx_orientations, zod, aod)
    704 
    705         # Transform arrival angles to the LCS
    706         s = tf.shape(rx_orientations)
    707         shape = tf.concat([[s[0],1],[s[1],1,1,s[-1]]], 0)
    708         rx_orientations = tf.reshape(rx_orientations, shape)
    709         zoa_prime, aoa_prime = self._gcs_to_lcs(rx_orientations, zoa, aoa)
    710 
    711         # Compute transmitted and received field strength for all antennas
    712         # in the LCS  and convert to GCS
    713         f_tx_pol1_prime = tf.stack(self._tx_array.ant_pol1.field(zod_prime,
    714                                                             aod_prime), axis=-1)
    715         f_rx_pol1_prime = tf.stack(self._rx_array.ant_pol1.field(zoa_prime,
    716                                                             aoa_prime), axis=-1)
    717 
    718         f_tx_pol1 = self._l2g_response(f_tx_pol1_prime, tx_orientations,
    719             zod, aod)
    720 
    721         f_rx_pol1 = self._l2g_response(f_rx_pol1_prime, rx_orientations,
    722             zoa, aoa)
    723 
    724         if self._tx_array.polarization == 'dual':
    725             f_tx_pol2_prime = tf.stack(self._tx_array.ant_pol2.field(
    726                 zod_prime, aod_prime), axis=-1)
    727             f_tx_pol2 = self._l2g_response(f_tx_pol2_prime, tx_orientations,
    728                 zod, aod)
    729 
    730         if self._rx_array.polarization == 'dual':
    731             f_rx_pol2_prime = tf.stack(self._rx_array.ant_pol2.field(
    732                 zoa_prime, aoa_prime), axis=-1)
    733             f_rx_pol2 = self._l2g_response(f_rx_pol2_prime, rx_orientations,
    734                 zoa, aoa)
    735 
    736         # Fill the full channel matrix with field responses
    737         pol1_tx = tf.matmul(h_phase, tf.complex(f_tx_pol1,
    738             tf.constant(0., self._dtype.real_dtype)))
    739         if self._tx_array.polarization == 'dual':
    740             pol2_tx = tf.matmul(h_phase, tf.complex(f_tx_pol2, tf.constant(0.,
    741                                             self._dtype.real_dtype)))
    742 
    743         num_ant_tx = self._tx_array.num_ant
    744         if self._tx_array.polarization == 'single':
    745             # Each BS antenna gets the polarization 1 response
    746             f_tx_array = tf.tile(tf.expand_dims(pol1_tx, 0),
    747                 tf.concat([[num_ant_tx], tf.ones([tf.rank(pol1_tx)], tf.int32)],
    748                 axis=0))
    749         else:
    750             # Assign polarization reponse according to polarization to each
    751             # antenna
    752             pol_tx = tf.stack([pol1_tx, pol2_tx], 0) # pylint: disable=possibly-used-before-assignment
    753             ant_ind_pol2 = self._tx_array.ant_ind_pol2
    754             num_ant_pol2 = ant_ind_pol2.shape[0]
    755             # O = Pol 1, 1 = Pol 2, we only scatter the indices for Pol 1,
    756             # the other elements are already 0
    757             gather_ind = tf.scatter_nd(tf.reshape(ant_ind_pol2, [-1,1]),
    758                 tf.ones([num_ant_pol2], tf.int32), [num_ant_tx])
    759             f_tx_array = tf.gather(pol_tx, gather_ind, axis=0)
    760 
    761         num_ant_rx = self._rx_array.num_ant
    762         if self._rx_array.polarization == 'single':
    763             # Each UT antenna gets the polarization 1 response
    764             f_rx_array = tf.tile(tf.expand_dims(f_rx_pol1, 0),
    765                 tf.concat([[num_ant_rx], tf.ones([tf.rank(f_rx_pol1)],
    766                                                  tf.int32)], axis=0))
    767             f_rx_array = tf.complex(f_rx_array,
    768                                     tf.constant(0., self._dtype.real_dtype))
    769         else:
    770             # Assign polarization response according to polarization to each
    771             # antenna
    772             pol_rx = tf.stack([f_rx_pol1, f_rx_pol2], 0) # pylint: disable=possibly-used-before-assignment
    773             ant_ind_pol2 = self._rx_array.ant_ind_pol2
    774             num_ant_pol2 = ant_ind_pol2.shape[0]
    775             # O = Pol 1, 1 = Pol 2, we only scatter the indices for Pol 1,
    776             # the other elements are already 0
    777             gather_ind = tf.scatter_nd(tf.reshape(ant_ind_pol2, [-1,1]),
    778                 tf.ones([num_ant_pol2], tf.int32), [num_ant_rx])
    779             f_rx_array = tf.complex(tf.gather(pol_rx, gather_ind, axis=0),
    780                             tf.constant(0., self._dtype.real_dtype))
    781 
    782         # Compute the scalar product between the field vectors through
    783         # reduce_sum and transpose to put antenna dimensions last
    784         h_field = tf.reduce_sum(tf.expand_dims(f_rx_array, 1)*tf.expand_dims(
    785             f_tx_array, 0), [-2,-1])
    786         h_field = tf.transpose(h_field, tf.roll(tf.range(tf.rank(h_field)),
    787             -2, 0))
    788 
    789         return h_field
    790 
    791     def _step_11_nlos(self, phi, topology, rays, t):
    792         # pylint: disable=line-too-long
    793         r"""
    794         Compute the full NLOS channel matrix (7.5-28)
    795 
    796         Input
    797         -----
    798         phi: [batch size, num TXs, num RXs, num clusters, num rays, 4], tf.float
    799             Random initial phases [radian]
    800 
    801         topology : Topology
    802             Topology of the network
    803 
    804         rays : Rays
    805             Rays
    806 
    807         t : [num time samples], tf.float
    808             Time samples
    809 
    810         Output
    811         ------
    812         h_full : [batch size, num_tx, num rx, num clusters, num rays, num rx antennas, num tx antennas, num time steps], tf.complex
    813             NLoS channel matrix
    814         """
    815 
    816         h_phase = self._step_11_phase_matrix(phi, rays)
    817         h_field = self._step_11_field_matrix(topology, rays.aoa, rays.aod,
    818                                                     rays.zoa, rays.zod, h_phase)
    819         h_array = self._step_11_array_offsets(topology, rays.aoa, rays.aod,
    820                                                             rays.zoa, rays.zod)
    821         h_doppler = self._step_11_doppler_matrix(topology, rays.aoa, rays.zoa,
    822                                                                             t)
    823         h_full = tf.expand_dims(h_field*h_array, -1) * tf.expand_dims(
    824             tf.expand_dims(h_doppler, -2), -2)
    825 
    826         power_scaling = tf.complex(tf.sqrt(rays.powers/
    827             tf.cast(tf.shape(h_full)[4], self._dtype.real_dtype)),
    828                             tf.constant(0., self._dtype.real_dtype))
    829         shape = tf.concat([tf.shape(power_scaling), tf.ones(
    830             [tf.rank(h_full)-tf.rank(power_scaling)], tf.int32)], 0)
    831         h_full *= tf.reshape(power_scaling, shape)
    832 
    833         return h_full
    834 
    835     def _step_11_reduce_nlos(self, h_full, rays, c_ds):
    836         # pylint: disable=line-too-long
    837         r"""
    838         Compute the final NLOS matrix in (7.5-27)
    839 
    840         Input
    841         ------
    842         h_full : [batch size, num_tx, num rx, num clusters, num rays, num rx antennas, num tx antennas, num time steps], tf.complex
    843             NLoS channel matrix
    844 
    845         rays : Rays
    846             Rays
    847 
    848         c_ds : [batch size, num TX, num RX], tf.float
    849             Cluster delay spread
    850 
    851         Output
    852         -------
    853         h_nlos : [batch size, num_tx, num rx, num clusters, num rx antennas, num tx antennas, num time steps], tf.complex
    854             Paths NLoS coefficients
    855 
    856         delays_nlos : [batch size, num_tx, num rx, num clusters], tf.float
    857             Paths NLoS delays
    858         """
    859 
    860         if self._subclustering:
    861 
    862             powers = rays.powers
    863             delays = rays.delays
    864 
    865             # Sort all clusters along their power
    866             strongest_clusters = tf.argsort(powers, axis=-1,
    867                 direction="DESCENDING")
    868 
    869             # Sort delays according to the same ordering
    870             delays_sorted = tf.gather(delays, strongest_clusters,
    871                 batch_dims=3, axis=3)
    872 
    873             # Split into delays for strong and weak clusters
    874             delays_strong = delays_sorted[...,:2]
    875             delays_weak = delays_sorted[...,2:]
    876 
    877             # Compute delays for sub-clusters
    878             offsets = tf.reshape(self._sub_cl_delay_offsets,
    879                 (delays_strong.shape.rank-1)*[1]+[-1]+[1])
    880             delays_sub_cl = (tf.expand_dims(delays_strong, -2) +
    881                 offsets*tf.expand_dims(tf.expand_dims(c_ds, axis=-1), axis=-1))
    882             delays_sub_cl = tf.reshape(delays_sub_cl,
    883                 tf.concat([tf.shape(delays_sub_cl)[:-2], [-1]],0))
    884 
    885             # Select the strongest two clusters for sub-cluster splitting
    886             h_strong = tf.gather(h_full, strongest_clusters[...,:2],
    887                 batch_dims=3, axis=3)
    888 
    889             # The other clusters are the weak clusters
    890             h_weak = tf.gather(h_full, strongest_clusters[...,2:],
    891                 batch_dims=3, axis=3)
    892 
    893             # Sum specific rays for each sub-cluster
    894             h_sub_cl_1 = tf.reduce_sum(tf.gather(h_strong,
    895                 self._sub_cl_1_ind, axis=4), axis=4)
    896             h_sub_cl_2 = tf.reduce_sum(tf.gather(h_strong,
    897                 self._sub_cl_2_ind, axis=4), axis=4)
    898             h_sub_cl_3 = tf.reduce_sum(tf.gather(h_strong,
    899                 self._sub_cl_3_ind, axis=4), axis=4)
    900 
    901             # Sum all rays for the weak clusters
    902             h_weak = tf.reduce_sum(h_weak, axis=4)
    903 
    904             # Concatenate the channel and delay tensors
    905             h_nlos = tf.concat([h_sub_cl_1, h_sub_cl_2, h_sub_cl_3, h_weak],
    906                 axis=3)
    907             delays_nlos = tf.concat([delays_sub_cl, delays_weak], axis=3)
    908         else:
    909             # Sum over rays
    910             h_nlos = tf.reduce_sum(h_full, axis=4)
    911             delays_nlos = rays.delays
    912 
    913         # Order the delays in ascending orders
    914         delays_ind = tf.argsort(delays_nlos, axis=-1,
    915             direction="ASCENDING")
    916         delays_nlos = tf.gather(delays_nlos, delays_ind, batch_dims=3,
    917             axis=3)
    918 
    919         # # Order the channel clusters according to the delay, too
    920         h_nlos = tf.gather(h_nlos, delays_ind, batch_dims=3, axis=3)
    921 
    922         return h_nlos, delays_nlos
    923 
    924     def _step_11_los(self, topology, t):
    925         # pylint: disable=line-too-long
    926         r"""Compute the LOS channels from (7.5-29)
    927 
    928         Intput
    929         ------
    930         topology : Topology
    931             Network topology
    932 
    933         t : [num time samples], tf.float
    934             Number of time samples
    935 
    936         Output
    937         ------
    938         h_los : [batch size, num_tx, num rx, 1, num rx antennas, num tx antennas, num time steps], tf.complex
    939             Paths LoS coefficients
    940         """
    941 
    942         aoa = topology.los_aoa
    943         aod = topology.los_aod
    944         zoa = topology.los_zoa
    945         zod = topology.los_zod
    946 
    947          # LoS departure and arrival angles
    948         aoa = tf.expand_dims(tf.expand_dims(aoa, axis=3), axis=4)
    949         zoa = tf.expand_dims(tf.expand_dims(zoa, axis=3), axis=4)
    950         aod = tf.expand_dims(tf.expand_dims(aod, axis=3), axis=4)
    951         zod = tf.expand_dims(tf.expand_dims(zod, axis=3), axis=4)
    952 
    953         # Field matrix
    954         h_phase = tf.reshape(tf.constant([[1.,0.],
    955                                          [0.,-1.]],
    956                                          self._dtype),
    957                                          [1,1,1,1,1,2,2])
    958         h_field = self._step_11_field_matrix(topology, aoa, aod, zoa, zod,
    959                                                                     h_phase)
    960 
    961         # Array offset matrix
    962         h_array = self._step_11_array_offsets(topology, aoa, aod, zoa, zod)
    963 
    964         # Doppler matrix
    965         h_doppler = self._step_11_doppler_matrix(topology, aoa, zoa, t)
    966 
    967         # Phase shift due to propagation delay
    968         d3d = topology.distance_3d
    969         lambda_0 = self._lambda_0
    970         h_delay = tf.exp(tf.complex(tf.constant(0.,
    971                         self._dtype.real_dtype), 2*PI*d3d/lambda_0))
    972 
    973         # Combining all to compute channel coefficient
    974         h_field = tf.expand_dims(tf.squeeze(h_field, axis=4), axis=-1)
    975         h_array = tf.expand_dims(tf.squeeze(h_array, axis=4), axis=-1)
    976         h_doppler = tf.expand_dims(h_doppler, axis=4)
    977         h_delay = tf.expand_dims(tf.expand_dims(tf.expand_dims(
    978             tf.expand_dims(h_delay, axis=3), axis=4), axis=5), axis=6)
    979 
    980         h_los = h_field*h_array*h_doppler*h_delay
    981         return h_los
    982 
    983     def _step_11(self, phi, topology, k_factor, rays, t, c_ds):
    984         # pylint: disable=line-too-long
    985         r"""
    986         Combine LOS and LOS components to compute (7.5-30)
    987 
    988         Input
    989         -----
    990         phi: [batch size, num TXs, num RXs, num clusters, num rays, 4], tf.float
    991             Random initial phases
    992 
    993         topology : Topology
    994             Network topology
    995 
    996         k_factor : [batch size, num TX, num RX], tf.float
    997             Rician K-factor
    998 
    999         rays : Rays
   1000             Rays
   1001 
   1002         t : [num time samples], tf.float
   1003             Number of time samples
   1004 
   1005         c_ds : [batch size, num TX, num RX], tf.float
   1006             Cluster delay spread
   1007         """
   1008 
   1009         h_full = self._step_11_nlos(phi, topology, rays, t)
   1010         h_nlos, delays_nlos = self._step_11_reduce_nlos(h_full, rays, c_ds)
   1011 
   1012         ####  LoS scenario
   1013 
   1014         h_los_los_comp = self._step_11_los(topology, t)
   1015         k_factor = tf.reshape(k_factor, tf.concat([tf.shape(k_factor),
   1016             tf.ones([tf.rank(h_los_los_comp)-tf.rank(k_factor)], tf.int32)],0))
   1017         k_factor = tf.complex(k_factor, tf.constant(0.,
   1018                                             self._dtype.real_dtype))
   1019 
   1020         # Scale NLOS and LOS components according to K-factor
   1021         h_los_los_comp = h_los_los_comp*tf.sqrt(k_factor/(k_factor+1))
   1022         h_los_nlos_comp = h_nlos*tf.sqrt(1/(k_factor+1))
   1023 
   1024         # Add the LOS component to the zero-delay NLOS cluster
   1025         h_los_cl = h_los_los_comp + tf.expand_dims(
   1026             h_los_nlos_comp[:,:,:,0,...], 3)
   1027 
   1028         # Combine all clusters into a single tensor
   1029         h_los = tf.concat([h_los_cl, h_los_nlos_comp[:,:,:,1:,...]], axis=3)
   1030 
   1031         #### LoS or NLoS CIR according to link configuration
   1032         los_indicator = tf.reshape(topology.los,
   1033             tf.concat([tf.shape(topology.los), [1,1,1,1]], axis=0))
   1034         h = tf.where(los_indicator, h_los, h_nlos)
   1035 
   1036         return h, delays_nlos