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

rays.py (33182B)


      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 rays following 3GPP TR38.901 specifications and giving a
      7 channel simulation scenario and LSPs.
      8 """
      9 
     10 import tensorflow as tf
     11 
     12 from sionna import config
     13 from sionna.utils import log10
     14 from sionna.channel.utils import deg_2_rad, wrap_angle_0_360
     15 
     16 class Rays:
     17     # pylint: disable=line-too-long
     18     r"""
     19     Class for conveniently storing rays
     20 
     21     Parameters
     22     -----------
     23 
     24     delays : [batch size, number of BSs, number of UTs, number of clusters], tf.float
     25         Paths delays [s]
     26 
     27     powers : [batch size, number of BSs, number of UTs, number of clusters], tf.float
     28         Normalized path powers
     29 
     30     aoa : (batch size, number of BSs, number of UTs, number of clusters, number of rays], tf.float
     31         Azimuth angles of arrival [radian]
     32 
     33     aod : [batch size, number of BSs, number of UTs, number of clusters, number of rays], tf.float
     34         Azimuth angles of departure [radian]
     35 
     36     zoa : [batch size, number of BSs, number of UTs, number of clusters, number of rays], tf.float
     37         Zenith angles of arrival [radian]
     38 
     39     zod : [batch size, number of BSs, number of UTs, number of clusters, number of rays], tf.float
     40         Zenith angles of departure [radian]
     41 
     42     xpr [batch size, number of BSs, number of UTs, number of clusters, number of rays], tf.float
     43         Coss-polarization power ratios.
     44     """
     45 
     46     def __init__(self, delays, powers, aoa, aod, zoa, zod, xpr):
     47         self.delays = delays
     48         self.powers = powers
     49         self.aoa = aoa
     50         self.aod = aod
     51         self.zoa = zoa
     52         self.zod = zod
     53         self.xpr = xpr
     54 
     55 
     56 class RaysGenerator:
     57     """
     58     Sample rays according to a given channel scenario and large scale
     59     parameters (LSP).
     60 
     61     This class implements steps 6 to 9 from the TR 38.901 specifications,
     62     (section 7.5).
     63 
     64     Note that a global scenario is set for the entire batches when instantiating
     65     this class (UMa, UMi, or RMa). However, each UT-BS link can have its
     66     specific state (LoS, NLoS, or indoor).
     67 
     68     The batch size is set by the ``scenario`` given as argument when
     69     constructing the class.
     70 
     71     Parameters
     72     ----------
     73     scenario : :class:`~sionna.channel.tr38901.SystemLevelScenario``
     74         Scenario used to generate LSPs
     75 
     76     Input
     77     -----
     78     lsp : :class:`~sionna.channel.tr38901.LSP`
     79         LSPs samples
     80 
     81     Output
     82     ------
     83     rays : :class:`~sionna.channel.tr38901.Rays`
     84         Rays samples
     85     """
     86 
     87     def __init__(self, scenario):
     88         # Scenario
     89         self._scenario = scenario
     90 
     91         # For AoA, AoD, ZoA, and ZoD, offset to add to cluster angles to get ray
     92         # angles. This is hardcoded from table 7.5-3 for 3GPP 38.901
     93         # specification.
     94         self._ray_offsets = tf.constant([0.0447, -0.0447,
     95                                          0.1413, -0.1413,
     96                                          0.2492, -0.2492,
     97                                          0.3715, -0.3715,
     98                                          0.5129, -0.5129,
     99                                          0.6797, -0.6797,
    100                                          0.8844, -0.8844,
    101                                          1.1481, -0.1481,
    102                                          1.5195, -1.5195,
    103                                          2.1551, -2.1551],
    104                                          self._scenario.dtype.real_dtype)
    105 
    106     #########################################
    107     # Public methods and properties
    108     #########################################
    109 
    110     def __call__(self, lsp):
    111         # Sample cluster delays
    112         delays, delays_unscaled = self._cluster_delays(lsp.ds, lsp.k_factor)
    113 
    114         # Sample cluster powers
    115         powers, powers_for_angles_gen = self._cluster_powers(lsp.ds,
    116                                             lsp.k_factor, delays_unscaled)
    117 
    118         # Sample AoA
    119         aoa = self._azimuth_angles_of_arrival(lsp.asa, lsp.k_factor,
    120                                                 powers_for_angles_gen)
    121 
    122         # Sample AoD
    123         aod = self._azimuth_angles_of_departure(lsp.asd, lsp.k_factor,
    124                                                 powers_for_angles_gen)
    125 
    126         # Sample ZoA
    127         zoa = self._zenith_angles_of_arrival(lsp.zsa, lsp.k_factor,
    128                                                 powers_for_angles_gen)
    129 
    130         # Sample ZoD
    131         zod = self._zenith_angles_of_departure(lsp.zsd, lsp.k_factor,
    132                                                 powers_for_angles_gen)
    133 
    134         # XPRs
    135         xpr = self._cross_polarization_power_ratios()
    136 
    137         # Random coupling
    138         aoa, aod, zoa, zod = self._random_coupling(aoa, aod, zoa, zod)
    139 
    140         # Convert angles of arrival and departure from degree to radian
    141         aoa = deg_2_rad(aoa)
    142         aod = deg_2_rad(aod)
    143         zoa = deg_2_rad(zoa)
    144         zod = deg_2_rad(zod)
    145 
    146         # Storing and returning rays
    147         rays = Rays(delays = delays,
    148                     powers = powers,
    149                     aoa    = aoa,
    150                     aod    = aod,
    151                     zoa    = zoa,
    152                     zod    = zod,
    153                     xpr    = xpr)
    154 
    155         return rays
    156 
    157     def topology_updated_callback(self):
    158         """
    159         Updates internal quantities. Must be called at every update of the
    160         scenario that changes the state of UTs or their locations.
    161 
    162         Input
    163         ------
    164         None
    165 
    166         Output
    167         ------
    168         None
    169         """
    170         self._compute_clusters_mask()
    171 
    172     ########################################
    173     # Internal utility methods
    174     ########################################
    175 
    176     def _compute_clusters_mask(self):
    177         """
    178         Given a scenario (UMi, UMa, RMa), the number of clusters is different
    179         for different state of UT-BS links (LoS, NLoS, indoor).
    180 
    181         Because we use tensors with predefined dimension size (not ragged), the
    182         cluster dimension is always set to the maximum number of clusters the
    183         scenario requires. A mask is then used to discard not required tensors,
    184         depending on the state of each UT-BS link.
    185 
    186         This function computes and stores this mask of size
    187         [batch size, number of BSs, number of UTs, maximum number of cluster]
    188         where an element equals 0 if the cluster is used, 1 otherwise.
    189         """
    190 
    191         scenario = self._scenario
    192         num_clusters_los = scenario.num_clusters_los
    193         num_clusters_nlos = scenario.num_clusters_nlos
    194         num_clusters_o2i = scenario.num_clusters_indoor
    195         num_clusters_max = tf.reduce_max([num_clusters_los, num_clusters_nlos,
    196             num_clusters_o2i])
    197 
    198 
    199         # Initialize an empty mask
    200         mask = tf.zeros(shape=[scenario.batch_size, scenario.num_bs,
    201             scenario.num_ut, num_clusters_max],
    202             dtype=self._scenario.dtype.real_dtype)
    203 
    204         # Indoor mask
    205         mask_indoor = tf.concat((tf.zeros([num_clusters_o2i],
    206                                           self._scenario.dtype.real_dtype),
    207                                  tf.ones([num_clusters_max-num_clusters_o2i],
    208                                     self._scenario.dtype.real_dtype)), axis=0)
    209         mask_indoor = tf.reshape(mask_indoor, [1, 1, 1, num_clusters_max])
    210         indoor = tf.expand_dims(scenario.indoor, axis=1) # Broadcasting with BS
    211         o2i_slice_mask = tf.cast(indoor, self._scenario.dtype.real_dtype)
    212         o2i_slice_mask = tf.expand_dims(o2i_slice_mask, axis=3)
    213         mask = mask + o2i_slice_mask*mask_indoor
    214 
    215         # LoS
    216         mask_los = tf.concat([tf.zeros([num_clusters_los],
    217             self._scenario.dtype.real_dtype),
    218             tf.ones([num_clusters_max-num_clusters_los],
    219             self._scenario.dtype.real_dtype)], axis=0)
    220         mask_los = tf.reshape(mask_los, [1, 1, 1, num_clusters_max])
    221         los_slice_mask = scenario.los
    222         los_slice_mask = tf.cast(los_slice_mask,
    223                                     self._scenario.dtype.real_dtype)
    224         los_slice_mask = tf.expand_dims(los_slice_mask, axis=3)
    225         mask = mask + los_slice_mask*mask_los
    226 
    227         # NLoS
    228         mask_nlos = tf.concat([tf.zeros([num_clusters_nlos],
    229             self._scenario.dtype.real_dtype),
    230             tf.ones([num_clusters_max-num_clusters_nlos],
    231             self._scenario.dtype.real_dtype)], axis=0)
    232         mask_nlos = tf.reshape(mask_nlos, [1, 1, 1, num_clusters_max])
    233         nlos_slice_mask = tf.logical_and(tf.logical_not(scenario.los),
    234             tf.logical_not(indoor))
    235         nlos_slice_mask = tf.cast(nlos_slice_mask,
    236                                     self._scenario.dtype.real_dtype)
    237         nlos_slice_mask = tf.expand_dims(nlos_slice_mask, axis=3)
    238         mask = mask + nlos_slice_mask*mask_nlos
    239 
    240         # Save the mask
    241         self._cluster_mask = mask
    242 
    243     def _cluster_delays(self, delay_spread, rician_k_factor):
    244         # pylint: disable=line-too-long
    245         """
    246         Generate cluster delays.
    247         See step 5 of section 7.5 from TR 38.901 specification.
    248 
    249         Input
    250         ------
    251         delay_spread : [batch size, num of BSs, num of UTs], tf.float
    252             RMS delay spread of each BS-UT link.
    253 
    254         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    255             Rician K-factor of each BS-UT link. Used only for LoS links.
    256 
    257         Output
    258         -------
    259         delays : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    260             Path delays [s]
    261 
    262         unscaled_delays [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    263             Unscaled path delays [s]
    264         """
    265 
    266         scenario = self._scenario
    267 
    268         batch_size = scenario.batch_size
    269         num_bs = scenario.num_bs
    270         num_ut = scenario.num_ut
    271 
    272         num_clusters_max = scenario.num_clusters_max
    273 
    274         # Getting scaling parameter according to each BS-UT link scenario
    275         delay_scaling_parameter = scenario.get_param("rTau")
    276         delay_scaling_parameter = tf.expand_dims(delay_scaling_parameter,
    277             axis=3)
    278 
    279         # Generating random cluster delays
    280         # We don't start at 0 to avoid numerical errors
    281         delay_spread = tf.expand_dims(delay_spread, axis=3)
    282         x = config.tf_rng.uniform(shape=[batch_size, num_bs, num_ut,
    283                                          num_clusters_max],
    284                                   minval=1e-6, maxval=1.0,
    285             dtype=self._scenario.dtype.real_dtype)
    286 
    287         # Moving to linear domain
    288         unscaled_delays = -delay_scaling_parameter*delay_spread*tf.math.log(x)
    289         # Forcing the cluster that should not exist to huge delays (1s)
    290         unscaled_delays = (unscaled_delays*(1.-self._cluster_mask)
    291             + self._cluster_mask)
    292 
    293         # Normalizing and sorting the delays
    294         unscaled_delays = unscaled_delays - tf.reduce_min(unscaled_delays,
    295             axis=3, keepdims=True)
    296         unscaled_delays = tf.sort(unscaled_delays, axis=3)
    297 
    298         # Additional scaling applied to LoS links
    299         rician_k_factor_db = 10.0*log10(rician_k_factor) # to dB
    300         scaling_factor = (0.7705 - 0.0433*rician_k_factor_db
    301             + 0.0002*tf.square(rician_k_factor_db)
    302             + 0.000017*tf.math.pow(rician_k_factor_db, tf.constant(3.,
    303             self._scenario.dtype.real_dtype)))
    304         scaling_factor = tf.expand_dims(scaling_factor, axis=3)
    305         delays = tf.where(tf.expand_dims(scenario.los, axis=3),
    306             unscaled_delays / scaling_factor, unscaled_delays)
    307 
    308         return delays, unscaled_delays
    309 
    310     def _cluster_powers(self, delay_spread, rician_k_factor, unscaled_delays):
    311         # pylint: disable=line-too-long
    312         """
    313         Generate cluster powers.
    314         See step 6 of section 7.5 from TR 38.901 specification.
    315 
    316         Input
    317         ------
    318         delays : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    319             Path delays [s]
    320 
    321         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    322             Rician K-factor of each BS-UT link. Used only for LoS links.
    323 
    324         unscaled_delays [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    325             Unscaled path delays [s]. Required to compute the path powers.
    326 
    327         Output
    328         -------
    329         powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    330             Normalized path powers
    331         """
    332 
    333         scenario = self._scenario
    334 
    335         batch_size = scenario.batch_size
    336         num_bs = scenario.num_bs
    337         num_ut = scenario.num_ut
    338 
    339         num_clusters_max = scenario.num_clusters_max
    340 
    341         delay_scaling_parameter = scenario.get_param("rTau")
    342         cluster_shadowing_std_db = scenario.get_param("zeta")
    343         delay_spread = tf.expand_dims(delay_spread, axis=3)
    344         cluster_shadowing_std_db = tf.expand_dims(cluster_shadowing_std_db,
    345             axis=3)
    346         delay_scaling_parameter = tf.expand_dims(delay_scaling_parameter,
    347             axis=3)
    348 
    349         # Generate unnormalized cluster powers
    350         z = config.tf_rng.normal(shape=[batch_size, num_bs, num_ut,
    351             num_clusters_max], mean=0.0, stddev=cluster_shadowing_std_db,
    352             dtype=self._scenario.dtype.real_dtype)
    353 
    354         # Moving to linear domain
    355         powers_unnormalized = (tf.math.exp(-unscaled_delays*
    356             (delay_scaling_parameter - 1.0)/
    357             (delay_scaling_parameter*delay_spread))*tf.math.pow(tf.constant(10.,
    358             self._scenario.dtype.real_dtype), -z/10.0))
    359 
    360         # Force the power of unused cluster to zero
    361         powers_unnormalized = powers_unnormalized*(1.-self._cluster_mask)
    362 
    363         # Normalizing cluster powers
    364         powers = (powers_unnormalized/
    365             tf.reduce_sum(powers_unnormalized, axis=3, keepdims=True))
    366 
    367         # Additional specular component for LoS
    368         rician_k_factor = tf.expand_dims(rician_k_factor, axis=3)
    369         p_nlos_scaling = 1.0/(rician_k_factor + 1.0)
    370         p_1_los = rician_k_factor*p_nlos_scaling
    371         powers_1 = p_nlos_scaling*powers[:,:,:,:1] + p_1_los
    372         powers_n = p_nlos_scaling*powers[:,:,:,1:]
    373         powers_for_angles_gen = tf.where(tf.expand_dims(scenario.los, axis=3),
    374             tf.concat([powers_1, powers_n], axis=3), powers)
    375 
    376         return powers, powers_for_angles_gen
    377 
    378     def _azimuth_angles(self, azimuth_spread, rician_k_factor, cluster_powers,
    379                         angle_type):
    380         # pylint: disable=line-too-long
    381         """
    382         Generate departure or arrival azimuth angles (degrees).
    383         See step 7 of section 7.5 from TR 38.901 specification.
    384 
    385         Input
    386         ------
    387         azimuth_spread : [batch size, num of BSs, num of UTs], tf.float
    388             Angle spread, (ASD or ASA) depending on ``angle_type`` [deg]
    389 
    390         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    391             Rician K-factor of each BS-UT link. Used only for LoS links.
    392 
    393         cluster_powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    394             Normalized path powers
    395 
    396         angle_type : str
    397             Type of angle to compute. Must be 'aoa' or 'aod'.
    398 
    399         Output
    400         -------
    401         azimuth_angles : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    402             Paths azimuth angles wrapped within (-180, 180) [degree]. Either the AoA or AoD depending on ``angle_type``.
    403         """
    404 
    405         scenario = self._scenario
    406 
    407         batch_size = scenario.batch_size
    408         num_bs = scenario.num_bs
    409         num_ut = scenario.num_ut
    410 
    411         num_clusters_max = scenario.num_clusters_max
    412 
    413         azimuth_spread = tf.expand_dims(azimuth_spread, axis=3)
    414 
    415         # Loading the angle spread
    416         if angle_type == 'aod':
    417             azimuth_angles_los = scenario.los_aod
    418             cluster_angle_spread = scenario.get_param('cASD')
    419         else:
    420             azimuth_angles_los = scenario.los_aoa
    421             cluster_angle_spread = scenario.get_param('cASA')
    422         # Adding cluster dimension for broadcasting
    423         azimuth_angles_los = tf.expand_dims(azimuth_angles_los, axis=3)
    424         cluster_angle_spread = tf.expand_dims(tf.expand_dims(
    425             cluster_angle_spread, axis=3), axis=4)
    426 
    427         # Compute C-phi constant
    428         rician_k_factor = tf.expand_dims(rician_k_factor, axis=3)
    429         rician_k_factor_db = 10.0*log10(rician_k_factor) # to dB
    430         c_phi_nlos = tf.expand_dims(scenario.get_param("CPhiNLoS"), axis=3)
    431         c_phi_los = c_phi_nlos*(1.1035- 0.028*rician_k_factor_db
    432             - 0.002*tf.square(rician_k_factor_db)
    433             + 0.0001*tf.math.pow(rician_k_factor_db, 3.))
    434         c_phi = tf.where(tf.expand_dims(scenario.los, axis=3),
    435             c_phi_los, c_phi_nlos)
    436 
    437         # Inverse Gaussian function
    438         z = cluster_powers/tf.reduce_max(cluster_powers, axis=3, keepdims=True)
    439         z = tf.clip_by_value(z, 1e-6, 1.0)
    440         azimuth_angles_prime = (2.*azimuth_spread/1.4)*(tf.sqrt(-tf.math.log(z)
    441                                                                 )/c_phi)
    442 
    443         # Introducing random variation
    444         random_sign = config.tf_rng.uniform(shape=[batch_size, num_bs, 1,
    445             num_clusters_max], minval=0, maxval=2, dtype=tf.int32)
    446         random_sign = 2*random_sign - 1
    447         random_sign = tf.cast(random_sign, self._scenario.dtype.real_dtype)
    448         random_comp = config.tf_rng.normal(shape=[batch_size, num_bs, num_ut,
    449             num_clusters_max], mean=0.0, stddev=azimuth_spread/7.0,
    450             dtype=self._scenario.dtype.real_dtype)
    451         azimuth_angles = (random_sign*azimuth_angles_prime + random_comp
    452             + azimuth_angles_los)
    453         azimuth_angles = (azimuth_angles -
    454             tf.where(tf.expand_dims(scenario.los, axis=3),
    455             random_sign[:,:,:,:1]*azimuth_angles_prime[:,:,:,:1]
    456             + random_comp[:,:,:,:1], 0.0))
    457 
    458         # Add offset angles to cluster angles to get the ray angles
    459         ray_offsets = self._ray_offsets[:scenario.rays_per_cluster]
    460         # Add dimensions for batch size, num bs, num ut, num clusters
    461         ray_offsets = tf.reshape(ray_offsets, [1,1,1,1,
    462                                                 scenario.rays_per_cluster])
    463         # Rays angles
    464         azimuth_angles = tf.expand_dims(azimuth_angles, axis=4)
    465         azimuth_angles = azimuth_angles + cluster_angle_spread*ray_offsets
    466 
    467         # Wrapping to (-180, 180)
    468         azimuth_angles = wrap_angle_0_360(azimuth_angles)
    469         azimuth_angles = tf.where(tf.math.greater(azimuth_angles, 180.),
    470             azimuth_angles-360., azimuth_angles)
    471 
    472         return azimuth_angles
    473 
    474     def _azimuth_angles_of_arrival(self, azimuth_spread_arrival,
    475                                    rician_k_factor, cluster_powers):
    476         # pylint: disable=line-too-long
    477         """
    478         Compute the azimuth angle of arrival (AoA)
    479         See step 7 of section 7.5 from TR 38.901 specification.
    480 
    481         Input
    482         ------
    483         azimuth_spread_arrival : [batch size, num of BSs, num of UTs], tf.float
    484             Azimuth angle spread of arrival (ASA) [deg]
    485 
    486         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    487             Rician K-factor of each BS-UT link. Used only for LoS links.
    488 
    489         cluster_powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    490             Normalized path powers
    491 
    492         Output
    493         -------
    494         aoa : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    495             Paths azimuth angles of arrival (AoA) wrapped within (-180,180) [degree]
    496         """
    497         return self._azimuth_angles(azimuth_spread_arrival,
    498                                     rician_k_factor, cluster_powers, 'aoa')
    499 
    500     def _azimuth_angles_of_departure(self, azimuth_spread_departure,
    501                                      rician_k_factor, cluster_powers):
    502         # pylint: disable=line-too-long
    503         """
    504         Compute the azimuth angle of departure (AoD)
    505         See step 7 of section 7.5 from TR 38.901 specification.
    506 
    507         Input
    508         ------
    509         azimuth_spread_departure : [batch size, num of BSs, num of UTs], tf.float
    510             Azimuth angle spread of departure (ASD) [deg]
    511 
    512         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    513             Rician K-factor of each BS-UT link. Used only for LoS links.
    514 
    515         cluster_powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    516             Normalized path powers
    517 
    518         Output
    519         -------
    520         aod : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    521             Paths azimuth angles of departure (AoD) wrapped within (-180,180) [degree]
    522         """
    523         return self._azimuth_angles(azimuth_spread_departure,
    524                                     rician_k_factor, cluster_powers, 'aod')
    525 
    526     def _zenith_angles(self, zenith_spread, rician_k_factor, cluster_powers,
    527                        angle_type):
    528         # pylint: disable=line-too-long
    529         """
    530         Generate departure or arrival zenith angles (degrees).
    531         See step 7 of section 7.5 from TR 38.901 specification.
    532 
    533         Input
    534         ------
    535         zenith_spread : [batch size, num of BSs, num of UTs], tf.float
    536             Angle spread, (ZSD or ZSA) depending on ``angle_type`` [deg]
    537 
    538         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    539             Rician K-factor of each BS-UT link. Used only for LoS links.
    540 
    541         cluster_powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    542             Normalized path powers
    543 
    544         angle_type : str
    545             Type of angle to compute. Must be 'zoa' or 'zod'.
    546 
    547         Output
    548         -------
    549         zenith_angles : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    550             Paths zenith angles wrapped within (0,180) [degree]. Either the ZoA or ZoD depending on ``angle_type``.
    551         """
    552 
    553         scenario = self._scenario
    554 
    555         batch_size = scenario.batch_size
    556         num_bs = scenario.num_bs
    557         num_ut = scenario.num_ut
    558 
    559         # Tensors giving UTs states
    560         los = scenario.los
    561         indoor_uts = tf.expand_dims(scenario.indoor, axis=1)
    562         los_uts = tf.logical_and(los, tf.logical_not(indoor_uts))
    563         nlos_uts = tf.logical_and(tf.logical_not(los),
    564                             tf.logical_not(indoor_uts))
    565 
    566         num_clusters_max = scenario.num_clusters_max
    567 
    568         # Adding cluster dimension for broadcasting
    569         zenith_spread = tf.expand_dims(zenith_spread, axis=3)
    570         rician_k_factor = tf.expand_dims(rician_k_factor, axis=3)
    571         indoor_uts = tf.expand_dims(indoor_uts, axis=3)
    572         los_uts = tf.expand_dims(los_uts, axis=3)
    573         nlos_uts = tf.expand_dims(nlos_uts, axis=3)
    574 
    575         # Loading angle spread
    576         if angle_type == 'zod':
    577             zenith_angles_los = scenario.los_zod
    578             cluster_angle_spread = (3./8.)*tf.math.pow(tf.constant(10.,
    579                 self._scenario.dtype.real_dtype),
    580                 scenario.lsp_log_mean[:,:,:,6])
    581         else:
    582             cluster_angle_spread = scenario.get_param('cZSA')
    583             zenith_angles_los = scenario.los_zoa
    584         zod_offset = scenario.zod_offset
    585         # Adding cluster dimension for broadcasting
    586         zod_offset = tf.expand_dims(zod_offset, axis=3)
    587         zenith_angles_los = tf.expand_dims(zenith_angles_los, axis=3)
    588         cluster_angle_spread = tf.expand_dims(cluster_angle_spread, axis=3)
    589 
    590         # Compute the C_theta
    591         rician_k_factor_db = 10.0*log10(rician_k_factor) # to dB
    592         c_theta_nlos = tf.expand_dims(scenario.get_param("CThetaNLoS"),axis=3)
    593         c_theta_los = c_theta_nlos*(1.3086 + 0.0339*rician_k_factor_db
    594             - 0.0077*tf.square(rician_k_factor_db)
    595             + 0.0002*tf.math.pow(rician_k_factor_db, 3.))
    596         c_theta = tf.where(los_uts, c_theta_los, c_theta_nlos)
    597 
    598         # Inverse Laplacian function
    599         z = cluster_powers/tf.reduce_max(cluster_powers, axis=3, keepdims=True)
    600         z = tf.clip_by_value(z, 1e-6, 1.0)
    601         zenith_angles_prime = -zenith_spread*tf.math.log(z)/c_theta
    602 
    603         # Random component
    604         random_sign = config.tf_rng.uniform(shape=[batch_size, num_bs, 1,
    605             num_clusters_max], minval=0, maxval=2, dtype=tf.int32)
    606         random_sign = 2*random_sign - 1
    607         random_sign = tf.cast(random_sign, self._scenario.dtype.real_dtype)
    608         random_comp = config.tf_rng.normal(shape=[batch_size, num_bs, num_ut,
    609             num_clusters_max], mean=0.0, stddev=zenith_spread/7.0,
    610             dtype=self._scenario.dtype.real_dtype)
    611 
    612         # The center cluster angles depend on the UT scenario
    613         zenith_angles = random_sign*zenith_angles_prime + random_comp
    614         los_additinoal_comp = -(random_sign[:,:,:,:1]*
    615             zenith_angles_prime[:,:,:,:1] + random_comp[:,:,:,:1]
    616             - zenith_angles_los)
    617         if angle_type == 'zod':
    618             additional_comp = tf.where(los_uts, los_additinoal_comp,
    619                 zenith_angles_los + zod_offset)
    620         else:
    621             additional_comp = tf.where(los_uts, los_additinoal_comp,
    622                 0.0)
    623             additional_comp = tf.where(nlos_uts, zenith_angles_los,
    624                 additional_comp)
    625             additional_comp = tf.where(indoor_uts, tf.constant(90.0,
    626                 self._scenario.dtype.real_dtype),
    627                 additional_comp)
    628         zenith_angles = zenith_angles + additional_comp
    629 
    630         # Generating rays for every cluster
    631         # Add offset angles to cluster angles to get the ray angles
    632         ray_offsets = self._ray_offsets[:scenario.rays_per_cluster]
    633         # # Add dimensions for batch size, num bs, num ut, num clusters
    634         ray_offsets = tf.reshape(ray_offsets, [1,1,1,1,
    635                                                 scenario.rays_per_cluster])
    636         # Adding ray dimension for broadcasting
    637         zenith_angles = tf.expand_dims(zenith_angles, axis=4)
    638         cluster_angle_spread = tf.expand_dims(cluster_angle_spread, axis=4)
    639         zenith_angles = zenith_angles + cluster_angle_spread*ray_offsets
    640 
    641         # Wrapping to (0, 180)
    642         zenith_angles = wrap_angle_0_360(zenith_angles)
    643         zenith_angles = tf.where(tf.math.greater(zenith_angles, 180.),
    644             360.-zenith_angles, zenith_angles)
    645 
    646         return zenith_angles
    647 
    648     def _zenith_angles_of_arrival(self, zenith_spread_arrival, rician_k_factor,
    649         cluster_powers):
    650         # pylint: disable=line-too-long
    651         """
    652         Compute the zenith angle of arrival (ZoA)
    653         See step 7 of section 7.5 from TR 38.901 specification.
    654 
    655         Input
    656         ------
    657         zenith_spread_arrival : [batch size, num of BSs, num of UTs], tf.float
    658             Zenith angle spread of arrival (ZSA) [deg]
    659 
    660         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    661             Rician K-factor of each BS-UT link. Used only for LoS links.
    662 
    663         cluster_powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    664             Normalized path powers
    665 
    666         Output
    667         -------
    668         zoa : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    669             Paths zenith angles of arrival (ZoA) wrapped within (0,180) [degree]
    670         """
    671         return self._zenith_angles(zenith_spread_arrival, rician_k_factor,
    672                                    cluster_powers, 'zoa')
    673 
    674     def _zenith_angles_of_departure(self, zenith_spread_departure,
    675                                     rician_k_factor, cluster_powers):
    676         # pylint: disable=line-too-long
    677         """
    678         Compute the zenith angle of departure (ZoD)
    679         See step 7 of section 7.5 from TR 38.901 specification.
    680 
    681         Input
    682         ------
    683         zenith_spread_departure : [batch size, num of BSs, num of UTs], tf.float
    684             Zenith angle spread of departure (ZSD) [deg]
    685 
    686         rician_k_factor : [batch size, num of BSs, num of UTs], tf.float
    687             Rician K-factor of each BS-UT link. Used only for LoS links.
    688 
    689         cluster_powers : [batch size, num of BSs, num of UTs, maximum number of clusters], tf.float
    690             Normalized path powers
    691 
    692         Output
    693         -------
    694         zod : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    695             Paths zenith angles of departure (ZoD) wrapped within (0,180) [degree]
    696         """
    697         return self._zenith_angles(zenith_spread_departure, rician_k_factor,
    698                                    cluster_powers, 'zod')
    699 
    700     def _shuffle_angles(self, angles):
    701         # pylint: disable=line-too-long
    702         """
    703         Randomly shuffle a tensor carrying azimuth/zenith angles
    704         of arrival/departure.
    705 
    706         Input
    707         ------
    708         angles : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    709             Angles to shuffle
    710 
    711         Output
    712         -------
    713         shuffled_angles : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    714             Shuffled ``angles``
    715         """
    716 
    717         scenario = self._scenario
    718 
    719         batch_size = scenario.batch_size
    720         num_bs = scenario.num_bs
    721         num_ut = scenario.num_ut
    722 
    723         # Create randomly shuffled indices by arg-sorting samples from a random
    724         # normal distribution
    725         random_numbers = config.tf_rng.normal([batch_size, num_bs, 1,
    726                 scenario.num_clusters_max, scenario.rays_per_cluster])
    727         shuffled_indices = tf.argsort(random_numbers)
    728         shuffled_indices = tf.tile(shuffled_indices, [1, 1, num_ut, 1, 1])
    729         # Shuffling the angles
    730         shuffled_angles = tf.gather(angles,shuffled_indices, batch_dims=4)
    731         return shuffled_angles
    732 
    733     def _random_coupling(self, aoa, aod, zoa, zod):
    734         # pylint: disable=line-too-long
    735         """
    736         Randomly couples the angles within a cluster for both azimuth and
    737         elevation.
    738 
    739         Step 8 in TR 38.901 specification.
    740 
    741         Input
    742         ------
    743         aoa : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    744             Paths azimuth angles of arrival [degree] (AoA)
    745 
    746         aod : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    747             Paths azimuth angles of departure (AoD) [degree]
    748 
    749         zoa : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    750             Paths zenith angles of arrival [degree] (ZoA)
    751 
    752         zod : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    753             Paths zenith angles of departure [degree] (ZoD)
    754 
    755         Output
    756         -------
    757         shuffled_aoa : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    758             Shuffled `aoa`
    759 
    760         shuffled_aod : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    761             Shuffled `aod`
    762 
    763         shuffled_zoa : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    764             Shuffled `zoa`
    765 
    766         shuffled_zod : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    767             Shuffled `zod`
    768         """
    769         shuffled_aoa = self._shuffle_angles(aoa)
    770         shuffled_aod = self._shuffle_angles(aod)
    771         shuffled_zoa = self._shuffle_angles(zoa)
    772         shuffled_zod = self._shuffle_angles(zod)
    773 
    774         return shuffled_aoa, shuffled_aod, shuffled_zoa, shuffled_zod
    775 
    776     def _cross_polarization_power_ratios(self):
    777         # pylint: disable=line-too-long
    778         """
    779         Generate cross-polarization power ratios.
    780 
    781         Step 9 in TR 38.901 specification.
    782 
    783         Input
    784         ------
    785         None
    786 
    787         Output
    788         -------
    789         cross_polarization_power_ratios : [batch size, num of BSs, num of UTs, maximum number of clusters, number of rays], tf.float
    790             Polarization power ratios
    791         """
    792 
    793         scenario = self._scenario
    794 
    795         batch_size = scenario.batch_size
    796         num_bs = scenario.num_bs
    797         num_ut = scenario.num_ut
    798         num_clusters = scenario.num_clusters_max
    799         num_rays_per_cluster = scenario.rays_per_cluster
    800 
    801         # Loading XPR mean and standard deviation
    802         mu_xpr = scenario.get_param("muXPR")
    803         std_xpr = scenario.get_param("sigmaXPR")
    804         # Expanding for broadcasting with clusters and rays dims
    805         mu_xpr = tf.expand_dims(tf.expand_dims(mu_xpr, axis=3), axis=4)
    806         std_xpr = tf.expand_dims(tf.expand_dims(std_xpr, axis=3), axis=4)
    807 
    808         # XPR are assumed to follow a log-normal distribution.
    809         # Generate XPR in log-domain
    810         x = config.tf_rng.normal(shape=[batch_size, num_bs, num_ut, num_clusters,
    811             num_rays_per_cluster], mean=mu_xpr, stddev=std_xpr,
    812             dtype=self._scenario.dtype.real_dtype)
    813         # To linear domain
    814         cross_polarization_power_ratios = tf.math.pow(tf.constant(10.,
    815             self._scenario.dtype.real_dtype), x/10.0)
    816         return cross_polarization_power_ratios