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

utils.py (58411B)


      1 #
      2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
      3 # SPDX-License-Identifier: Apache-2.0
      4 #
      5 """Utility functions for the channel module"""
      6 
      7 import tensorflow as tf
      8 import warnings
      9 
     10 from sionna import PI
     11 from sionna import config
     12 from sionna.utils import expand_to_rank
     13 
     14 
     15 def subcarrier_frequencies(num_subcarriers, subcarrier_spacing,
     16                            dtype=tf.complex64):
     17     # pylint: disable=line-too-long
     18     r"""
     19     Compute the baseband frequencies of ``num_subcarrier`` subcarriers spaced by
     20     ``subcarrier_spacing``, i.e.,
     21 
     22     >>> # If num_subcarrier is even:
     23     >>> frequencies = [-num_subcarrier/2, ..., 0, ..., num_subcarrier/2-1] * subcarrier_spacing
     24     >>>
     25     >>> # If num_subcarrier is odd:
     26     >>> frequencies = [-(num_subcarrier-1)/2, ..., 0, ..., (num_subcarrier-1)/2] * subcarrier_spacing
     27 
     28 
     29     Input
     30     ------
     31     num_subcarriers : int
     32         Number of subcarriers
     33 
     34     subcarrier_spacing : float
     35         Subcarrier spacing [Hz]
     36 
     37     dtype : tf.DType
     38         Datatype to use for internal processing and output.
     39         If a complex datatype is provided, the corresponding precision of
     40         real components is used.
     41         Defaults to `tf.complex64` (`tf.float32`).
     42 
     43     Output
     44     ------
     45         frequencies : [``num_subcarrier``], tf.float
     46             Baseband frequencies of subcarriers
     47     """
     48 
     49     if dtype.is_complex:
     50         real_dtype = dtype.real_dtype
     51     elif dtype.if_floating:
     52         real_dtype = dtype
     53     else:
     54         raise AssertionError("dtype must be a complex or floating datatype")
     55 
     56     if tf.equal(tf.math.floormod(num_subcarriers, 2), 0):
     57         start=-num_subcarriers/2
     58         limit=num_subcarriers/2
     59     else:
     60         start=-(num_subcarriers-1)/2
     61         limit=(num_subcarriers-1)/2+1
     62 
     63     frequencies = tf.range( start=start,
     64                             limit=limit,
     65                             dtype=real_dtype)
     66     frequencies = frequencies*subcarrier_spacing
     67     return frequencies
     68 
     69 def time_frequency_vector(num_samples, sample_duration, dtype=tf.float32):
     70     # pylint: disable=line-too-long
     71     r"""
     72     Compute the time and frequency vector for a given number of samples
     73     and duration per sample in normalized time unit.
     74 
     75     >>> t = tf.cast(tf.linspace(-n_min, n_max, num_samples), dtype) * sample_duration
     76     >>> f = tf.cast(tf.linspace(-n_min, n_max, num_samples), dtype) * 1/(sample_duration*num_samples)
     77 
     78     Input
     79     ------
     80         num_samples : int
     81             Number of samples
     82 
     83         sample_duration : float
     84             Sample duration in normalized time
     85 
     86         dtype : tf.DType
     87             Datatype to use for internal processing and output.
     88             Defaults to `tf.float32`.
     89 
     90     Output
     91     ------
     92         t : [``num_samples``], ``dtype``
     93             Time vector
     94 
     95         f : [``num_samples``], ``dtype``
     96             Frequency vector
     97     """
     98 
     99     num_samples = int(num_samples)
    100 
    101     if tf.math.mod(num_samples, 2) == 0:  # if even
    102         n_min = tf.cast(-(num_samples) / 2, dtype=tf.int32)
    103         n_max = tf.cast((num_samples) / 2 - 1, dtype=tf.int32)
    104     else:  # if odd
    105         n_min = tf.cast(-(num_samples-1) / 2, dtype=tf.int32)
    106         n_max = tf.cast((num_samples+1) / 2 - 1, dtype=tf.int32)
    107 
    108     # Time vector
    109     t = tf.cast(tf.linspace(n_min, n_max, num_samples), dtype) \
    110         * tf.cast(sample_duration, dtype)
    111 
    112     # Frequency vector
    113     df = tf.cast(1.0/sample_duration, dtype)/tf.cast(num_samples, dtype)
    114     f = tf.cast(tf.linspace(n_min, n_max, num_samples), dtype) \
    115         * tf.cast(df, dtype)
    116 
    117     return t, f
    118 
    119 def time_lag_discrete_time_channel(bandwidth, maximum_delay_spread=3e-6):
    120     # pylint: disable=line-too-long
    121     r"""
    122     Compute the smallest and largest time-lag for the descrete complex baseband
    123     channel, i.e., :math:`L_{\text{min}}` and :math:`L_{\text{max}}`.
    124 
    125     The smallest time-lag (:math:`L_{\text{min}}`) returned is always -6, as this value
    126     was found small enough for all models included in Sionna.
    127 
    128     The largest time-lag (:math:`L_{\text{max}}`) is computed from the ``bandwidth``
    129     and ``maximum_delay_spread`` as follows:
    130 
    131     .. math::
    132         L_{\text{max}} = \lceil W \tau_{\text{max}} \rceil + 6
    133 
    134     where :math:`L_{\text{max}}` is the largest time-lag, :math:`W` the ``bandwidth``,
    135     and :math:`\tau_{\text{max}}` the ``maximum_delay_spread``.
    136 
    137     The default value for the ``maximum_delay_spread`` is 3us, which was found
    138     to be large enough to include most significant paths with all channel models
    139     included in Sionna assuming a nominal delay spread of 100ns.
    140 
    141     Note
    142     ----
    143     The values of :math:`L_{\text{min}}` and :math:`L_{\text{max}}` computed
    144     by this function are only recommended values.
    145     :math:`L_{\text{min}}` and :math:`L_{\text{max}}` should be set according to
    146     the considered channel model. For OFDM systems, one also needs to be careful
    147     that the effective length of the complex baseband channel is not larger than
    148     the cyclic prefix length.
    149 
    150     Input
    151     ------
    152     bandwidth : float
    153         Bandwith (:math:`W`) [Hz]
    154 
    155     maximum_delay_spread : float
    156         Maximum delay spread [s]. Defaults to 3us.
    157 
    158     Output
    159     -------
    160     l_min : int
    161         Smallest time-lag (:math:`L_{\text{min}}`) for the descrete complex baseband
    162         channel. Set to -6, , as this value was found small enough for all models
    163         included in Sionna.
    164 
    165     l_max : int
    166         Largest time-lag (:math:`L_{\text{max}}`) for the descrete complex baseband
    167         channel
    168     """
    169     l_min = tf.cast(-6, tf.int32)
    170     l_max = tf.math.ceil(maximum_delay_spread*bandwidth) + 6
    171     l_max = tf.cast(l_max, tf.int32)
    172     return l_min, l_max
    173 
    174 def cir_to_ofdm_channel(frequencies, a, tau, normalize=False):
    175     # pylint: disable=line-too-long
    176     r"""
    177     Compute the frequency response of the channel at ``frequencies``.
    178 
    179     Given a channel impulse response
    180     :math:`(a_{m}, \tau_{m}), 0 \leq m \leq M-1` (inputs ``a`` and ``tau``),
    181     the channel frequency response for the frequency :math:`f`
    182     is computed as follows:
    183 
    184     .. math::
    185         \widehat{h}(f) = \sum_{m=0}^{M-1} a_{m} e^{-j2\pi f \tau_{m}}
    186 
    187     Input
    188     ------
    189     frequencies : [fft_size], tf.float
    190         Frequencies at which to compute the channel response
    191 
    192     a : [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_paths, num_time_steps], tf.complex
    193         Path coefficients
    194 
    195     tau : [batch size, num_rx, num_tx, num_paths] or [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_paths], tf.float
    196         Path delays
    197 
    198     normalize : bool
    199         If set to `True`, the channel is normalized over the resource grid
    200         to ensure unit average energy per resource element. Defaults to `False`.
    201 
    202     Output
    203     -------
    204     h_f : [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_time_steps, fft_size], tf.complex
    205         Channel frequency responses at ``frequencies``
    206     """
    207 
    208     real_dtype = tau.dtype
    209 
    210     if len(tau.shape) == 4:
    211         # Expand dims to broadcast with h. Add the following dimensions:
    212         #  - number of rx antennas (2)
    213         #  - number of tx antennas (4)
    214         tau = tf.expand_dims(tf.expand_dims(tau, axis=2), axis=4)
    215         # Broadcast is not supported yet by TF for such high rank tensors.
    216         # We therefore do part of it manually
    217         tau = tf.tile(tau, [1, 1, 1, 1, a.shape[4], 1])
    218 
    219     # Add a time samples dimension for broadcasting
    220     tau = tf.expand_dims(tau, axis=6)
    221 
    222     # Bring all tensors to broadcastable shapes
    223     tau = tf.expand_dims(tau, axis=-1)
    224     h = tf.expand_dims(a, axis=-1)
    225     frequencies = expand_to_rank(frequencies, tf.rank(tau), axis=0)
    226 
    227     ## Compute the Fourier transforms of all cluster taps
    228     # Exponential component
    229     e = tf.exp(tf.complex(tf.constant(0, real_dtype),
    230         -2*PI*frequencies*tau))
    231     h_f = h*e
    232     # Sum over all clusters to get the channel frequency responses
    233     h_f = tf.reduce_sum(h_f, axis=-3)
    234 
    235     if normalize:
    236         # Normalization is performed such that for each batch example and
    237         # link the energy per resource grid is one.
    238         # Average over TX antennas, RX antennas, OFDM symbols and
    239         # subcarriers.
    240         c = tf.reduce_mean( tf.square(tf.abs(h_f)), axis=(2,4,5,6),
    241                             keepdims=True)
    242         c = tf.complex(tf.sqrt(c), tf.constant(0., real_dtype))
    243         h_f = tf.math.divide_no_nan(h_f, c)
    244 
    245     return h_f
    246 
    247 def cir_to_time_channel(bandwidth, a, tau, l_min, l_max, normalize=False):
    248     # pylint: disable=line-too-long
    249     r"""
    250     Compute the channel taps forming the discrete complex-baseband
    251     representation of the channel from the channel impulse response
    252     (``a``, ``tau``).
    253 
    254     This function assumes that a sinc filter is used for pulse shaping and receive
    255     filtering. Therefore, given a channel impulse response
    256     :math:`(a_{m}(t), \tau_{m}), 0 \leq m \leq M-1`, the channel taps
    257     are computed as follows:
    258 
    259     .. math::
    260         \bar{h}_{b, \ell}
    261         = \sum_{m=0}^{M-1} a_{m}\left(\frac{b}{W}\right)
    262             \text{sinc}\left( \ell - W\tau_{m} \right)
    263 
    264     for :math:`\ell` ranging from ``l_min`` to ``l_max``, and where :math:`W` is
    265     the ``bandwidth``.
    266 
    267     Input
    268     ------
    269     bandwidth : float
    270         Bandwidth [Hz]
    271 
    272     a : [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_paths, num_time_steps], tf.complex
    273         Path coefficients
    274 
    275     tau : [batch size, num_rx, num_tx, num_paths] or [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_paths], tf.float
    276         Path delays [s]
    277 
    278     l_min : int
    279         Smallest time-lag for the discrete complex baseband channel (:math:`L_{\text{min}}`)
    280 
    281     l_max : int
    282         Largest time-lag for the discrete complex baseband channel (:math:`L_{\text{max}}`)
    283 
    284     normalize : bool
    285         If set to `True`, the channel is normalized over the block size
    286         to ensure unit average energy per time step. Defaults to `False`.
    287 
    288     Output
    289     -------
    290     hm :  [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_time_steps, l_max - l_min + 1], tf.complex
    291         Channel taps coefficients
    292     """
    293 
    294     real_dtype = tau.dtype
    295 
    296     if len(tau.shape) == 4:
    297         # Expand dims to broadcast with h. Add the following dimensions:
    298         #  - number of rx antennas (2)
    299         #  - number of tx antennas (4)
    300         tau = tf.expand_dims(tf.expand_dims(tau, axis=2), axis=4)
    301         # Broadcast is not supported by TF for such high rank tensors.
    302         # We therefore do part of it manually
    303         tau = tf.tile(tau, [1, 1, 1, 1, a.shape[4], 1])
    304 
    305     # Add a time samples dimension for broadcasting
    306     tau = tf.expand_dims(tau, axis=6)
    307 
    308     # Time lags for which to compute the channel taps
    309     l = tf.range(l_min, l_max+1, dtype=real_dtype)
    310 
    311     # Bring tau and l to broadcastable shapes
    312     tau = tf.expand_dims(tau, axis=-1)
    313     l = expand_to_rank(l, tau.shape.rank, axis=0)
    314 
    315     # sinc pulse shaping
    316     g = tf.experimental.numpy.sinc(l-tau*bandwidth)
    317     g = tf.complex(g, tf.constant(0., real_dtype))
    318     a = tf.expand_dims(a, axis=-1)
    319 
    320     # For every tap, sum the sinc-weighted coefficients
    321     hm = tf.reduce_sum(a*g, axis=-3)
    322 
    323     if normalize:
    324         # Normalization is performed such that for each batch example and
    325         # link the energy per block is one.
    326         # The total energy of a channel response is the sum of the squared
    327         # norm over the channel taps.
    328         # Average over block size, RX antennas, and TX antennas
    329         c = tf.reduce_mean(tf.reduce_sum(tf.square(tf.abs(hm)),
    330                                          axis=6, keepdims=True),
    331                            axis=(2,4,5), keepdims=True)
    332         c = tf.complex(tf.sqrt(c), tf.constant(0., real_dtype))
    333         hm = tf.math.divide_no_nan(hm, c)
    334 
    335     return hm
    336 
    337 def time_to_ofdm_channel(h_t, rg, l_min):
    338     # pylint: disable=line-too-long
    339     r"""
    340     Compute the channel frequency response from the discrete complex-baseband
    341     channel impulse response.
    342 
    343     Given a discrete complex-baseband channel impulse response
    344     :math:`\bar{h}_{b,\ell}`, for :math:`\ell` ranging from :math:`L_\text{min}\le 0`
    345     to :math:`L_\text{max}`, the discrete channel frequency response is computed as
    346 
    347     .. math::
    348 
    349         \hat{h}_{b,n} = \sum_{k=0}^{L_\text{max}} \bar{h}_{b,k} e^{-j \frac{2\pi kn}{N}} + \sum_{k=L_\text{min}}^{-1} \bar{h}_{b,k} e^{-j \frac{2\pi n(N+k)}{N}}, \quad n=0,\dots,N-1
    350 
    351     where :math:`N` is the FFT size and :math:`b` is the time step.
    352 
    353     This function only produces one channel frequency response per OFDM symbol, i.e.,
    354     only values of :math:`b` corresponding to the start of an OFDM symbol (after
    355     cyclic prefix removal) are considered.
    356 
    357     Input
    358     ------
    359     h_t : [...num_time_steps,l_max-l_min+1], tf.complex
    360         Tensor of discrete complex-baseband channel impulse responses
    361 
    362     resource_grid : :class:`~sionna.ofdm.ResourceGrid`
    363         Resource grid
    364 
    365     l_min : int
    366         Smallest time-lag for the discrete complex baseband
    367         channel impulse response (:math:`L_{\text{min}}`)
    368 
    369     Output
    370     ------
    371     h_f : [...,num_ofdm_symbols,fft_size], tf.complex
    372         Tensor of discrete complex-baseband channel frequency responses
    373 
    374     Note
    375     ----
    376     Note that the result of this function is generally different from the
    377     output of :meth:`~sionna.channel.utils.cir_to_ofdm_channel` because
    378     the discrete complex-baseband channel impulse response is truncated
    379     (see :meth:`~sionna.channel.utils.cir_to_time_channel`). This effect
    380     can be observed in the example below.
    381 
    382     Examples
    383     --------
    384     .. code-block:: Python
    385 
    386         # Setup resource grid and channel model
    387         sm = StreamManagement(np.array([[1]]), 1)
    388         rg = ResourceGrid(num_ofdm_symbols=1,
    389                           fft_size=1024,
    390                           subcarrier_spacing=15e3)
    391         tdl = TDL("A", 100e-9, 3.5e9)
    392 
    393         # Generate CIR
    394         cir = tdl(batch_size=1, num_time_steps=1, sampling_frequency=rg.bandwidth)
    395 
    396         # Generate OFDM channel from CIR
    397         frequencies = subcarrier_frequencies(rg.fft_size, rg.subcarrier_spacing)
    398         h_freq = tf.squeeze(cir_to_ofdm_channel(frequencies, *cir, normalize=True))
    399 
    400         # Generate time channel from CIR
    401         l_min, l_max = time_lag_discrete_time_channel(rg.bandwidth)
    402         h_time = cir_to_time_channel(rg.bandwidth, *cir, l_min=l_min, l_max=l_max, normalize=True)
    403 
    404         # Generate OFDM channel from time channel
    405         h_freq_hat = tf.squeeze(time_to_ofdm_channel(h_time, rg, l_min))
    406 
    407         # Visualize results
    408         plt.figure()
    409         plt.plot(np.real(h_freq), "-")
    410         plt.plot(np.real(h_freq_hat), "--")
    411         plt.plot(np.imag(h_freq), "-")
    412         plt.plot(np.imag(h_freq_hat), "--")
    413         plt.xlabel("Subcarrier index")
    414         plt.ylabel(r"Channel frequency response")
    415         plt.legend(["OFDM Channel (real)", "OFDM Channel from time (real)", "OFDM Channel (imag)", "OFDM Channel from time (imag)"])
    416 
    417     .. image:: ../figures/time_to_ofdm_channel.png
    418     """
    419     # Totla length of an OFDM symbol including cyclic prefix
    420     ofdm_length = rg.fft_size + rg.cyclic_prefix_length
    421 
    422     # Downsample the impulse respons to one sample per OFDM symbol
    423     h_t = h_t[...,rg.cyclic_prefix_length:rg.num_time_samples:ofdm_length, :]
    424 
    425     # Pad channel impulse response with zeros to the FFT size
    426     pad_dims = rg.fft_size - tf.shape(h_t)[-1]
    427     pad_shape = tf.concat([tf.shape(h_t)[:-1], [pad_dims]], axis=-1)
    428     h_t = tf.concat([h_t, tf.zeros(pad_shape, dtype=h_t.dtype)], axis=-1)
    429 
    430     # Circular shift of negative time lags so that the channel impulse reponse
    431     # starts with h_{b,0}
    432     h_t = tf.roll(h_t, l_min, axis=-1)
    433 
    434     # Compute FFT
    435     h_f = tf.signal.fft(h_t)
    436 
    437     # Move the zero subcarrier to the center of the spectrum
    438     h_f = tf.signal.fftshift(h_f, axes=-1)
    439 
    440     return h_f
    441 
    442 
    443 def deg_2_rad(x):
    444     r"""
    445     Convert degree to radian
    446 
    447     Input
    448     ------
    449         x : Tensor
    450             Angles in degree
    451 
    452     Output
    453     -------
    454         y : Tensor
    455             Angles ``x`` converted to radian
    456     """
    457     return x*tf.constant(PI/180.0, x.dtype)
    458 
    459 def rad_2_deg(x):
    460     r"""
    461     Convert radian to degree
    462 
    463     Input
    464     ------
    465         x : Tensor
    466             Angles in radian
    467 
    468     Output
    469     -------
    470         y : Tensor
    471             Angles ``x`` converted to degree
    472     """
    473     return x*tf.constant(180.0/PI, x.dtype)
    474 
    475 def wrap_angle_0_360(angle):
    476     r"""
    477     Wrap ``angle`` to (0,360)
    478 
    479     Input
    480     ------
    481         angle : Tensor
    482             Input to wrap
    483 
    484     Output
    485     -------
    486         y : Tensor
    487             ``angle`` wrapped to (0,360)
    488     """
    489     return tf.math.mod(angle, 360.)
    490 
    491 def sample_bernoulli(shape, p, dtype=tf.float32):
    492     r"""
    493     Sample a tensor with shape ``shape`` from a Bernoulli distribution with
    494     probability ``p``
    495 
    496     Input
    497     --------
    498     shape : Tensor shape
    499         Shape of the tensor to sample
    500 
    501     p : Broadcastable with ``shape``, tf.float
    502         Probability
    503 
    504     dtype : tf.DType
    505         Datatype to use for internal processing and output.
    506 
    507     Output
    508     --------
    509     : Tensor of shape ``shape``, bool
    510         Binary samples
    511     """
    512     z = config.tf_rng.uniform(shape=shape, minval=0.0, maxval=1.0, dtype=dtype)
    513     z = tf.math.less(z, p)
    514     return z
    515 
    516 def drop_uts_in_sector(batch_size,
    517                        num_ut,
    518                        min_bs_ut_dist,
    519                        isd,
    520                        bs_height=0.,
    521                        ut_height=0.,
    522                        dtype=tf.complex64):
    523     r"""
    524     Uniformly sample UT locations from a sector
    525 
    526     The sector from which UTs are sampled is shown in the following figure.
    527     The BS is assumed to be located at the origin (0,0) of the coordinate
    528     system.
    529 
    530     .. figure:: ../figures/drop_uts_in_sector.png
    531         :align: center
    532         :scale: 30%
    533 
    534     Input
    535     --------
    536     batch_size : int
    537         Batch size
    538 
    539     num_ut : int
    540         Number of UTs to sample per batch example
    541 
    542     min_bs_ut_dist : tf.float
    543         Minimum BS-UT distance [m]
    544 
    545     isd : tf.float
    546         Inter-site distance, i.e., the distance between two adjacent BSs [m]
    547 
    548     bs_height : tf.float
    549         BS height, i.e., distance between the BS and the X-Y plane [m].
    550         Defaults to 0.
    551 
    552     ut_height : tf.float
    553         UT height, i.e., distance between the BS and the X-Y plane [m].
    554         Defaults to 0.
    555 
    556     dtype : tf.DType
    557         Datatype to use for internal processing and output.
    558         If a complex datatype is provided, the corresponding precision of
    559         real components is used.
    560         Defaults to `tf.complex64` (`tf.float32`).
    561 
    562     Output
    563     ------
    564     ut_loc : [batch_size, num_ut, 2], tf.float
    565         UT locations in the X-Y plane
    566     """
    567 
    568     if dtype.is_complex:
    569         real_dtype = dtype.real_dtype
    570     elif dtype.if_floating:
    571         real_dtype = dtype
    572     else:
    573         raise AssertionError("dtype must be a complex or floating datatype")
    574 
    575     # Cast to real_dtype
    576     d_min = tf.cast(min_bs_ut_dist, real_dtype)
    577     bs_height = tf.cast(bs_height, real_dtype)
    578     ut_height = tf.cast(ut_height, real_dtype)
    579 
    580     # Force the minimum BS-UT distance >= their height difference
    581     d_min = tf.maximum(d_min, tf.abs(bs_height - ut_height))
    582 
    583     r = tf.cast(isd*0.5, real_dtype)
    584 
    585     # Minimum squared distance between BS and UT on the X-Y plane
    586     r_min2 = d_min**2 - (bs_height - ut_height)**2
    587 
    588     # Angles from (-pi/6, pi/6), covering half of the sector and denoted by
    589     # alpha_half, are randomly sampled for all UTs.
    590     # Then, the maximum distance of UTs from the BS, denoted by r_max,
    591     # is computed for each angle.
    592     # For each angles, UT positions are sampled such that their *squared* from
    593     # the BTS is uniformly sampled from the range (r_min**2, r_max**2)
    594     # Each UT is then randomly and uniformly pushed to a half of the sector
    595     # by adding either PI/6 or PI/2 to the angle alpha_half
    596 
    597     # Sample angles for half of the sector (which half will be decided randomly)
    598     alpha_half = config.tf_rng.uniform(shape=[batch_size, num_ut],
    599                                        minval=-PI/6.,
    600                                        maxval=PI/6.,
    601                                        dtype=real_dtype)
    602 
    603     # Maximum distance (computed on the X-Y plane) from BS to a point in
    604     # the sector, at each angle in alpha_half
    605     r_max = r / tf.math.cos(alpha_half)
    606 
    607     # To ensure the UT distribution to be uniformly distributed across the
    608     # sector, we sample positions such that their *squared* distance from
    609     # the BS is uniformly distributed within (r_min**2, r_max**2)
    610     distance2 = config.tf_rng.uniform(shape=[batch_size, num_ut],
    611                                       minval=r_min2,
    612                                       maxval=r_max**2,
    613                                       dtype=real_dtype)
    614     distance = tf.sqrt(distance2)
    615 
    616     # Randomly assign the UTs to one of the two half of the sector
    617     side = sample_bernoulli([batch_size, num_ut],
    618                             tf.cast(0.5, real_dtype),
    619                             real_dtype)
    620     side = tf.cast(side, real_dtype)
    621     side = 2.*side + 1.
    622     alpha = alpha_half + side*PI/6.
    623 
    624     # Compute UT locations in the X-Y plane
    625     ut_loc = tf.stack([distance*tf.math.cos(alpha),
    626                        distance*tf.math.sin(alpha)], axis=-1)
    627 
    628     return ut_loc
    629 
    630 def set_3gpp_scenario_parameters(   scenario,
    631                                     min_bs_ut_dist=None,
    632                                     isd=None,
    633                                     bs_height=None,
    634                                     min_ut_height=None,
    635                                     max_ut_height=None,
    636                                     indoor_probability = None,
    637                                     min_ut_velocity=None,
    638                                     max_ut_velocity=None,
    639                                     dtype=tf.complex64):
    640     r"""
    641     Set valid parameters for a specified 3GPP system level ``scenario``
    642     (RMa, UMi, or UMa).
    643 
    644     If a parameter is given, then it is returned. If it is set to `None`,
    645     then a parameter valid according to the chosen scenario is returned
    646     (see [TR38901]_).
    647 
    648     Input
    649     --------
    650     scenario : str
    651         System level model scenario. Must be one of "rma", "umi", or "uma".
    652 
    653     min_bs_ut_dist : None or tf.float
    654         Minimum BS-UT distance [m]
    655 
    656     isd : None or tf.float
    657         Inter-site distance [m]
    658 
    659     bs_height : None or tf.float
    660         BS elevation [m]
    661 
    662     min_ut_height : None or tf.float
    663         Minimum UT elevation [m]
    664 
    665     max_ut_height : None or tf.float
    666         Maximum UT elevation [m]
    667 
    668     indoor_probability : None or tf.float
    669         Probability of a UT to be indoor
    670 
    671     min_ut_velocity : None or tf.float
    672         Minimum UT velocity [m/s]
    673 
    674     max_ut_velocity : None or tf.float
    675         Maximim UT velocity [m/s]
    676 
    677     dtype : tf.DType
    678         Datatype to use for internal processing and output.
    679         If a complex datatype is provided, the corresponding precision of
    680         real components is used.
    681         Defaults to `tf.complex64` (`tf.float32`).
    682     Output
    683     --------
    684     min_bs_ut_dist : tf.float
    685         Minimum BS-UT distance [m]
    686 
    687     isd : tf.float
    688         Inter-site distance [m]
    689 
    690     bs_height : tf.float
    691         BS elevation [m]
    692 
    693     min_ut_height : tf.float
    694         Minimum UT elevation [m]
    695 
    696     max_ut_height : tf.float
    697         Maximum UT elevation [m]
    698 
    699     indoor_probability : tf.float
    700         Probability of a UT to be indoor
    701 
    702     min_ut_velocity : tf.float
    703         Minimum UT velocity [m/s]
    704 
    705     max_ut_velocity : tf.float
    706         Maximim UT velocity [m/s]
    707     """
    708 
    709     assert scenario in ('umi', 'uma', 'rma'),\
    710         "`scenario` must be one of 'umi', 'uma', 'rma'"
    711 
    712     if dtype.is_complex:
    713         real_dtype = dtype.real_dtype
    714     elif dtype.if_floating:
    715         real_dtype = dtype
    716     else:
    717         raise AssertionError("dtype must be a complex or floating datatype")
    718 
    719     # Default values for scenario parameters.
    720     # From TR38.901, sections 7.2 and 7.4.
    721     # All distances and heights are in meters
    722     # All velocities are in meters per second.
    723     default_scenario_par = {'umi' : {
    724                                 'min_bs_ut_dist' : tf.constant(10., real_dtype),
    725                                 'isd' : tf.constant(200., real_dtype),
    726                                 'bs_height' : tf.constant(10., real_dtype),
    727                                 'min_ut_height' : tf.constant(1.5, real_dtype),
    728                                 'max_ut_height' : tf.constant(1.5, real_dtype),
    729                                 'indoor_probability' : tf.constant(0.8,
    730                                                                     real_dtype),
    731                                 'min_ut_velocity' : tf.constant(0.0,
    732                                                                     real_dtype),
    733                                 'max_ut_velocity' :tf.constant(0.0, real_dtype)
    734                             },
    735                             'uma' : {
    736                                 'min_bs_ut_dist' : tf.constant(35., real_dtype),
    737                                 'isd' : tf.constant(500., real_dtype),
    738                                 'bs_height' : tf.constant(25., real_dtype),
    739                                 'min_ut_height' : tf.constant(1.5, real_dtype),
    740                                 'max_ut_height' : tf.constant(1.5, real_dtype),
    741                                 'indoor_probability' : tf.constant(0.8,
    742                                                                     real_dtype),
    743                                 'min_ut_velocity' : tf.constant(0.0,
    744                                                                     real_dtype),
    745                                 'max_ut_velocity' : tf.constant(0.0,
    746                                                                     real_dtype),
    747                             },
    748                             'rma' : {
    749                                 'min_bs_ut_dist' : tf.constant(35., real_dtype),
    750                                 'isd' : tf.constant(5000., real_dtype),
    751                                 'bs_height' : tf.constant(35., real_dtype),
    752                                 'min_ut_height' : tf.constant(1.5, real_dtype),
    753                                 'max_ut_height' : tf.constant(1.5, real_dtype),
    754                                 'indoor_probability' : tf.constant(0.5,
    755                                                                     real_dtype),
    756                                 'min_ut_velocity' : tf.constant(0.0,
    757                                                                     real_dtype),
    758                                 'max_ut_velocity' : tf.constant(0.0,
    759                                                                     real_dtype)
    760                             }
    761                         }
    762 
    763     # Setting the scenario parameters
    764     if min_bs_ut_dist is None:
    765         min_bs_ut_dist = default_scenario_par[scenario]['min_bs_ut_dist']
    766     if isd is None:
    767         isd = default_scenario_par[scenario]['isd']
    768     if bs_height is None:
    769         bs_height = default_scenario_par[scenario]['bs_height']
    770     if min_ut_height is None:
    771         min_ut_height = default_scenario_par[scenario]['min_ut_height']
    772     if max_ut_height is None:
    773         max_ut_height = default_scenario_par[scenario]['max_ut_height']
    774     if indoor_probability is None:
    775         indoor_probability =default_scenario_par[scenario]['indoor_probability']
    776     if min_ut_velocity is None:
    777         min_ut_velocity = default_scenario_par[scenario]['min_ut_velocity']
    778     if max_ut_velocity is None:
    779         max_ut_velocity = default_scenario_par[scenario]['max_ut_velocity']
    780 
    781     return min_bs_ut_dist, isd, bs_height, min_ut_height, max_ut_height,\
    782             indoor_probability, min_ut_velocity, max_ut_velocity
    783 
    784 def relocate_uts(ut_loc, sector_id, cell_loc):
    785     # pylint: disable=line-too-long
    786     r"""
    787     Relocate the UTs by rotating them into the sector with index ``sector_id``
    788     and transposing them to the cell centered on ``cell_loc``.
    789 
    790     ``sector_id`` gives the index of the sector to which the UTs are
    791     rotated to. The picture below shows how the three sectors of a cell are
    792     indexed.
    793 
    794     .. figure:: ../figures/panel_array_sector_id.png
    795         :align: center
    796         :scale: 30%
    797 
    798         Indexing of sectors
    799 
    800     If ``sector_id`` is a scalar, then all UTs are relocated to the same
    801     sector indexed by ``sector_id``.
    802     If ``sector_id`` is a tensor, it should be broadcastable with
    803     [``batch_size``, ``num_ut``], and give the sector in which each UT or
    804     batch example is relocated to.
    805 
    806     When calling the function, ``ut_loc`` gives the locations of the UTs to
    807     relocate, which are all assumed to be in sector with index 0, and in the
    808     cell centered on the origin (0,0).
    809 
    810     Input
    811     --------
    812     ut_loc : [batch_size, num_ut, 2], tf.float
    813         UTs locations in the X-Y plan
    814 
    815     sector_id : Tensor broadcastable with [batch_size, num_ut], int
    816         Indexes of the sector to which to relocate the UTs
    817 
    818     cell_loc : Tensor broadcastable with [batch_size, num_ut], tf.float
    819         Center of the cell to which to transpose the UTs
    820 
    821     Output
    822     ------
    823     ut_loc : [batch_size, num_ut, 2], tf.float
    824         Relocated UTs locations in the X-Y plan
    825     """
    826 
    827     # Expand the rank of sector_id such that is is broadcastable with
    828     # (batch size, num_ut)
    829     sector_id = tf.cast(sector_id, ut_loc.dtype)
    830     sector_id = expand_to_rank(sector_id, 2, 0)
    831 
    832     # Expant
    833     cell_loc = tf.cast(cell_loc, ut_loc.dtype)
    834     cell_loc = expand_to_rank(cell_loc, tf.rank(ut_loc), 0)
    835 
    836     # Rotation matrix tensor, broadcastable with [batch size, num uts, 2, 2]
    837     rotation_angle = sector_id*2.*PI/3.0
    838     rotation_matrix = tf.stack([tf.math.cos(rotation_angle),
    839                                 -tf.math.sin(rotation_angle),
    840                                 tf.math.sin(rotation_angle),
    841                                 tf.math.cos(rotation_angle)],
    842                                axis=-1)
    843     rotation_matrix = tf.reshape(rotation_matrix,
    844                                  tf.concat([tf.shape(rotation_angle),
    845                                             [2,2]], axis=-1))
    846     rotation_matrix = tf.cast(rotation_matrix, ut_loc.dtype)
    847 
    848     # Applying the rotation matrix
    849     ut_loc = tf.expand_dims(ut_loc, axis=-1)
    850     ut_loc_rotated = tf.squeeze(rotation_matrix@ut_loc, axis=-1)
    851 
    852     # Translate to the BS location
    853     ut_loc_rotated_translated = ut_loc_rotated + cell_loc
    854 
    855     return ut_loc_rotated_translated
    856 
    857 def generate_uts_topology(  batch_size,
    858                             num_ut,
    859                             drop_area,
    860                             cell_loc_xy,
    861                             min_bs_ut_dist,
    862                             isd,
    863                             min_ut_height,
    864                             max_ut_height,
    865                             indoor_probability,
    866                             min_ut_velocity,
    867                             max_ut_velocity,
    868                             dtype=tf.complex64):
    869     # pylint: disable=line-too-long
    870     r"""
    871     Sample UTs location from a sector or a cell
    872 
    873     Input
    874     --------
    875     batch_size : int
    876         Batch size
    877 
    878     num_ut : int
    879         Number of UTs to sample per batch example
    880 
    881     drop_area : str
    882         'sector' or 'cell'. If set to 'sector', UTs are sampled from the
    883         sector with index 0 in the figure below
    884 
    885         .. figure:: ../figures/panel_array_sector_id.png
    886             :align: center
    887             :scale: 30%
    888 
    889     Indexing of sectors
    890 
    891     cell_loc_xy : Tensor broadcastable with[batch_size, num_ut, 3], tf.float
    892         Center of the cell(s)
    893 
    894     min_bs_ut_dist : None or tf.float
    895         Minimum BS-UT distance [m]
    896 
    897     isd : None or tf.float
    898         Inter-site distance [m]
    899 
    900     min_ut_height : None or tf.float
    901         Minimum UT elevation [m]
    902 
    903     max_ut_height : None or tf.float
    904         Maximum UT elevation [m]
    905 
    906     indoor_probability : None or tf.float
    907         Probability of a UT to be indoor
    908 
    909     min_ut_velocity : None or tf.float
    910         Minimum UT velocity [m/s]
    911 
    912     max_ut_velocity : None or tf.float
    913         Maximum UT velocity [m/s]
    914 
    915     dtype : tf.DType
    916         Datatype to use for internal processing and output.
    917         If a complex datatype is provided, the corresponding precision of
    918         real components is used.
    919         Defaults to `tf.complex64` (`tf.float32`).
    920 
    921     Output
    922     ------
    923     ut_loc : [batch_size, num_ut, 3], tf.float
    924         UTs locations
    925 
    926     ut_orientations : [batch_size, num_ut, 3], tf.float
    927         UTs orientations [radian]
    928 
    929     ut_velocities : [batch_size, num_ut, 3], tf.float
    930         UTs velocities [m/s]
    931 
    932     in_state : [batch_size, num_ut], tf.float
    933         Indoor/outdoor state of UTs. `True` means indoor, `False` means
    934         outdoor.
    935     """
    936 
    937     assert drop_area in ('sector', 'cell'),\
    938         "Drop area must be either 'sector' or 'cell'"
    939 
    940     if dtype.is_complex:
    941         real_dtype = dtype.real_dtype
    942     elif dtype.if_floating:
    943         real_dtype = dtype
    944     else:
    945         raise AssertionError("dtype must be a complex or floating datatype")
    946 
    947     # Randomly generating the UT locations
    948     ut_loc_xy = drop_uts_in_sector(batch_size,
    949                                    num_ut,
    950                                    min_bs_ut_dist,
    951                                    isd,
    952                                    dtype=dtype)
    953     if drop_area == 'sector':
    954         sectors = tf.constant(0, tf.int32)
    955     else: # 'cell'
    956         sectors = config.tf_rng.uniform(shape=[batch_size, num_ut],
    957                                         minval=0,
    958                                         maxval=3,
    959                                         dtype=tf.int32)
    960     ut_loc_xy = relocate_uts(ut_loc_xy,
    961                              sectors,
    962                              cell_loc_xy)
    963 
    964     ut_loc_z = config.tf_rng.uniform(shape=[batch_size, num_ut, 1],
    965                                      minval=min_ut_height,
    966                                      maxval=max_ut_height,
    967                                      dtype=real_dtype)
    968     ut_loc = tf.concat([ut_loc_xy,
    969                         ut_loc_z], axis=-1)
    970 
    971     # Randomly generating the UT indoor/outdoor state
    972     in_state = sample_bernoulli([batch_size, num_ut], indoor_probability,
    973                                 real_dtype)
    974 
    975     # Randomly generate the UT velocities
    976     ut_vel_angle = config.tf_rng.uniform([batch_size, num_ut],
    977                                          minval=-PI,
    978                                          maxval=PI,
    979                                          dtype=real_dtype)
    980     ut_vel_norm = config.tf_rng.uniform([batch_size, num_ut],
    981                                         minval=min_ut_velocity,
    982                                         maxval=max_ut_velocity,
    983                                         dtype=real_dtype)
    984     ut_velocities = tf.stack([ut_vel_norm*tf.math.cos(ut_vel_angle),
    985                               ut_vel_norm*tf.math.sin(ut_vel_angle),
    986                               tf.zeros([batch_size, num_ut], real_dtype)],
    987                              axis=-1)
    988 
    989     # Randomly generate the UT orientations
    990     ut_bearing = config.tf_rng.uniform([batch_size, num_ut], minval=-0.5*PI,
    991                                 maxval=0.5*PI, dtype=real_dtype)
    992     ut_downtilt = config.tf_rng.uniform([batch_size, num_ut], minval=-0.5*PI,
    993                                  maxval=0.5*PI, dtype=real_dtype)
    994     ut_slant = config.tf_rng.uniform([batch_size, num_ut], minval=-0.5*PI,
    995                               maxval=0.5*PI, dtype=real_dtype)
    996     ut_orientations = tf.stack([ut_bearing, ut_downtilt, ut_slant], axis=-1)
    997 
    998     return ut_loc, ut_orientations, ut_velocities, in_state
    999 
   1000 def gen_single_sector_topology( batch_size,
   1001                                 num_ut,
   1002                                 scenario,
   1003                                 min_bs_ut_dist=None,
   1004                                 isd=None,
   1005                                 bs_height=None,
   1006                                 min_ut_height=None,
   1007                                 max_ut_height=None,
   1008                                 indoor_probability = None,
   1009                                 min_ut_velocity=None,
   1010                                 max_ut_velocity=None,
   1011                                 dtype=tf.complex64):
   1012     # pylint: disable=line-too-long
   1013     r"""
   1014     Generate a batch of topologies consisting of a single BS located at the
   1015     origin and ``num_ut`` UTs randomly and uniformly dropped in a cell sector.
   1016 
   1017     The following picture shows the sector from which UTs are sampled.
   1018 
   1019     .. figure:: ../figures/drop_uts_in_sector.png
   1020         :align: center
   1021         :scale: 30%
   1022 
   1023     UTs orientations are randomly and uniformly set, whereas the BS orientation
   1024     is set such that the it is oriented towards the center of the sector.
   1025 
   1026     The drop configuration can be controlled through the optional parameters.
   1027     Parameters set to `None` are set to valid values according to the chosen
   1028     ``scenario`` (see [TR38901]_).
   1029 
   1030     The returned batch of topologies can be used as-is with the
   1031     :meth:`set_topology` method of the system level models, i.e.
   1032     :class:`~sionna.channel.tr38901.UMi`, :class:`~sionna.channel.tr38901.UMa`,
   1033     and :class:`~sionna.channel.tr38901.RMa`.
   1034 
   1035     Example
   1036     --------
   1037     >>> # Create antenna arrays
   1038     >>> bs_array = PanelArray(num_rows_per_panel = 4,
   1039     ...                      num_cols_per_panel = 4,
   1040     ...                      polarization = 'dual',
   1041     ...                      polarization_type = 'VH',
   1042     ...                      antenna_pattern = '38.901',
   1043     ...                      carrier_frequency = 3.5e9)
   1044     >>>
   1045     >>> ut_array = PanelArray(num_rows_per_panel = 1,
   1046     ...                       num_cols_per_panel = 1,
   1047     ...                       polarization = 'single',
   1048     ...                       polarization_type = 'V',
   1049     ...                       antenna_pattern = 'omni',
   1050     ...                       carrier_frequency = 3.5e9)
   1051     >>> # Create channel model
   1052     >>> channel_model = UMi(carrier_frequency = 3.5e9,
   1053     ...                     o2i_model = 'low',
   1054     ...                     ut_array = ut_array,
   1055     ...                     bs_array = bs_array,
   1056     ...                     direction = 'uplink')
   1057     >>> # Generate the topology
   1058     >>> topology = gen_single_sector_topology(batch_size = 100,
   1059     ...                                       num_ut = 4,
   1060     ...                                       scenario = 'umi')
   1061     >>> # Set the topology
   1062     >>> ut_loc, bs_loc, ut_orientations, bs_orientations, ut_velocities, in_state = topology
   1063     >>> channel_model.set_topology(ut_loc,
   1064     ...                            bs_loc,
   1065     ...                            ut_orientations,
   1066     ...                            bs_orientations,
   1067     ...                            ut_velocities,
   1068     ...                            in_state)
   1069     >>> channel_model.show_topology()
   1070 
   1071     .. image:: ../figures/drop_uts_in_sector_topology.png
   1072 
   1073     Input
   1074     --------
   1075     batch_size : int
   1076         Batch size
   1077 
   1078     num_ut : int
   1079         Number of UTs to sample per batch example
   1080 
   1081     scenario : str
   1082         System leven model scenario. Must be one of "rma", "umi", or "uma".
   1083 
   1084     min_bs_ut_dist : None or tf.float
   1085         Minimum BS-UT distance [m]
   1086 
   1087     isd : None or tf.float
   1088         Inter-site distance [m]
   1089 
   1090     bs_height : None or tf.float
   1091         BS elevation [m]
   1092 
   1093     min_ut_height : None or tf.float
   1094         Minimum UT elevation [m]
   1095 
   1096     max_ut_height : None or tf.float
   1097         Maximum UT elevation [m]
   1098 
   1099     indoor_probability : None or tf.float
   1100         Probability of a UT to be indoor
   1101 
   1102     min_ut_velocity : None or tf.float
   1103         Minimum UT velocity [m/s]
   1104 
   1105     max_ut_velocity : None or tf.float
   1106         Maximim UT velocity [m/s]
   1107 
   1108     dtype : tf.DType
   1109         Datatype to use for internal processing and output.
   1110         If a complex datatype is provided, the corresponding precision of
   1111         real components is used.
   1112         Defaults to `tf.complex64` (`tf.float32`).
   1113 
   1114     Output
   1115     ------
   1116     ut_loc : [batch_size, num_ut, 3], tf.float
   1117         UTs locations
   1118 
   1119     bs_loc : [batch_size, 1, 3], tf.float
   1120         BS location. Set to (0,0,0) for all batch examples.
   1121 
   1122     ut_orientations : [batch_size, num_ut, 3], tf.float
   1123         UTs orientations [radian]
   1124 
   1125     bs_orientations : [batch_size, 1, 3], tf.float
   1126         BS orientations [radian]. Oriented towards the center of the sector.
   1127 
   1128     ut_velocities : [batch_size, num_ut, 3], tf.float
   1129         UTs velocities [m/s]
   1130 
   1131     in_state : [batch_size, num_ut], tf.float
   1132         Indoor/outdoor state of UTs. `True` means indoor, `False` means
   1133         outdoor.
   1134     """
   1135 
   1136     params = set_3gpp_scenario_parameters(  scenario,
   1137                                             min_bs_ut_dist,
   1138                                             isd,
   1139                                             bs_height,
   1140                                             min_ut_height,
   1141                                             max_ut_height,
   1142                                             indoor_probability,
   1143                                             min_ut_velocity,
   1144                                             max_ut_velocity,
   1145                                             dtype)
   1146     min_bs_ut_dist, isd, bs_height, min_ut_height, max_ut_height,\
   1147             indoor_probability, min_ut_velocity, max_ut_velocity = params
   1148 
   1149     real_dtype = dtype.real_dtype
   1150 
   1151     # Setting BS to (0,0,bs_height)
   1152     bs_loc = tf.stack([ tf.zeros([batch_size, 1], real_dtype),
   1153                         tf.zeros([batch_size, 1], real_dtype),
   1154                         tf.fill( [batch_size, 1], bs_height)], axis=-1)
   1155 
   1156     # Setting the BS orientation such that it is downtilted towards the center
   1157     # of the sector
   1158     sector_center = (min_bs_ut_dist + 0.5*isd)*0.5
   1159     bs_downtilt = 0.5*PI - tf.math.atan(sector_center/bs_height)
   1160     bs_yaw = tf.constant(PI/3.0, real_dtype)
   1161     bs_orientation = tf.stack([ tf.fill([batch_size, 1], bs_yaw),
   1162                                 tf.fill([batch_size, 1], bs_downtilt),
   1163                                 tf.zeros([batch_size, 1], real_dtype)], axis=-1)
   1164 
   1165     # Generating the UTs
   1166     ut_topology = generate_uts_topology(    batch_size,
   1167                                             num_ut,
   1168                                             'sector',
   1169                                             tf.zeros([2], real_dtype),
   1170                                             min_bs_ut_dist,
   1171                                             isd,
   1172                                             min_ut_height,
   1173                                             max_ut_height,
   1174                                             indoor_probability,
   1175                                             min_ut_velocity,
   1176                                             max_ut_velocity,
   1177                                             dtype)
   1178     ut_loc, ut_orientations, ut_velocities, in_state = ut_topology
   1179 
   1180     return ut_loc, bs_loc, ut_orientations, bs_orientation, ut_velocities,\
   1181             in_state
   1182 
   1183 def gen_single_sector_topology_interferers( batch_size,
   1184                                             num_ut,
   1185                                             num_interferer,
   1186                                             scenario,
   1187                                             min_bs_ut_dist=None,
   1188                                             isd=None,
   1189                                             bs_height=None,
   1190                                             min_ut_height=None,
   1191                                             max_ut_height=None,
   1192                                             indoor_probability = None,
   1193                                             min_ut_velocity=None,
   1194                                             max_ut_velocity=None,
   1195                                             dtype=tf.complex64):
   1196     # pylint: disable=line-too-long
   1197     r"""
   1198     Generate a batch of topologies consisting of a single BS located at the
   1199     origin, ``num_ut`` UTs randomly and uniformly dropped in a cell sector, and
   1200     ``num_interferer`` interfering UTs randomly dropped in the adjacent cells.
   1201 
   1202     The following picture shows how UTs are sampled
   1203 
   1204     .. figure:: ../figures/drop_uts_in_sector_interferers.png
   1205         :align: center
   1206         :scale: 30%
   1207 
   1208     UTs orientations are randomly and uniformly set, whereas the BS orientation
   1209     is set such that it is oriented towards the center of the sector it
   1210     serves.
   1211 
   1212     The drop configuration can be controlled through the optional parameters.
   1213     Parameters set to `None` are set to valid values according to the chosen
   1214     ``scenario`` (see [TR38901]_).
   1215 
   1216     The returned batch of topologies can be used as-is with the
   1217     :meth:`set_topology` method of the system level models, i.e.
   1218     :class:`~sionna.channel.tr38901.UMi`, :class:`~sionna.channel.tr38901.UMa`,
   1219     and :class:`~sionna.channel.tr38901.RMa`.
   1220 
   1221     In the returned ``ut_loc``, ``ut_orientations``, ``ut_velocities``, and
   1222     ``in_state`` tensors, the first ``num_ut`` items along the axis with index
   1223     1 correspond to the served UTs, whereas the remaining ``num_interferer``
   1224     items correspond to the interfering UTs.
   1225 
   1226     Example
   1227     --------
   1228     >>> # Create antenna arrays
   1229     >>> bs_array = PanelArray(num_rows_per_panel = 4,
   1230     ...                      num_cols_per_panel = 4,
   1231     ...                      polarization = 'dual',
   1232     ...                      polarization_type = 'VH',
   1233     ...                      antenna_pattern = '38.901',
   1234     ...                      carrier_frequency = 3.5e9)
   1235     >>>
   1236     >>> ut_array = PanelArray(num_rows_per_panel = 1,
   1237     ...                       num_cols_per_panel = 1,
   1238     ...                       polarization = 'single',
   1239     ...                       polarization_type = 'V',
   1240     ...                       antenna_pattern = 'omni',
   1241     ...                       carrier_frequency = 3.5e9)
   1242     >>> # Create channel model
   1243     >>> channel_model = UMi(carrier_frequency = 3.5e9,
   1244     ...                     o2i_model = 'low',
   1245     ...                     ut_array = ut_array,
   1246     ...                     bs_array = bs_array,
   1247     ...                     direction = 'uplink')
   1248     >>> # Generate the topology
   1249     >>> topology = gen_single_sector_topology_interferers(batch_size = 100,
   1250     ...                                                   num_ut = 4,
   1251     ...                                                   num_interferer = 4,
   1252     ...                                                   scenario = 'umi')
   1253     >>> # Set the topology
   1254     >>> ut_loc, bs_loc, ut_orientations, bs_orientations, ut_velocities, in_state = topology
   1255     >>> channel_model.set_topology(ut_loc,
   1256     ...                            bs_loc,
   1257     ...                            ut_orientations,
   1258     ...                            bs_orientations,
   1259     ...                            ut_velocities,
   1260     ...                            in_state)
   1261     >>> channel_model.show_topology()
   1262 
   1263     .. image:: ../figures/drop_uts_in_sector_topology_inter.png
   1264 
   1265     Input
   1266     --------
   1267     batch_size : int
   1268         Batch size
   1269 
   1270     num_ut : int
   1271         Number of UTs to sample per batch example
   1272 
   1273     num_interferer : int
   1274         Number of interfeering UTs per batch example
   1275 
   1276     scenario : str
   1277         System leven model scenario. Must be one of "rma", "umi", or "uma".
   1278 
   1279     min_bs_ut_dist : None or tf.float
   1280         Minimum BS-UT distance [m]
   1281 
   1282     isd : None or tf.float
   1283         Inter-site distance [m]
   1284 
   1285     bs_height : None or tf.float
   1286         BS elevation [m]
   1287 
   1288     min_ut_height : None or tf.float
   1289         Minimum UT elevation [m]
   1290 
   1291     max_ut_height : None or tf.float
   1292         Maximum UT elevation [m]
   1293 
   1294     indoor_probability : None or tf.float
   1295         Probability of a UT to be indoor
   1296 
   1297     min_ut_velocity : None or tf.float
   1298         Minimum UT velocity [m/s]
   1299 
   1300     max_ut_velocity : None or tf.float
   1301         Maximim UT velocity [m/s]
   1302 
   1303     dtype : tf.DType
   1304         Datatype to use for internal processing and output.
   1305         If a complex datatype is provided, the corresponding precision of
   1306         real components is used.
   1307         Defaults to `tf.complex64` (`tf.float32`).
   1308 
   1309     Output
   1310     ------
   1311     ut_loc : [batch_size, num_ut, 3], tf.float
   1312         UTs locations. The first ``num_ut`` items along the axis with index
   1313         1 correspond to the served UTs, whereas the remaining
   1314         ``num_interferer`` items correspond to the interfeering UTs.
   1315 
   1316     bs_loc : [batch_size, 1, 3], tf.float
   1317         BS location. Set to (0,0,0) for all batch examples.
   1318 
   1319     ut_orientations : [batch_size, num_ut, 3], tf.float
   1320         UTs orientations [radian]. The first ``num_ut`` items along the
   1321         axis with index 1 correspond to the served UTs, whereas the
   1322         remaining ``num_interferer`` items correspond to the interfeering
   1323         UTs.
   1324 
   1325     bs_orientations : [batch_size, 1, 3], tf.float
   1326         BS orientation [radian]. Oriented towards the center of the sector.
   1327 
   1328     ut_velocities : [batch_size, num_ut, 3], tf.float
   1329         UTs velocities [m/s]. The first ``num_ut`` items along the axis
   1330         with index 1 correspond to the served UTs, whereas the remaining
   1331         ``num_interferer`` items correspond to the interfeering UTs.
   1332 
   1333     in_state : [batch_size, num_ut], tf.float
   1334         Indoor/outdoor state of UTs. `True` means indoor, `False` means
   1335         outdoor. The first ``num_ut`` items along the axis with
   1336         index 1 correspond to the served UTs, whereas the remaining
   1337         ``num_interferer`` items correspond to the interfering UTs.
   1338     """
   1339 
   1340     params = set_3gpp_scenario_parameters(  scenario,
   1341                                             min_bs_ut_dist,
   1342                                             isd,
   1343                                             bs_height,
   1344                                             min_ut_height,
   1345                                             max_ut_height,
   1346                                             indoor_probability,
   1347                                             min_ut_velocity,
   1348                                             max_ut_velocity,
   1349                                             dtype)
   1350     min_bs_ut_dist, isd, bs_height, min_ut_height, max_ut_height,\
   1351             indoor_probability, min_ut_velocity, max_ut_velocity = params
   1352 
   1353     real_dtype = dtype.real_dtype
   1354 
   1355     # Setting BS to (0,0,bs_height)
   1356     bs_loc = tf.stack([ tf.zeros([batch_size, 1], real_dtype),
   1357                         tf.zeros([batch_size, 1], real_dtype),
   1358                         tf.fill( [batch_size, 1], bs_height)], axis=-1)
   1359 
   1360     # Setting the BS orientation such that it is downtilted towards the center
   1361     # of the sector
   1362     sector_center = (min_bs_ut_dist + 0.5*isd)*0.5
   1363     bs_downtilt = 0.5*PI - tf.math.atan(sector_center/bs_height)
   1364     bs_yaw = tf.constant(PI/3.0, real_dtype)
   1365     bs_orientation = tf.stack([ tf.fill([batch_size, 1], bs_yaw),
   1366                                 tf.fill([batch_size, 1], bs_downtilt),
   1367                                 tf.zeros([batch_size, 1], real_dtype)], axis=-1)
   1368 
   1369     # Generating the UTs located in the UTs served by the BS
   1370     ut_topology = generate_uts_topology(    batch_size,
   1371                                             num_ut,
   1372                                             'sector',
   1373                                             tf.zeros([2], real_dtype),
   1374                                             min_bs_ut_dist,
   1375                                             isd,
   1376                                             min_ut_height,
   1377                                             max_ut_height,
   1378                                             indoor_probability,
   1379                                             min_ut_velocity,
   1380                                             max_ut_velocity,
   1381                                             dtype)
   1382     ut_loc, ut_orientations, ut_velocities, in_state = ut_topology
   1383 
   1384 
   1385     ## Generating the UTs located in the adjacent cells
   1386 
   1387     # Users are randomly dropped in one of the two adjacent cells
   1388     inter_cell_center = tf.stack([[0.0, isd],
   1389                                   [isd*tf.math.cos(PI/6.0),
   1390                                    isd*tf.math.sin(PI/6.0)]],
   1391                                  axis=0)
   1392     cell_index = config.tf_rng.uniform(shape=[batch_size, num_interferer],
   1393                                 minval=0, maxval=2, dtype=tf.int32)
   1394     inter_cells = tf.gather(inter_cell_center, cell_index)
   1395 
   1396     inter_topology = generate_uts_topology(     batch_size,
   1397                                                 num_interferer,
   1398                                                 'cell',
   1399                                                 inter_cells,
   1400                                                 min_bs_ut_dist,
   1401                                                 isd,
   1402                                                 min_ut_height,
   1403                                                 max_ut_height,
   1404                                                 indoor_probability,
   1405                                                 min_ut_velocity,
   1406                                                 max_ut_velocity,
   1407                                                 dtype)
   1408     inter_loc, inter_orientations, inter_velocities, inter_in_state \
   1409         = inter_topology
   1410 
   1411     ut_loc = tf.concat([ut_loc, inter_loc], axis=1)
   1412     ut_orientations = tf.concat([ut_orientations, inter_orientations], axis=1)
   1413     ut_velocities = tf.concat([ut_velocities, inter_velocities], axis=1)
   1414     in_state = tf.concat([in_state, inter_in_state], axis=1)
   1415 
   1416     return ut_loc, bs_loc, ut_orientations, bs_orientation, ut_velocities,\
   1417             in_state
   1418 
   1419 
   1420 
   1421 def exp_corr_mat(a, n, dtype=tf.complex64):
   1422     r"""Generate exponential correlation matrices.
   1423 
   1424     This function computes for every element :math:`a` of a complex-valued
   1425     tensor :math:`\mathbf{a}` the corresponding :math:`n\times n` exponential
   1426     correlation matrix :math:`\mathbf{R}(a,n)`, defined as (Eq. 1, [MAL2018]_):
   1427 
   1428     .. math::
   1429         \mathbf{R}(a,n)_{i,j} = \begin{cases}
   1430                     1 & \text{if } i=j\\
   1431                     a^{i-j}  & \text{if } i>j\\
   1432                     (a^\star)^{j-i}  & \text{if } j<i, j=1,\dots,n\\
   1433                   \end{cases}
   1434 
   1435     where :math:`|a|<1` and :math:`\mathbf{R}\in\mathbb{C}^{n\times n}`.
   1436 
   1437     Input
   1438     -----
   1439     a : [n_0, ..., n_k], tf.complex
   1440         A tensor of arbitrary rank whose elements
   1441         have an absolute value smaller than one.
   1442 
   1443     n : int
   1444         Number of dimensions of the output correlation matrices.
   1445 
   1446     dtype : tf.complex64, tf.complex128
   1447         The dtype of the output.
   1448 
   1449     Output
   1450     ------
   1451     R : [n_0, ..., n_k, n, n], tf.complex
   1452         A tensor of the same dtype as the input tensor :math:`\mathbf{a}`.
   1453     """
   1454     # Cast to desired output dtype and expand last dimension for broadcasting
   1455     a = tf.cast(a, dtype=dtype)
   1456     a = tf.expand_dims(a, -1)
   1457 
   1458     # Check that a is valid
   1459     msg = "The absolute value of the elements of `a` must be smaller than one"
   1460     tf.debugging.assert_less(tf.abs(a), tf.cast(1, a.dtype.real_dtype), msg)
   1461 
   1462     # Vector of exponents, adapt dtype and dimensions for broadcasting
   1463     exp = tf.range(0, n)
   1464     exp = tf.cast(exp, dtype=dtype)
   1465     exp = expand_to_rank(exp, tf.rank(a), 0)
   1466 
   1467     # First column of R
   1468     col = tf.math.pow(a, exp)
   1469 
   1470     # For a=0, one needs to remove the resulting nans due to 0**0=nan
   1471     cond = tf.math.is_nan(tf.math.real(col))
   1472     col = tf.where(cond, tf.ones_like(col), col)
   1473 
   1474     # First row of R (equal to complex-conjugate of the first column)
   1475     row = tf.math.conj(col)
   1476 
   1477     # Create Toeplitz operator
   1478     operator = tf.linalg.LinearOperatorToeplitz(col, row)
   1479 
   1480     # Generate dense tensor from operator
   1481     r = operator.to_dense()
   1482 
   1483     return r
   1484 
   1485 def one_ring_corr_mat(phi_deg, num_ant, d_h=0.5, sigma_phi_deg=15,
   1486                       dtype=tf.complex64):
   1487     r"""Generate covariance matrices from the one-ring model.
   1488 
   1489     This function generates approximate covariance matrices for the
   1490     so-called `one-ring` model (Eq. 2.24) [BHS2017]_. A uniform
   1491     linear array (ULA) with uniform antenna spacing is assumed. The elements
   1492     of the covariance matrices are computed as:
   1493 
   1494     .. math::
   1495         \mathbf{R}_{\ell,m} =
   1496               \exp\left( j2\pi d_\text{H} (\ell -m)\sin(\varphi) \right)
   1497               \exp\left( -\frac{\sigma_\varphi^2}{2}
   1498               \left( 2\pi d_\text{H}(\ell -m)\cos(\varphi) \right)^2 \right)
   1499 
   1500     for :math:`\ell,m = 1,\dots, M`, where :math:`M` is the number of antennas,
   1501     :math:`\varphi` is the angle of arrival, :math:`d_\text{H}` is the antenna
   1502     spacing in multiples of the wavelength,
   1503     and :math:`\sigma^2_\varphi` is the angular standard deviation.
   1504 
   1505     Input
   1506     -----
   1507     phi_deg : [n_0, ..., n_k], tf.float
   1508         A tensor of arbitrary rank containing azimuth angles (deg) of arrival.
   1509 
   1510     num_ant : int
   1511         Number of antennas
   1512 
   1513     d_h : float
   1514         Antenna spacing in multiples of the wavelength. Defaults to 0.5.
   1515 
   1516     sigma_phi_deg : float
   1517         Angular standard deviation (deg). Defaults to 15 (deg). Values greater
   1518         than 15 should not be used as the approximation becomes invalid.
   1519 
   1520     dtype : tf.complex64, tf.complex128
   1521         The dtype of the output.
   1522 
   1523     Output
   1524     ------
   1525     R : [n_0, ..., n_k, num_ant, nun_ant], `dtype`
   1526         Tensor containing the covariance matrices of the desired dtype.
   1527     """
   1528 
   1529     if sigma_phi_deg>15:
   1530         warnings.warn("sigma_phi_deg should be smaller than 15.")
   1531 
   1532     # Convert all inputs to radians
   1533     phi_deg = tf.cast(phi_deg, dtype=dtype.real_dtype)
   1534     sigma_phi_deg = tf.cast(sigma_phi_deg, dtype=dtype.real_dtype)
   1535     phi = deg_2_rad(phi_deg)
   1536     sigma_phi = deg_2_rad(sigma_phi_deg)
   1537 
   1538     # Add dimensions for broadcasting
   1539     phi = tf.expand_dims(phi, -1)
   1540     sigma_phi = tf.expand_dims(sigma_phi, -1)
   1541 
   1542     # Compute first column
   1543     c = tf.constant(2*PI*d_h, dtype=dtype.real_dtype)
   1544     d = c*tf.range(0, num_ant, dtype=dtype.real_dtype)
   1545     d = expand_to_rank(d, tf.rank(phi), 0)
   1546 
   1547     a = tf.complex(tf.cast(0, dtype=dtype.real_dtype), d*tf.sin(phi))
   1548     exp_a = tf.exp(a) # First exponential term
   1549 
   1550     b = -tf.cast(0.5, dtype=dtype.real_dtype)*(sigma_phi*d*tf.cos(phi))**2
   1551     exp_b = tf.cast(tf.exp(b), dtype=dtype) # Second exponetial term
   1552 
   1553     col = exp_a*exp_b # First column
   1554 
   1555     # First row is just the complex conjugate of first column
   1556     row = tf.math.conj(col)
   1557 
   1558     # Create Toeplitz operator
   1559     operator = tf.linalg.LinearOperatorToeplitz(col, row)
   1560 
   1561     # Generate dense tensor from operator
   1562     r = operator.to_dense()
   1563 
   1564     return r