#
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
"""Class for sampling rays following 3GPP TR38.901 specifications given a
channel simulation scenario and LSPs."""
from typing import TYPE_CHECKING, Optional
import torch
from sionna.phy.object import Object
from sionna.phy.channel.utils import deg_2_rad, wrap_angle_0_360
from sionna.phy.utils import normal, rand, randint
from .spatial_consistency import (
spatial_consistency_correlation_matrix,
spatial_consistency_matrix_sqrt,
)
from .utils import update_topology_buffer
if TYPE_CHECKING:
from .blockage import BlockageModelA, BlockageModelB
from .lsp import LSP
__all__ = ["Rays", "RaysGenerator"]
[docs]
class Rays:
r"""Class for conveniently storing rays
:param delays: Paths delays [s], shape
[batch size, number of base stations, number of UTs, number of clusters]
:param powers: Normalized path powers, shape
[batch size, number of base stations, number of UTs, number of clusters]
:param aoa: Azimuth angles of arrival [radian], shape
[batch size, number of base stations, number of UTs, number of clusters, number of rays]
:param aod: Azimuth angles of departure [radian], shape
[batch size, number of base stations, number of UTs, number of clusters, number of rays]
:param zoa: Zenith angles of arrival [radian], shape
[batch size, number of base stations, number of UTs, number of clusters, number of rays]
:param zod: Zenith angles of departure [radian], shape
[batch size, number of base stations, number of UTs, number of clusters, number of rays]
:param xpr: Cross-polarization power ratios, shape
[batch size, number of base stations, number of UTs, number of clusters, number of rays]
:param phases: Optional initial random phases [radian] for the four
polarization combinations, shape
[batch size, number of base stations, number of UTs, number of clusters,
number of rays, 4]. If `None`, phases are generated by
:class:`~sionna.phy.channel.tr38901.ChannelCoefficientsGenerator`.
Defaults to `None`.
:param blockage_loss_db: Optional blockage attenuation [dB], shape
[batch size, number of base stations, number of UTs, number of clusters,
number of rays].
:param los_blockage_loss_db: Optional blockage attenuation [dB] for the
deterministic LOS component, shape
[batch size, number of base stations, number of UTs].
:param blockage_loss_applied_to_powers: If `True`, ``blockage_loss_db`` is
informational and has already been included in ``powers``.
:param cluster_sort_indices: Optional permutation mapping the generated
clusters to increasing delay order, shape
[batch size, number of base stations, number of UTs, number of clusters].
:param strongest_cluster_indices: Optional indices of the two strongest
diffuse clusters before blockage is applied, shape
[batch size, number of base stations, number of UTs, 2].
"""
def __init__(
self,
delays: torch.Tensor,
powers: torch.Tensor,
aoa: torch.Tensor,
aod: torch.Tensor,
zoa: torch.Tensor,
zod: torch.Tensor,
xpr: torch.Tensor,
phases: Optional[torch.Tensor] = None,
blockage_loss_db: Optional[torch.Tensor] = None,
los_blockage_loss_db: Optional[torch.Tensor] = None,
blockage_loss_applied_to_powers: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
strongest_cluster_indices: Optional[torch.Tensor] = None,
) -> None:
self.delays = delays
self.powers = powers
self.aoa = aoa
self.aod = aod
self.zoa = zoa
self.zod = zod
self.xpr = xpr
self.phases = phases
self.blockage_loss_db = blockage_loss_db
self.los_blockage_loss_db = los_blockage_loss_db
self.blockage_loss_applied_to_powers = blockage_loss_applied_to_powers
self.cluster_sort_indices = cluster_sort_indices
self.strongest_cluster_indices = strongest_cluster_indices
[docs]
class RaysGenerator(Object):
r"""Sample rays according to a given channel scenario and large scale
parameters (LSP).
This class implements steps 6 to 9 from the TR 38.901 specifications,
(section 7.5).
Note that a global scenario is set for the entire batches when instantiating
this class (UMa, UMi, RMa, InH, or InF). However, each UT-BS link can have
its
specific state (LoS, NLoS, or indoor).
The batch size is set by the ``scenario`` given as argument when
constructing the class.
:param scenario: Scenario used to generate LSPs
:param enable_spatial_consistency: If `True`, generate the cluster- and
ray-specific random variables of Steps 5 to 10 from spatially
consistent random fields according to Sections 7.6.3.1 and 7.6.3.4 of
:cite:p:`TR38901V1920`. This includes cluster delays,
cluster shadowing, cluster angle signs and offsets, random coupling,
XPR, and initial random phases. Fields are sampled for the current
topology snapshot and are not retained across topology updates. The
feature is disabled by default for backwards compatibility.
:param blockage_model: Optional blockage model applied after Step 9 and
before Step 10, as specified by Section 7.6.4 of
:cite:p:`TR38901V1920`.
:input lsp: :class:`~sionna.phy.channel.tr38901.LSP`.
LSPs samples.
:output rays: :class:`~sionna.phy.channel.tr38901.Rays`.
Rays samples.
.. rubric:: Examples
.. code-block:: python
# Assuming scenario is a SystemLevelScenario instance
rays_generator = RaysGenerator(scenario)
rays = rays_generator(lsp)
"""
def __init__(
self,
scenario,
enable_spatial_consistency: bool = False,
blockage_model: Optional["BlockageModelA | BlockageModelB"] = None,
) -> None:
super().__init__(precision=scenario.precision, device=scenario.device)
self._scenario = scenario
self._enable_spatial_consistency = bool(enable_spatial_consistency)
self._blockage_model = blockage_model
self._spatial_consistency_matrix_sqrt = None
self._spatial_sign_consistency_matrix_sqrt = None
# For AoA, AoD, ZoA, and ZoD, offset to add to cluster angles to get ray
# angles. This is hardcoded from table 7.5-3 for 3GPP 38.901
# specification.
# Register as buffer for CUDAGraph compatibility
self.register_buffer("_ray_offsets", torch.tensor(
[
0.0447, -0.0447,
0.1413, -0.1413,
0.2492, -0.2492,
0.3715, -0.3715,
0.5129, -0.5129,
0.6797, -0.6797,
0.8844, -0.8844,
1.1481, -1.1481,
1.5195, -1.5195,
2.1551, -2.1551,
],
dtype=self.dtype,
device=self.device,
))
self.register_buffer(
"_subcluster_1_indices",
torch.tensor(
[0, 1, 2, 3, 4, 5, 6, 7, 18, 19],
dtype=torch.int64,
device=self.device,
),
)
self.register_buffer(
"_subcluster_2_indices",
torch.tensor(
[8, 9, 10, 11, 16, 17],
dtype=torch.int64,
device=self.device,
),
)
self.register_buffer(
"_subcluster_3_indices",
torch.tensor(
[12, 13, 14, 15],
dtype=torch.int64,
device=self.device,
),
)
def __call__(self, lsp: "LSP") -> Rays:
"""Generate rays from LSPs."""
# Sample cluster delays
delays, delays_unscaled, cluster_sort_indices = self._cluster_delays(
lsp.ds, lsp.k_factor
)
# Sample cluster powers
powers, powers_for_angles_gen = self._cluster_powers(
lsp.ds,
lsp.k_factor,
delays_unscaled,
cluster_sort_indices,
)
# Sample AoA
aoa, cluster_aoa = self._azimuth_angles_of_arrival(
lsp.asa, lsp.k_factor, powers_for_angles_gen,
return_cluster_angles=True,
cluster_sort_indices=cluster_sort_indices,
)
# Sample AoD
aod = self._azimuth_angles_of_departure(
lsp.asd,
lsp.k_factor,
powers_for_angles_gen,
cluster_sort_indices=cluster_sort_indices,
)
# Sample ZoA
zoa, cluster_zoa = self._zenith_angles_of_arrival(
lsp.zsa, lsp.k_factor, powers_for_angles_gen,
return_cluster_angles=True,
cluster_sort_indices=cluster_sort_indices,
)
# Sample ZoD
zod = self._zenith_angles_of_departure(
lsp.zsd,
lsp.k_factor,
powers_for_angles_gen,
cluster_sort_indices=cluster_sort_indices,
)
# XPRs
xpr = self._cross_polarization_power_ratios(cluster_sort_indices)
# Random coupling
strongest_cluster_indices = torch.topk(
powers, k=min(2, powers.shape[-1]), dim=-1
).indices
aoa, aod, zoa, zod = self._random_coupling(
aoa,
aod,
zoa,
zod,
strongest_cluster_indices,
cluster_sort_indices,
)
blockage_loss_db = None
los_blockage_loss_db = None
blockage_loss_applied_to_powers = False
if self._blockage_model is not None:
if self._blockage_model.requires_ray_angles:
blockage_loss_db, los_blockage_loss_db = self._blockage_model(
aoa,
zoa,
self._scenario.los_aoa,
self._scenario.los_zoa,
)
else:
blockage_cluster_loss_db, los_blockage_loss_db = self._blockage_model(
cluster_aoa,
cluster_zoa,
self._scenario.los_aoa,
self._scenario.los_zoa,
)
powers = powers * torch.pow(
torch.tensor(10.0, dtype=self.dtype, device=self.device),
-blockage_cluster_loss_db / 10.0,
)
blockage_loss_db = blockage_cluster_loss_db.unsqueeze(-1).expand(
-1, -1, -1, -1, self._scenario.rays_per_cluster
)
blockage_loss_applied_to_powers = True
# Initial random phases for step 10. These are only generated here
# when spatial consistency is enabled. Otherwise, they are generated by
# ChannelCoefficientsGenerator to preserve the historical RNG sequence.
phases = None
if self._enable_spatial_consistency:
phases = self._random_initial_phases(cluster_sort_indices)
# Convert angles of arrival and departure from degree to radian
aoa = deg_2_rad(aoa)
aod = deg_2_rad(aod)
zoa = deg_2_rad(zoa)
zod = deg_2_rad(zod)
# Storing and returning rays
rays = Rays(
delays=delays,
powers=powers,
aoa=aoa,
aod=aod,
zoa=zoa,
zod=zod,
xpr=xpr,
phases=phases,
blockage_loss_db=blockage_loss_db,
los_blockage_loss_db=los_blockage_loss_db,
blockage_loss_applied_to_powers=blockage_loss_applied_to_powers,
cluster_sort_indices=cluster_sort_indices,
strongest_cluster_indices=strongest_cluster_indices,
)
return rays
[docs]
def topology_updated_callback(self) -> None:
"""Updates internal quantities when the scenario topology changes."""
self._compute_clusters_mask()
if self._enable_spatial_consistency:
self._compute_spatial_consistency_matrix_sqrt()
self._compute_spatial_sign_consistency_matrix_sqrt()
if self._blockage_model is not None:
self._blockage_model.topology_updated_callback()
@property
def spatial_consistency_enabled(self) -> bool:
"""`True` if spatial consistency is enabled for ray generation."""
return self._enable_spatial_consistency
[docs]
def reset_topology(self) -> None:
"""Reset topology-dependent buffers."""
if hasattr(self, "_cluster_mask"):
delattr(self, "_cluster_mask")
if hasattr(self, "_spatial_consistency_matrix_sqrt"):
delattr(self, "_spatial_consistency_matrix_sqrt")
self._spatial_consistency_matrix_sqrt = None
if hasattr(self, "_spatial_sign_consistency_matrix_sqrt"):
delattr(self, "_spatial_sign_consistency_matrix_sqrt")
self._spatial_sign_consistency_matrix_sqrt = None
if self._blockage_model is not None:
self._blockage_model.reset_topology()
[docs]
def allocate_topology_tensors(
self,
batch_size: int,
num_bs: int,
num_ut: int,
) -> None:
"""Pre-allocate topology-dependent buffers."""
self.reset_topology()
for name in (
"_cluster_mask",
"_spatial_consistency_matrix_sqrt",
"_spatial_sign_consistency_matrix_sqrt",
):
if hasattr(self, name) and name not in self._buffers:
delattr(self, name)
self.register_buffer(
"_cluster_mask",
torch.zeros(
batch_size,
num_bs,
num_ut,
self._scenario.num_clusters_max,
dtype=self.dtype,
device=self.device,
),
)
if self._enable_spatial_consistency:
self.register_buffer(
"_spatial_consistency_matrix_sqrt",
torch.zeros(
batch_size,
num_bs,
num_ut,
num_ut,
dtype=self.dtype,
device=self.device,
),
)
self.register_buffer(
"_spatial_sign_consistency_matrix_sqrt",
torch.zeros(
batch_size,
num_bs,
num_ut,
num_ut,
dtype=self.dtype,
device=self.device,
),
)
if (
self._blockage_model is not None
and hasattr(self._blockage_model, "allocate_topology_tensors")
):
self._blockage_model.allocate_topology_tensors(
batch_size,
num_bs,
num_ut,
)
########################################
# Internal utility methods
########################################
def _spatial_consistency_distances(self) -> tuple[float, float, float]:
"""Return correlation distances for cluster/ray random variables."""
kind = self._scenario.scenario_kind
if kind == "rma":
return 50.0, 60.0, 15.0
if kind == "umi":
return 12.0, 15.0, 15.0
if kind == "uma":
return 40.0, 50.0, 15.0
if kind == "inh":
return 10.0, 10.0, 10.0
if kind == "inf":
return 10.0, 10.0, 10.0
raise ValueError(
"Spatial consistency is not configured for scenario "
f"{type(self._scenario).__name__}"
)
def _spatial_consistency_states(self) -> torch.Tensor:
"""Link-state labels used to mask spatial consistency."""
scenario = self._scenario
indoor = scenario.indoor.unsqueeze(1).expand(
-1, scenario.num_bs, -1
)
los = scenario.los
states = torch.where(
los,
torch.zeros((), dtype=torch.int64, device=self.device),
torch.ones((), dtype=torch.int64, device=self.device),
)
if scenario.use_indoor_lsp_params:
states = torch.where(
indoor,
torch.full((), 2, dtype=torch.int64, device=self.device),
states,
)
region_ids = scenario.ut_spatial_region_ids.unsqueeze(1)
return states + 3 * region_ids
def _compute_spatial_matrix_sqrt(self, buffer_name: str) -> None:
"""Compute and store a spatial-consistency matrix square root."""
scenario = self._scenario
d_los, d_nlos, d_o2i = self._spatial_consistency_distances()
correlation_distance = scenario.broadcast_params(d_los, d_nlos, d_o2i)
states = self._spatial_consistency_states()
distance_2d = scenario.matrix_ut_distance_2d.unsqueeze(1)
correlation = spatial_consistency_correlation_matrix(
distance_2d,
correlation_distance,
states=states,
correlation_distance_layout="per_terminal",
precision=self.precision,
device=self.device,
)
chol = spatial_consistency_matrix_sqrt(
correlation,
precision=self.precision,
device=self.device,
)
self._update_buffer(buffer_name, chol)
def _compute_spatial_consistency_matrix_sqrt(self) -> None:
"""Compute matrices for continuous spatial random variables."""
self._compute_spatial_matrix_sqrt("_spatial_consistency_matrix_sqrt")
def _compute_spatial_sign_consistency_matrix_sqrt(self) -> None:
"""Compute matrices for track-fixed discrete random variables."""
self._compute_spatial_matrix_sqrt(
"_spatial_sign_consistency_matrix_sqrt"
)
def _update_buffer(self, name: str, value: torch.Tensor) -> None:
"""Update or register a topology-dependent buffer."""
update_topology_buffer(self, name, value)
def _apply_spatial_consistency(
self,
samples: torch.Tensor,
use_sign_matrix: bool = False,
) -> torch.Tensor:
"""Apply the UT-domain spatial filter to site-specific samples."""
if self._enable_spatial_consistency:
if use_sign_matrix:
if self._spatial_sign_consistency_matrix_sqrt is None:
self._compute_spatial_sign_consistency_matrix_sqrt()
matrix_sqrt = self._spatial_sign_consistency_matrix_sqrt
else:
if self._spatial_consistency_matrix_sqrt is None:
self._compute_spatial_consistency_matrix_sqrt()
matrix_sqrt = self._spatial_consistency_matrix_sqrt
batch_size, num_bs, num_ut = samples.shape[:3]
tail_shape = samples.shape[3:]
samples_flat = samples.reshape(batch_size, num_bs, num_ut, -1)
samples_flat = torch.matmul(matrix_sqrt, samples_flat)
samples = samples_flat.reshape(
batch_size, num_bs, num_ut, *tail_shape
)
track_ids = self._scenario.spatial_consistency_track_ids
if use_sign_matrix and track_ids is not None:
same_track = track_ids.unsqueeze(-1) == track_ids.unsqueeze(-2)
representative = torch.argmax(same_track.to(torch.int64), dim=-1)
gather_index = representative.reshape(
batch_size, 1, num_ut, *([1] * len(tail_shape))
)
gather_index = gather_index.expand(
batch_size, num_bs, num_ut, *tail_shape
)
samples = torch.gather(samples, dim=2, index=gather_index)
return self._scenario.share_by_bs_site(samples)
def _spatial_normal(
self,
shape: tuple[int, ...],
use_sign_matrix: bool = False,
) -> torch.Tensor:
"""Generate normal variables with optional spatial consistency."""
samples = normal(
shape,
dtype=self.dtype,
device=self.device,
generator=self.torch_rng,
)
return self._apply_spatial_consistency(samples, use_sign_matrix)
def _spatial_uniform(
self,
shape: tuple[int, ...],
use_sign_matrix: bool = False,
) -> torch.Tensor:
"""Generate uniform variables with optional spatial consistency."""
if not self._enable_spatial_consistency:
samples = rand(
shape,
dtype=self.dtype,
device=self.device,
generator=self.torch_rng,
)
return self._scenario.share_by_bs_site(samples)
samples = self._spatial_normal(shape, use_sign_matrix)
sqrt_two = torch.sqrt(
torch.tensor(2.0, dtype=self.dtype, device=self.device)
)
samples = 0.5 * torch.erfc(
-samples / sqrt_two
)
eps = torch.finfo(self.dtype).eps
return samples.clamp(eps, 1.0 - eps)
def _spatial_binary_sign(self, shape: tuple[int, ...]) -> torch.Tensor:
"""Generate signs in {-1, 1} with optional spatial consistency."""
if not self._enable_spatial_consistency:
samples = randint(
0,
2,
shape,
dtype=torch.int32,
device=self.device,
generator=self.torch_rng,
)
samples = (2 * samples - 1).to(self.dtype)
return self._scenario.share_by_bs_site(samples)
samples = torch.where(
self._spatial_uniform(shape, use_sign_matrix=True) < 0.5,
torch.tensor(-1.0, dtype=self.dtype, device=self.device),
torch.tensor(1.0, dtype=self.dtype, device=self.device),
)
return samples
def _sort_clusters_by_delay(
self,
samples: torch.Tensor,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Reorder cluster-indexed samples according to the delay ordering."""
if cluster_sort_indices is None:
return samples
if samples.shape[2] != self._scenario.num_ut:
return samples
indices = cluster_sort_indices
while indices.dim() < samples.dim():
indices = indices.unsqueeze(-1)
return torch.gather(samples, dim=3, index=indices.expand_as(samples))
def _random_initial_phases(
self,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Generate spatially consistent initial random phases for step 10."""
scenario = self._scenario
phases = self._spatial_uniform(
(
scenario.batch_size,
scenario.num_bs,
scenario.num_ut,
scenario.num_clusters_max,
scenario.rays_per_cluster,
4,
)
)
phases = self._sort_clusters_by_delay(phases, cluster_sort_indices)
return phases * 2.0 * torch.pi - torch.pi
def _compute_clusters_mask(self) -> None:
"""
Given a scenario (UMi, UMa, RMa, InH, InF), the number of clusters is different
for different state of UT-BS links (LoS, NLoS, indoor).
Because we use tensors with predefined dimension size (not ragged), the
cluster dimension is always set to the maximum number of clusters the
scenario requires. A mask is then used to discard not required tensors,
depending on the state of each UT-BS link.
This function computes and stores this mask of size
[batch size, number of base stations, number of UTs, maximum number of cluster]
where an element equals 0 if the cluster is used, 1 otherwise.
"""
scenario = self._scenario
num_clusters_los = scenario.num_clusters_los
num_clusters_nlos = scenario.num_clusters_nlos
num_clusters_o2i = scenario.num_clusters_indoor
num_clusters_max = max(num_clusters_los, num_clusters_nlos, num_clusters_o2i)
# Initialize an empty mask
mask = torch.zeros(
scenario.batch_size,
scenario.num_bs,
scenario.num_ut,
num_clusters_max,
dtype=self.dtype,
device=self.device,
)
# Indoor mask
mask_indoor = torch.cat(
[
torch.zeros(num_clusters_o2i, dtype=self.dtype, device=self.device),
torch.ones(
num_clusters_max - num_clusters_o2i,
dtype=self.dtype,
device=self.device,
),
],
dim=0,
).reshape(1, 1, 1, num_clusters_max)
indoor = scenario.indoor.unsqueeze(1) # Broadcasting with BS
if scenario.use_indoor_lsp_params:
o2i_slice_mask = indoor.to(self.dtype).unsqueeze(3)
else:
o2i_slice_mask = torch.zeros(
scenario.batch_size,
scenario.num_bs,
scenario.num_ut,
1,
dtype=self.dtype,
device=self.device,
)
mask = mask + o2i_slice_mask * mask_indoor
# LoS
mask_los = torch.cat(
[
torch.zeros(num_clusters_los, dtype=self.dtype, device=self.device),
torch.ones(
num_clusters_max - num_clusters_los,
dtype=self.dtype,
device=self.device,
),
],
dim=0,
).reshape(1, 1, 1, num_clusters_max)
los_slice_mask = scenario.los.to(self.dtype).unsqueeze(3)
mask = mask + los_slice_mask * mask_los
# NLoS
mask_nlos = torch.cat(
[
torch.zeros(num_clusters_nlos, dtype=self.dtype, device=self.device),
torch.ones(
num_clusters_max - num_clusters_nlos,
dtype=self.dtype,
device=self.device,
),
],
dim=0,
).reshape(1, 1, 1, num_clusters_max)
if scenario.use_indoor_lsp_params:
nlos_slice_mask = (~scenario.los) & (~indoor)
else:
nlos_slice_mask = ~scenario.los
nlos_slice_mask = nlos_slice_mask.to(self.dtype).unsqueeze(3)
mask = mask + nlos_slice_mask * mask_nlos
# Save the mask
self._update_buffer("_cluster_mask", mask)
def _cluster_delays(
self, delay_spread: torch.Tensor, rician_k_factor: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Generate cluster delays (step 5 of section 7.5, TR 38.901).
:param delay_spread: RMS delay spread of each BS-UT link,
shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
"""
scenario = self._scenario
batch_size = scenario.batch_size
num_bs = scenario.num_bs
num_ut = scenario.num_ut
num_clusters_max = scenario.num_clusters_max
# Getting scaling parameter according to each BS-UT link scenario
delay_scaling_parameter = scenario.get_param("rTau").unsqueeze(3)
# Generating random cluster delays
# We don't start at 0 to avoid numerical errors
delay_spread = delay_spread.unsqueeze(3)
x = self._spatial_uniform((batch_size, num_bs, num_ut, num_clusters_max))
x = x * (1.0 - 1e-6) + 1e-6
# Moving to linear domain
unscaled_delays = -delay_scaling_parameter * delay_spread * torch.log(x)
# Forcing the cluster that should not exist to huge delays (1s)
unscaled_delays = unscaled_delays * (1.0 - self._cluster_mask) + self._cluster_mask
# Normalizing and sorting the delays
unscaled_delays = unscaled_delays - unscaled_delays.min(dim=3, keepdim=True).values
cluster_sort_indices = torch.argsort(unscaled_delays, dim=3)
unscaled_delays = torch.gather(
unscaled_delays, dim=3, index=cluster_sort_indices
)
# Additional scaling applied to LoS links
rician_k_factor_db = 10.0 * torch.log10(rician_k_factor) # to dB
scaling_factor = (
0.7705
- 0.0433 * rician_k_factor_db
+ 0.0002 * rician_k_factor_db.square()
+ 0.000017 * rician_k_factor_db.pow(3.0)
).unsqueeze(3)
delays = torch.where(
scenario.los.unsqueeze(3),
unscaled_delays / scaling_factor,
unscaled_delays,
)
return delays, unscaled_delays, cluster_sort_indices
def _cluster_powers(
self,
delay_spread: torch.Tensor,
rician_k_factor: torch.Tensor,
unscaled_delays: torch.Tensor,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Generate cluster powers (step 6 of section 7.5, TR 38.901).
:param delay_spread: RMS delay spread of each BS-UT link,
shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param unscaled_delays: Unscaled path delays [s],
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]. Required to compute the path powers.
:param cluster_sort_indices: Permutation mapping raw clusters to
increasing delay order.
"""
scenario = self._scenario
batch_size = scenario.batch_size
num_bs = scenario.num_bs
num_ut = scenario.num_ut
num_clusters_max = scenario.num_clusters_max
delay_scaling_parameter = scenario.get_param("rTau")
cluster_shadowing_std_db = scenario.get_param("zeta")
delay_spread = delay_spread.unsqueeze(3)
cluster_shadowing_std_db = cluster_shadowing_std_db.unsqueeze(3)
delay_scaling_parameter = delay_scaling_parameter.unsqueeze(3)
# Generate unnormalized cluster powers
z = self._spatial_normal((batch_size, num_bs, num_ut, num_clusters_max))
z = self._sort_clusters_by_delay(z, cluster_sort_indices)
z = z * cluster_shadowing_std_db
# Moving to linear domain
powers_unnormalized = torch.exp(
-unscaled_delays
* (delay_scaling_parameter - 1.0)
/ (delay_scaling_parameter * delay_spread)
) * torch.pow(torch.tensor(10.0, dtype=self.dtype, device=self.device), -z / 10.0)
# Force the power of unused cluster to zero
powers_unnormalized = powers_unnormalized * (1.0 - self._cluster_mask)
# Normalizing cluster powers
powers = powers_unnormalized / powers_unnormalized.sum(dim=3, keepdim=True)
rician_k_factor = rician_k_factor.unsqueeze(3)
p_nlos_scaling = 1.0 / (rician_k_factor + 1.0)
p_1_los = rician_k_factor * p_nlos_scaling
# Remove clusters more than 25 dB below the strongest cluster while
# keeping fixed tensor shapes. TR 38.901 Step 6 defines the threshold
# based on (7.5-6), i.e., before adding the LoS specular component.
power_threshold = (
powers.max(dim=3, keepdim=True).values
* torch.pow(torch.tensor(10.0, dtype=self.dtype, device=self.device), -2.5)
)
cluster_keep_mask = (
(powers >= power_threshold)
& (self._cluster_mask == 0.0)
).to(self.dtype)
powers = powers * cluster_keep_mask
powers_1 = p_nlos_scaling * powers[:, :, :, :1] + p_1_los
powers_n = p_nlos_scaling * powers[:, :, :, 1:]
powers_for_angles_gen = torch.where(
scenario.los.unsqueeze(3),
torch.cat([powers_1, powers_n], dim=3),
powers,
)
return powers, powers_for_angles_gen
def _azimuth_angles(
self,
azimuth_spread: torch.Tensor,
rician_k_factor: torch.Tensor,
cluster_powers: torch.Tensor,
angle_type: str,
return_cluster_angles: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Generate departure or arrival azimuth angles in degrees
(step 7 of section 7.5, TR 38.901).
:param azimuth_spread: Angle spread (ASD or ASA) depending on
``angle_type`` [deg],
shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param cluster_powers: Normalized path powers,
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]
:param angle_type: Type of angle to compute. Must be ``'aoa'`` or
``'aod'``.
"""
scenario = self._scenario
batch_size = scenario.batch_size
num_bs = scenario.num_bs
num_ut = scenario.num_ut
num_clusters_max = scenario.num_clusters_max
azimuth_spread = azimuth_spread.unsqueeze(3)
# Loading the angle spread
if angle_type == "aod":
azimuth_angles_los = scenario.los_aod
cluster_angle_spread = scenario.get_param("cASD")
else:
azimuth_angles_los = scenario.los_aoa
cluster_angle_spread = scenario.get_param("cASA")
# Adding cluster dimension for broadcasting
azimuth_angles_los = azimuth_angles_los.unsqueeze(3)
cluster_angle_spread = cluster_angle_spread.unsqueeze(3).unsqueeze(4)
# Compute C-phi constant
rician_k_factor = rician_k_factor.unsqueeze(3)
rician_k_factor_db = 10.0 * torch.log10(rician_k_factor) # to dB
c_phi_nlos = scenario.get_param("CPhiNLoS").unsqueeze(3)
c_phi_los = c_phi_nlos * (
1.1035
- 0.028 * rician_k_factor_db
- 0.002 * rician_k_factor_db.square()
+ 0.0001 * rician_k_factor_db.pow(3.0)
)
c_phi = torch.where(scenario.los.unsqueeze(3), c_phi_los, c_phi_nlos)
# Inverse Gaussian function
z = cluster_powers / cluster_powers.max(dim=3, keepdim=True).values
z = z.clamp(1e-6, 1.0)
azimuth_angles_prime = (2.0 * azimuth_spread / 1.4) * (
torch.sqrt(-torch.log(z)) / c_phi
)
# Introducing random variation
random_sign = self._spatial_binary_sign(
(batch_size, num_bs, num_ut, num_clusters_max)
)
random_sign = self._sort_clusters_by_delay(
random_sign, cluster_sort_indices
)
random_comp = (
self._spatial_normal((batch_size, num_bs, num_ut, num_clusters_max))
* azimuth_spread
/ 7.0
)
random_comp = self._sort_clusters_by_delay(
random_comp, cluster_sort_indices
)
azimuth_angles = random_sign * azimuth_angles_prime + random_comp + azimuth_angles_los
azimuth_angles = azimuth_angles - torch.where(
scenario.los.unsqueeze(3),
random_sign[:, :, :, :1] * azimuth_angles_prime[:, :, :, :1]
+ random_comp[:, :, :, :1],
torch.zeros(1, dtype=self.dtype, device=self.device),
)
cluster_azimuth_angles = wrap_angle_0_360(azimuth_angles)
cluster_azimuth_angles = torch.where(
cluster_azimuth_angles > 180.0,
cluster_azimuth_angles - 360.0,
cluster_azimuth_angles,
)
# Add offset angles to cluster angles to get the ray angles
ray_offsets = self._ray_offsets[: scenario.rays_per_cluster]
# Add dimensions for batch size, num bs, num ut, num clusters
ray_offsets = ray_offsets.reshape(1, 1, 1, 1, scenario.rays_per_cluster)
# Rays angles
azimuth_angles = azimuth_angles.unsqueeze(4)
azimuth_angles = azimuth_angles + cluster_angle_spread * ray_offsets
# Wrapping to (-180, 180)
azimuth_angles = wrap_angle_0_360(azimuth_angles)
azimuth_angles = torch.where(azimuth_angles > 180.0, azimuth_angles - 360.0, azimuth_angles)
if return_cluster_angles:
return azimuth_angles, cluster_azimuth_angles
return azimuth_angles
def _azimuth_angles_of_arrival(
self,
azimuth_spread_arrival: torch.Tensor,
rician_k_factor: torch.Tensor,
cluster_powers: torch.Tensor,
return_cluster_angles: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Compute azimuth angles of arrival (AoA)
(step 7 of section 7.5, TR 38.901).
:param azimuth_spread_arrival: Azimuth angle spread of arrival
(ASA) [deg], shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param cluster_powers: Normalized path powers,
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]
"""
return self._azimuth_angles(
azimuth_spread_arrival,
rician_k_factor,
cluster_powers,
"aoa",
return_cluster_angles=return_cluster_angles,
cluster_sort_indices=cluster_sort_indices,
)
def _azimuth_angles_of_departure(
self,
azimuth_spread_departure: torch.Tensor,
rician_k_factor: torch.Tensor,
cluster_powers: torch.Tensor,
return_cluster_angles: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Compute azimuth angles of departure (AoD)
(step 7 of section 7.5, TR 38.901).
:param azimuth_spread_departure: Azimuth angle spread of departure
(ASD) [deg], shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param cluster_powers: Normalized path powers,
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]
"""
return self._azimuth_angles(
azimuth_spread_departure,
rician_k_factor,
cluster_powers,
"aod",
return_cluster_angles=return_cluster_angles,
cluster_sort_indices=cluster_sort_indices,
)
def _zenith_angles(
self,
zenith_spread: torch.Tensor,
rician_k_factor: torch.Tensor,
cluster_powers: torch.Tensor,
angle_type: str,
return_cluster_angles: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Generate departure or arrival zenith angles in degrees
(step 7 of section 7.5, TR 38.901).
:param zenith_spread: Angle spread (ZSD or ZSA) depending on
``angle_type`` [deg],
shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param cluster_powers: Normalized path powers,
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]
:param angle_type: Type of angle to compute. Must be ``'zoa'`` or
``'zod'``.
"""
scenario = self._scenario
batch_size = scenario.batch_size
num_bs = scenario.num_bs
num_ut = scenario.num_ut
num_clusters_max = scenario.num_clusters_max
# Tensors giving UTs states
los = scenario.los
indoor_uts = scenario.indoor.unsqueeze(1)
if scenario.indoor_links_use_o2i_zenith_model:
o2i_uts = indoor_uts
else:
o2i_uts = torch.zeros_like(indoor_uts)
los_uts = los & (~o2i_uts)
nlos_uts = (~los) & (~o2i_uts)
# Adding cluster dimension for broadcasting
zenith_spread = zenith_spread.unsqueeze(3)
rician_k_factor = rician_k_factor.unsqueeze(3)
o2i_uts = o2i_uts.unsqueeze(3)
los_uts = los_uts.unsqueeze(3)
nlos_uts = nlos_uts.unsqueeze(3)
# Loading angle spread
if angle_type == "zod":
zenith_angles_los = scenario.los_zod
cluster_angle_spread = (3.0 / 8.0) * torch.pow(
torch.tensor(10.0, dtype=self.dtype, device=self.device),
scenario.lsp_log_mean[:, :, :, 6],
)
else:
cluster_angle_spread = scenario.get_param("cZSA")
zenith_angles_los = scenario.los_zoa
zod_offset = scenario.zod_offset
# Adding cluster dimension for broadcasting
zod_offset = zod_offset.unsqueeze(3)
zenith_angles_los = zenith_angles_los.unsqueeze(3)
cluster_angle_spread = cluster_angle_spread.unsqueeze(3)
# Compute the C_theta
rician_k_factor_db = 10.0 * torch.log10(rician_k_factor) # to dB
c_theta_nlos = scenario.get_param("CThetaNLoS").unsqueeze(3)
c_theta_los = c_theta_nlos * (
1.3086
+ 0.0339 * rician_k_factor_db
- 0.0077 * rician_k_factor_db.square()
+ 0.0002 * rician_k_factor_db.pow(3.0)
)
c_theta = torch.where(los_uts, c_theta_los, c_theta_nlos)
# Inverse Laplacian function
z = cluster_powers / cluster_powers.max(dim=3, keepdim=True).values
z = z.clamp(1e-6, 1.0)
zenith_angles_prime = -zenith_spread * torch.log(z) / c_theta
# Random component
random_sign = self._spatial_binary_sign(
(batch_size, num_bs, num_ut, num_clusters_max)
)
random_sign = self._sort_clusters_by_delay(
random_sign, cluster_sort_indices
)
random_comp = (
self._spatial_normal((batch_size, num_bs, num_ut, num_clusters_max))
* zenith_spread
/ 7.0
)
random_comp = self._sort_clusters_by_delay(
random_comp, cluster_sort_indices
)
# The center cluster angles depend on the UT scenario
zenith_angles = random_sign * zenith_angles_prime + random_comp
los_additional_comp = -(
random_sign[:, :, :, :1] * zenith_angles_prime[:, :, :, :1]
+ random_comp[:, :, :, :1]
- zenith_angles_los
)
if angle_type == "zod":
additional_comp = torch.where(
los_uts, los_additional_comp, zenith_angles_los + zod_offset
)
else:
additional_comp = torch.where(
los_uts,
los_additional_comp,
torch.zeros(1, dtype=self.dtype, device=self.device),
)
additional_comp = torch.where(nlos_uts, zenith_angles_los, additional_comp)
additional_comp = torch.where(
o2i_uts,
torch.tensor(90.0, dtype=self.dtype, device=self.device),
additional_comp,
)
zenith_angles = zenith_angles + additional_comp
cluster_zenith_angles = wrap_angle_0_360(zenith_angles)
cluster_zenith_angles = torch.where(
cluster_zenith_angles > 180.0,
360.0 - cluster_zenith_angles,
cluster_zenith_angles,
)
# Generating rays for every cluster
# Add offset angles to cluster angles to get the ray angles
ray_offsets = self._ray_offsets[: scenario.rays_per_cluster]
# Add dimensions for batch size, num bs, num ut, num clusters
ray_offsets = ray_offsets.reshape(1, 1, 1, 1, scenario.rays_per_cluster)
# Adding ray dimension for broadcasting
zenith_angles = zenith_angles.unsqueeze(4)
cluster_angle_spread = cluster_angle_spread.unsqueeze(4)
zenith_angles = zenith_angles + cluster_angle_spread * ray_offsets
# Wrapping to (0, 180)
zenith_angles = wrap_angle_0_360(zenith_angles)
zenith_angles = torch.where(zenith_angles > 180.0, 360.0 - zenith_angles, zenith_angles)
if return_cluster_angles:
return zenith_angles, cluster_zenith_angles
return zenith_angles
def _zenith_angles_of_arrival(
self,
zenith_spread_arrival: torch.Tensor,
rician_k_factor: torch.Tensor,
cluster_powers: torch.Tensor,
return_cluster_angles: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Compute zenith angles of arrival (ZoA)
(step 7 of section 7.5, TR 38.901).
:param zenith_spread_arrival: Zenith angle spread of arrival
(ZSA) [deg], shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param cluster_powers: Normalized path powers,
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]
"""
return self._zenith_angles(
zenith_spread_arrival,
rician_k_factor,
cluster_powers,
"zoa",
return_cluster_angles=return_cluster_angles,
cluster_sort_indices=cluster_sort_indices,
)
def _zenith_angles_of_departure(
self,
zenith_spread_departure: torch.Tensor,
rician_k_factor: torch.Tensor,
cluster_powers: torch.Tensor,
return_cluster_angles: bool = False,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Compute zenith angles of departure (ZoD)
(step 7 of section 7.5, TR 38.901).
:param zenith_spread_departure: Zenith angle spread of departure
(ZSD) [deg], shape [batch size, num of base stations, num of UTs]
:param rician_k_factor: Rician K-factor of each BS-UT link.
Used only for LoS links.
Shape [batch size, num of base stations, num of UTs].
:param cluster_powers: Normalized path powers,
shape [batch size, num of base stations, num of UTs, maximum number of
clusters]
"""
return self._zenith_angles(
zenith_spread_departure,
rician_k_factor,
cluster_powers,
"zod",
return_cluster_angles=return_cluster_angles,
cluster_sort_indices=cluster_sort_indices,
)
def _shuffle_angles(
self,
angles: torch.Tensor,
strongest_cluster_indices: Optional[torch.Tensor] = None,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Randomly shuffle a tensor carrying azimuth/zenith angles
of arrival/departure.
:param angles: Angles to shuffle, shape
[batch size, num of base stations, num of UTs, max num of clusters,
num of rays]
"""
scenario = self._scenario
batch_size = scenario.batch_size
num_bs = scenario.num_bs
num_ut = scenario.num_ut
# Create randomly shuffled indices by arg-sorting samples from a random
# normal distribution. With spatial consistency enabled, the random
# coupling is generated per UT from a spatial random field.
random_numbers = self._spatial_normal(
(
batch_size,
num_bs,
num_ut,
scenario.num_clusters_max,
scenario.rays_per_cluster,
),
use_sign_matrix=True,
)
random_numbers = self._sort_clusters_by_delay(
random_numbers, cluster_sort_indices
)
shuffled_indices = torch.argsort(random_numbers, dim=-1)
if strongest_cluster_indices is not None:
ray_indices = torch.arange(
scenario.rays_per_cluster,
dtype=torch.int64,
device=self.device,
)
grouped_indices = ray_indices.reshape(
*([1] * (random_numbers.dim() - 1)), -1
).expand_as(random_numbers).clone()
for group in (
self._subcluster_1_indices,
self._subcluster_2_indices,
self._subcluster_3_indices,
):
within_group = torch.argsort(
random_numbers.index_select(-1, group), dim=-1
)
group_permutation = group[within_group]
grouped_indices[..., group] = group_permutation
strongest = torch.zeros_like(
random_numbers[..., 0], dtype=torch.bool
)
strongest = strongest.scatter(
-1, strongest_cluster_indices, True
)
shuffled_indices = torch.where(
strongest.unsqueeze(-1), grouped_indices, shuffled_indices
)
# Shuffling the angles
shuffled_angles = torch.gather(angles, dim=4, index=shuffled_indices)
return shuffled_angles
def _random_coupling(
self,
aoa: torch.Tensor,
aod: torch.Tensor,
zoa: torch.Tensor,
zod: torch.Tensor,
strongest_cluster_indices: Optional[torch.Tensor] = None,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Randomly couple angles within a cluster for both azimuth and
elevation (step 8, TR 38.901).
:param aoa: Paths azimuth angles of arrival (AoA) [degree],
shape [batch size, num of base stations, num of UTs, max num of clusters,
num of rays]
:param aod: Paths azimuth angles of departure (AoD) [degree],
shape [batch size, num of base stations, num of UTs, max num of clusters,
num of rays]
:param zoa: Paths zenith angles of arrival (ZoA) [degree],
shape [batch size, num of base stations, num of UTs, max num of clusters,
num of rays]
:param zod: Paths zenith angles of departure (ZoD) [degree],
shape [batch size, num of base stations, num of UTs, max num of clusters,
num of rays]
"""
shuffled_aoa = self._shuffle_angles(
aoa, strongest_cluster_indices, cluster_sort_indices
)
shuffled_aod = self._shuffle_angles(
aod, strongest_cluster_indices, cluster_sort_indices
)
shuffled_zoa = self._shuffle_angles(
zoa, strongest_cluster_indices, cluster_sort_indices
)
shuffled_zod = self._shuffle_angles(
zod, strongest_cluster_indices, cluster_sort_indices
)
return shuffled_aoa, shuffled_aod, shuffled_zoa, shuffled_zod
def _cross_polarization_power_ratios(
self,
cluster_sort_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Generate cross-polarization power ratios (step 9, TR 38.901)."""
scenario = self._scenario
batch_size = scenario.batch_size
num_bs = scenario.num_bs
num_ut = scenario.num_ut
num_clusters = scenario.num_clusters_max
num_rays_per_cluster = scenario.rays_per_cluster
# Loading XPR mean and standard deviation
mu_xpr = scenario.get_param("muXPR")
std_xpr = scenario.get_param("sigmaXPR")
# Expanding for broadcasting with clusters and rays dims
mu_xpr = mu_xpr.unsqueeze(3).unsqueeze(4)
std_xpr = std_xpr.unsqueeze(3).unsqueeze(4)
# XPR are assumed to follow a log-normal distribution.
# Generate XPR in log-domain
x = (
self._spatial_normal(
(batch_size, num_bs, num_ut, num_clusters, num_rays_per_cluster)
)
* std_xpr
+ mu_xpr
)
x = self._sort_clusters_by_delay(x, cluster_sort_indices)
# To linear domain
cross_polarization_power_ratios = torch.pow(
torch.tensor(10.0, dtype=self.dtype, device=self.device), x / 10.0
)
return cross_polarization_power_ratios