#
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
"""Classes for spatially correlated flat-fading channels"""
from abc import abstractmethod
from typing import Optional
import torch
from sionna.phy.object import Object
from sionna.phy.utils import expand_to_rank
[docs]
class SpatialCorrelation(Object):
r"""Abstract class that defines an interface for spatial correlation functions
The :class:`~sionna.phy.channel.FlatFadingChannel` model can be configured with a
spatial correlation model.
:param precision: Precision used for internal calculations and outputs.
If set to `None`, :attr:`~sionna.phy.config.Config.precision` is used.
:param device: Device for computation (e.g., 'cpu', 'cuda:0').
If `None`, :attr:`~sionna.phy.config.Config.device` is used.
:input h: [...], `torch.complex`.
Tensor of arbitrary shape containing spatially uncorrelated
channel coefficients.
:output h_corr: [...], `torch.complex`.
Tensor of the same shape as ``h`` containing the spatially
correlated channel coefficients.
"""
@abstractmethod
def __call__(self, h: torch.Tensor) -> torch.Tensor:
raise NotImplementedError
[docs]
class KroneckerModel(SpatialCorrelation):
r"""Kronecker model for spatial correlation
Given a batch of matrices :math:`\mathbf{H}\in\mathbb{C}^{M\times K}`,
:math:`\mathbf{R}_\text{tx}\in\mathbb{C}^{K\times K}`, and
:math:`\mathbf{R}_\text{rx}\in\mathbb{C}^{M\times M}`, this function
will generate the following output:
.. math::
\mathbf{H}_\text{corr} = \mathbf{R}^{\frac12}_\text{rx} \mathbf{H} \mathbf{R}^{\frac12}_\text{tx}
Note that :math:`\mathbf{R}_\text{tx}\in\mathbb{C}^{K\times K}` and
:math:`\mathbf{R}_\text{rx}\in\mathbb{C}^{M\times M}`
must be positive semi-definite, such as the ones generated by
:meth:`~sionna.phy.channel.exp_corr_mat`.
:param r_tx: Transmit correlation matrices of shape [..., K, K]. If
the rank of ``r_tx`` is smaller than that of the input ``h``,
it will be broadcast.
:param r_rx: Receive correlation matrices of shape [..., M, M]. If
the rank of ``r_rx`` is smaller than that of the input ``h``,
it will be broadcast.
:param precision: Precision used for internal calculations and outputs.
If set to `None`, :attr:`~sionna.phy.config.Config.precision` is used.
:param device: Device for computation (e.g., 'cpu', 'cuda:0').
If `None`, :attr:`~sionna.phy.config.Config.device` is used.
:input h: [..., M, K], `torch.complex`.
Spatially uncorrelated channel coefficients.
:output h_corr: [..., M, K], `torch.complex`.
Spatially correlated channel coefficients.
.. rubric:: Examples
.. code-block:: python
import torch
from sionna.phy.channel import exp_corr_mat, KroneckerModel
# Create correlation matrices
r_tx = exp_corr_mat(0.4, 4) # 4x4 TX correlation
r_rx = exp_corr_mat(0.9, 16) # 16x16 RX correlation
# Create model
kron = KroneckerModel(r_tx, r_rx)
# Apply to channel matrix
h = torch.randn(32, 16, 4, dtype=torch.complex64)
h_corr = kron(h)
print(h_corr.shape)
# torch.Size([32, 16, 4])
# For use with torch.compile, update matrices in-place:
new_r_tx = exp_corr_mat(0.5, 4)
kron.r_tx.copy_(new_r_tx) # In-place update avoids graph breaks
"""
def __init__(
self,
r_tx: Optional[torch.Tensor] = None,
r_rx: Optional[torch.Tensor] = None,
precision: Optional[str] = None,
device: Optional[str] = None,
) -> None:
super().__init__(precision=precision, device=device)
# Register as buffers for proper device handling and state_dict support
# Use register_buffer with None for optional correlation matrices
if r_tx is not None:
r_tx = r_tx.to(device=self.device, dtype=self.cdtype)
if r_rx is not None:
r_rx = r_rx.to(device=self.device, dtype=self.cdtype)
self.register_buffer("_r_tx", r_tx)
self.register_buffer("_r_rx", r_rx)
@property
def r_tx(self) -> Optional[torch.Tensor]:
r"""Get/set transmit correlation matrices.
For use within ``torch.compile``, use in-place updates via
``model.r_tx.copy_(new_value)`` to avoid graph breaks.
"""
return self._r_tx
@r_tx.setter
def r_tx(self, value: Optional[torch.Tensor]) -> None:
if value is not None:
value = value.to(device=self.device, dtype=self.cdtype)
if self._r_tx is not None and self._r_tx.shape == value.shape:
# In-place update to avoid graph breaks in torch.compile
self._r_tx.copy_(value)
return
# Fall back to buffer registration for new shape or None
self.register_buffer("_r_tx", value)
@property
def r_rx(self) -> Optional[torch.Tensor]:
r"""Get/set receive correlation matrices.
For use within ``torch.compile``, use in-place updates via
``model.r_rx.copy_(new_value)`` to avoid graph breaks.
"""
return self._r_rx
@r_rx.setter
def r_rx(self, value: Optional[torch.Tensor]) -> None:
if value is not None:
value = value.to(device=self.device, dtype=self.cdtype)
if self._r_rx is not None and self._r_rx.shape == value.shape:
# In-place update to avoid graph breaks in torch.compile
self._r_rx.copy_(value)
return
# Fall back to buffer registration for new shape or None
self.register_buffer("_r_rx", value)
def __call__(self, h: torch.Tensor) -> torch.Tensor:
if self._r_tx is not None:
# Cast to input dtype for compatibility
r_tx = self._r_tx.to(dtype=h.dtype)
# Use cholesky_ex for CUDA graph compatibility
l_tx, _ = torch.linalg.cholesky_ex(r_tx, check_errors=False)
h = h @ l_tx.mH
if self._r_rx is not None:
# Cast to input dtype for compatibility
r_rx = self._r_rx.to(dtype=h.dtype)
# Use cholesky_ex for CUDA graph compatibility
l_rx, _ = torch.linalg.cholesky_ex(r_rx, check_errors=False)
h = l_rx @ h
return h
[docs]
class PerColumnModel(SpatialCorrelation):
r"""Per-column model for spatial correlation
Given a batch of matrices :math:`\mathbf{H}\in\mathbb{C}^{M\times K}`
and correlation matrices :math:`\mathbf{R}_k\in\mathbb{C}^{M\times M}, k=1,\dots,K`,
this function will generate the output :math:`\mathbf{H}_\text{corr}\in\mathbb{C}^{M\times K}`,
with columns
.. math::
\mathbf{h}^\text{corr}_k = \mathbf{R}^{\frac12}_k \mathbf{h}_k,\quad k=1, \dots, K
where :math:`\mathbf{h}_k` is the kth column of :math:`\mathbf{H}`.
Note that all :math:`\mathbf{R}_k\in\mathbb{C}^{M\times M}` must
be positive semi-definite, such as the ones generated
by :meth:`~sionna.phy.channel.one_ring_corr_mat`.
This model is typically used to simulate a MIMO channel between multiple
single-antenna users and a base station with multiple antennas.
The resulting SIMO channel for each user has a different spatial correlation.
:param r_rx: Receive correlation matrices of shape [..., M, M]. If
the rank of ``r_rx`` is smaller than that of the input ``h``,
it will be broadcast. For a typical use of this model, ``r_rx``
has shape [..., K, M, M], i.e., a different correlation matrix for each
column of ``h``.
:param precision: Precision used for internal calculations and outputs.
If set to `None`, :attr:`~sionna.phy.config.Config.precision` is used.
:param device: Device for computation (e.g., 'cpu', 'cuda:0').
If `None`, :attr:`~sionna.phy.config.Config.device` is used.
:input h: [..., M, K], `torch.complex`.
Spatially uncorrelated channel coefficients.
:output h_corr: [..., M, K], `torch.complex`.
Spatially correlated channel coefficients.
.. rubric:: Examples
.. code-block:: python
import torch
from sionna.phy.channel import one_ring_corr_mat, PerColumnModel
# Create per-column correlation matrices (4 users, 16 antennas)
r_rx = one_ring_corr_mat(torch.tensor([-45., -15., 0., 30.]), 16)
# Create model
per_col = PerColumnModel(r_rx)
# Apply to channel matrix
h = torch.randn(32, 16, 4, dtype=torch.complex64)
h_corr = per_col(h)
print(h_corr.shape)
# torch.Size([32, 16, 4])
"""
def __init__(
self,
r_rx: Optional[torch.Tensor] = None,
precision: Optional[str] = None,
device: Optional[str] = None,
) -> None:
super().__init__(precision=precision, device=device)
# Register as buffer for proper device handling and state_dict support
if r_rx is not None:
r_rx = r_rx.to(device=self.device, dtype=self.cdtype)
self.register_buffer("_r_rx", r_rx)
@property
def r_rx(self) -> Optional[torch.Tensor]:
r"""Get/set receive correlation matrices.
For use within ``torch.compile``, use in-place updates via
``model.r_rx.copy_(new_value)`` to avoid graph breaks.
"""
return self._r_rx
@r_rx.setter
def r_rx(self, value: Optional[torch.Tensor]) -> None:
if value is not None:
value = value.to(device=self.device, dtype=self.cdtype)
if self._r_rx is not None and self._r_rx.shape == value.shape:
# In-place update to avoid graph breaks in torch.compile
self._r_rx.copy_(value)
return
# Fall back to buffer registration for new shape or None
self.register_buffer("_r_rx", value)
def __call__(self, h: torch.Tensor) -> torch.Tensor:
if self._r_rx is not None:
# Cast to input dtype for compatibility
r_rx = self._r_rx.to(dtype=h.dtype)
# Use cholesky_ex for CUDA graph compatibility
l_rx, _ = torch.linalg.cholesky_ex(r_rx, check_errors=False)
# Swap last two axes: [..., M, K] -> [..., K, M]
h = h.swapaxes(-2, -1)
# Add dimension for matrix-vector product: [..., K, M] -> [..., K, M, 1]
h = h.unsqueeze(-1)
# Expand l_rx to match h rank
l_rx = expand_to_rank(l_rx, h.dim(), 0)
# Apply correlation: [..., K, M, M] @ [..., K, M, 1] -> [..., K, M, 1]
h = l_rx @ h
# Remove last dimension: [..., K, M, 1] -> [..., K, M]
h = h.squeeze(-1)
# Swap back: [..., K, M] -> [..., M, K]
h = h.swapaxes(-2, -1)
return h