#
# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0#
"""Classes and functions related to MIMO transmit precoding"""
import tensorflow as tf
import math
from sionna.phy import PI, dtypes, config
from sionna.phy.utils import expand_to_rank
[docs]
def rzf_precoding_matrix(h,
alpha=0.,
precision=None):
# pylint: disable=line-too-long
r"""Computes the Regularized Zero-Forcing (RZF) Precoder
This function computes the RZF precoding matrix for a MIMO link, assuming the
following model:
.. math::
\mathbf{y} = \mathbf{H}\mathbf{G}\mathbf{x} + \mathbf{n}
where :math:`\mathbf{y}\in\mathbb{C}^K` is the received signal vector,
:math:`\mathbf{H}\in\mathbb{C}^{K\times M}` is the known channel matrix,
:math:`\mathbf{G}\in\mathbb{C}^{M\times K}` is the precoding matrix,
:math:`\mathbf{x}\in\mathbb{C}^K` is the symbol vector to be precoded,
and :math:`\mathbf{n}\in\mathbb{C}^K` is a noise vector.
The precoding matrix :math:`\mathbf{G}` is defined as :
.. math::
\mathbf{G} = \mathbf{V}\mathbf{D}
where
.. math::
\mathbf{V} &= \mathbf{H}^{\mathsf{H}}\left(\mathbf{H}
\mathbf{H}^{\mathsf{H}} + \alpha \mathbf{I} \right)^{-1}\\
\mathbf{D} &= \mathop{\text{diag}}\left( \lVert \mathbf{v}_{k} \rVert_2^{-1}, k=0,\dots,K-1 \right)
where :math:`\alpha>0` is the regularization parameter. The matrix :math:`\mathbf{V}`
ensures that each stream is precoded with a unit-norm vector,
i.e., :math:`\mathop{\text{tr}}\left(\mathbf{G}\mathbf{G}^{\mathsf{H}}\right)=K`.
The function returns the matrix :math:`\mathbf{G}`.
Input
------
h : [...,K,M], `tf.complex`
2+D tensor containing the channel matrices
alpha : `0.` (default) | [...], `tf.float32`
Regularization parameter
precision : `None` (default) | "single" | "double"
Precision used for internal calculations and outputs.
If set to `None`,
:attr:`~sionna.phy.config.Config.precision` is used
Output
-------
g : [...,M,K], `tf.complex`
2+D tensor containing the precoding matrices
"""
# Cast inputs
if precision is None:
cdtype = config.tf_cdtype
else:
cdtype = dtypes[precision]["tf"]["cdtype"]
h = tf.cast(h, dtype=cdtype)
alpha = tf.cast(alpha, dtype=cdtype)
# Compute pseudo inverse for precoding
g = tf.matmul(h, h, adjoint_b=True)
alpha = expand_to_rank(alpha, tf.rank(g), axis=-1)
g = g + alpha * tf.eye(tf.shape(g)[-1],
batch_shape=tf.shape(g)[:-2], dtype=cdtype)
l = tf.linalg.cholesky(g)
g = tf.linalg.cholesky_solve(l, h)
g = tf.linalg.adjoint(g)
# Normalize each column to unit power
norm = tf.sqrt(tf.reduce_sum(tf.abs(g)**2, axis=-2, keepdims=True))
g = tf.math.divide_no_nan(g, tf.cast(norm, g.dtype))
return g
[docs]
def cbf_precoding_matrix(h, precision=None):
# pylint: disable=line-too-long
r"""Computes the conjugate beamforming (CBF) Precoder
This function computes the CBF precoding matrix for a MIMO link, assuming the
following model:
.. math::
\mathbf{y} = \mathbf{H}\mathbf{G}\mathbf{x} + \mathbf{n}
where :math:`\mathbf{y}\in\mathbb{C}^K` is the received signal vector,
:math:`\mathbf{H}\in\mathbb{C}^{K\times M}` is the known channel matrix,
:math:`\mathbf{G}\in\mathbb{C}^{M\times K}` is the precoding matrix,
:math:`\mathbf{x}\in\mathbb{C}^K` is the symbol vector to be precoded,
and :math:`\mathbf{n}\in\mathbb{C}^K` is a noise vector.
The precoding matrix :math:`\mathbf{G}` is defined as :
.. math::
\mathbf{G} = \mathbf{V}\mathbf{D}
where
.. math::
\mathbf{V} &= \mathbf{H}^{\mathsf{H}} \\
\mathbf{D} &= \mathop{\text{diag}}\left( \lVert \mathbf{v}_{k} \rVert_2^{-1}, k=0,\dots,K-1 \right).
The matrix :math:`\mathbf{V}`
ensures that each stream is precoded with a unit-norm vector,
i.e., :math:`\mathop{\text{tr}}\left(\mathbf{G}\mathbf{G}^{\mathsf{H}}\right)=K`.
The function returns the matrix :math:`\mathbf{G}`.
Input
------
h : [...,K,M], `tf.complex`
2+D tensor containing the channel matrices
precision : `None` (default) | "single" | "double"
Precision used for internal calculations and outputs.
If set to `None`,
:attr:`~sionna.phy.config.Config.precision` is used
Output
-------
g : [...,M,K], `tf.complex`
2+D tensor containing the precoding matrices
"""
# Cast inputs
if precision is None:
cdtype = config.tf_cdtype
else:
cdtype = dtypes[precision]["tf"]["cdtype"]
h = tf.cast(h, dtype=cdtype)
# Compute conjugate transpose of channel matrix
g = tf.linalg.adjoint(h)
# Normalize each column to unit power
norm = tf.sqrt(tf.reduce_sum(tf.abs(g)**2, axis=-2, keepdims=True))
g = tf.math.divide_no_nan(g, tf.cast(norm, g.dtype))
return g
[docs]
def rzf_precoder(x,
h,
alpha=0.,
return_precoding_matrix=False,
precision=None):
# pylint: disable=line-too-long
r"""Regularized Zero-Forcing (RZF) Precoder
This function implements RZF precoding for a MIMO link, assuming the
following model:
.. math::
\mathbf{y} = \mathbf{H}\mathbf{G}\mathbf{x} + \mathbf{n}
where :math:`\mathbf{y}\in\mathbb{C}^K` is the received signal vector,
:math:`\mathbf{H}\in\mathbb{C}^{K\times M}` is the known channel matrix,
:math:`\mathbf{G}\in\mathbb{C}^{M\times K}` is the precoding matrix,
:math:`\mathbf{x}\in\mathbb{C}^K` is the symbol vector to be precoded,
and :math:`\mathbf{n}\in\mathbb{C}^K` is a noise vector.
The precoding matrix :math:`\mathbf{G}` is defined as (Eq. 4.37) [BHS2017]_ :
.. math::
\mathbf{G} = \mathbf{V}\mathbf{D}
where
.. math::
\mathbf{V} &= \mathbf{H}^{\mathsf{H}}\left(\mathbf{H} \mathbf{H}^{\mathsf{H}} + \alpha \mathbf{I} \right)^{-1}\\
\mathbf{D} &= \mathop{\text{diag}}\left( \lVert \mathbf{v}_{k} \rVert_2^{-1}, k=0,\dots,K-1 \right)
where :math:`\alpha>0` is the regularization parameter.
This ensures that each stream is precoded with a unit-norm vector,
i.e., :math:`\mathop{\text{tr}}\left(\mathbf{G}\mathbf{G}^{\mathsf{H}}\right)=K`.
The function returns the precoded vector :math:`\mathbf{G}\mathbf{x}`.
Input
-----
x : [...,K], `tf.complex`
Symbol vectors to be precoded
h : [...,K,M], `tf.complex`
Channel matrices
alpha : `0.` (default) | [...], `tf.float32`
Regularization parameter
return_precoding_matrices : `bool`, (default, `False`)
Indicates if the precoding matrices should be returned or not
precision : `None` (default) | "single" | "double"
Precision used for internal calculations and outputs.
If set to `None`,
:attr:`~sionna.phy.config.Config.precision` is used
Output
-------
x_precoded : [...,M], `tf.complex`
Precoded symbol vectors
g : [...,M,K], `tf.complex`
Precoding matrices. Only returned
if ``return_precoding_matrices=True``.
"""
# Cast inputs
if precision is None:
cdtype = config.tf_cdtype
else:
cdtype = dtypes[precision]["tf"]["cdtype"]
x = tf.cast(x, dtype=cdtype)
h = tf.cast(h, dtype=cdtype)
# Compute the precoding matrix
g = rzf_precoding_matrix(h, alpha=alpha, precision=precision)
# Expand last dim of `x` for precoding
x_precoded = tf.expand_dims(x, -1)
# Precode
x_precoded = tf.squeeze(tf.matmul(g, x_precoded), -1)
if return_precoding_matrix:
return (x_precoded, g)
else:
return x_precoded
[docs]
def grid_of_beams_dft_ula(num_ant, oversmpl=1, precision=None):
# pylint: disable=line-too-long
r""" Computes the Discrete Fourier Transform (DFT) Grid of Beam (GoB)
coefficients for a uniform linear array (ULA)
The coefficient applied to antenna :math:`n` for beam :math:`m` is expressed
as:
.. math::
c_n^m = e^{\frac{2\pi n m}{N O}}, \quad n=0,\dots,N-1, \ m=0,\dots,NO
where :math:`N` is the number of antennas ``num_ant`` and :math:`O` is the oversampling
factor ``oversmpl``.
Note that the main lobe of beam :math:`m` points in the azimuth direction
:math:`\theta = \mathrm{arc sin} \left( 2\frac{m}{N} \right)` if :math:`m\le
N/2` and :math:`\theta = \mathrm{arc sin} \left( 2\frac{m-N}{N} \right)` if
:math:`m\ge N/2`, where :math:`\theta=0` defines the perpendicular to the
antenna array.
Input
------
num_ant : `int`
Number of antennas
oversmpl : `int`, (default 1)
Oversampling factor
precision : `None` (default) | "single" | "double"
Precision used for internal calculations and outputs.
If set to `None`,
:attr:`~sionna.phy.config.Config.precision` is used.
Output
-------
gob : [num_ant x oversmpl, num_ant], `tf.complex`
The :math:`m`-th row contains the `num_ant` antenna coefficients for
the :math:`m`-th DFT beam.
"""
if precision is None:
rdtype = config.tf_rdtype
else:
rdtype = dtypes[precision]["tf"]["rdtype"]
oversmpl = int(oversmpl)
# Beam indices: [0, .., num_ant * oversmpl - 1]
beam_ind = tf.range(num_ant * oversmpl, dtype=tf.float32)[:, tf.newaxis]
beam_ind = tf.cast(beam_ind, dtype=rdtype)
# Antenna indices: [0, .., num_ant - 1]
antenna_ind = tf.range(num_ant, dtype=tf.float32)[tf.newaxis, :]
antenna_ind = tf.cast(antenna_ind, dtype=rdtype)
# Combine real and imaginary part and normalize power to 1
phases = 2 * PI * beam_ind * antenna_ind / (num_ant * oversmpl)
gob = tf.complex(tf.cos(phases), tf.sin(phases)) / math.sqrt(num_ant)
return gob
[docs]
def grid_of_beams_dft(num_ant_v,
num_ant_h,
oversmpl_v=1,
oversmpl_h=1,
precision=None):
# pylint: disable=line-too-long
r""" Computes the Discrete Fourier Transform (DFT) Grid of Beam (GoB)
coefficients for a uniform rectangular array (URA)
GoB indices are arranged over a 2D grid indexed by :math:`(m_v,m_h)`.
The coefficient of the beam with index :math:`(m_v,m_h)` applied to the
antenna located at row :math:`n_v` and column :math:`n_h` of the rectangular
array is expressed as:
.. math::
c_{n_v,n_h}^{m_v,m_h} = e^{\frac{2\pi n_h m_v}{N_h O_h}} e^{\frac{2\pi n_h m_h}{N_v O_v}}
where :math:`n_v=0,\dots,N_v-1`, :math:`n_h=0,\dots,N_h-1`,
:math:`m_v=0,\dots,N_v O_v`, :math:`m_h=0,\dots,N_h O_h`, :math:`N` is the
number of antennas ``num_ant`` and :math:`O_v,O_h` are the oversampling
factor ``oversmpl_v``, ``oversmpl_h`` in the vertical and
horizontal direction, respectively.
We can rewrite more concisely the matrix coefficients
:math:`c^{m_v,m_h}` as follows:
.. math::
c^{m_v,m_h} = c^{m_v} \otimes c^{m_h}
where :math:`\otimes` denotes the Kronecker product and
:math:`c^{m_v},c^{m_h}` are the ULA DFT beams computed as in
:func:`~sionna.phy.mimo.grid_of_beams_dft_ula`.
Such a DFT GoB is, e.g., defined in Section 5.2.2.2.1 [3GPP38214]_.
Input
-----
num_ant_v : `int`
Number of antenna rows (i.e., in vertical direction) of the rectangular
array
num_ant_h : `int`
Number of antenna columns (i.e., in horizontal direction) of the
rectangular array.
oversmpl_v : `int`, (default 1)
Oversampling factor in vertical direction
oversmpl_h : `int`, (default 1)
Oversampling factor in horizontal direction
precision : `None` (default) | "single" | "double"
Precision used for internal calculations and outputs.
If set to `None`,
:attr:`~sionna.phy.config.Config.precision` is used.
Output
------
gob : [num_ant_v x oversmpl_v, num_ant_h x oversmpl_h, num_ant_v x num_ant_h], `tf.complex`
The elements :math:`[m_v,m_h,:]` contain the antenna coefficients of the
DFT beam with index pair :math:`(m_v,m_h)`.
"""
# Compute the DFT coefficients to be applied in the vertical direction
gob_v = grid_of_beams_dft_ula(num_ant_v,
oversmpl=oversmpl_v,
precision=precision)
gob_v = gob_v[:, tf.newaxis, :, tf.newaxis]
# Compute the DFT coefficients to be applied in the horizontal direction
gob_h = grid_of_beams_dft_ula(num_ant_h,
oversmpl=oversmpl_h,
precision=precision)
gob_h = gob_h[tf.newaxis, :, tf.newaxis, :]
# Kronecker product:
# [num_ant_v * oversmpl_v , num_ant_h * oversmpl_v, num_ant_v, num_ant_h]
coef_vh = tf.math.multiply(gob_h, gob_v)
# Flatten the last two dimensions to produce 1-dimensional precoding vectors
# [num_ant_v * oversmpl_v , num_ant_h * oversmpl_v, num_ant_v x num_ant_h]
coef_vh = flatten_precoding_mat(coef_vh)
return coef_vh
[docs]
def flatten_precoding_mat(precoding_mat, by_column=True):
# pylint: disable=line-too-long
r"""Flattens a [..., num_ant_v, num_ant_h] precoding matrix associated with
a rectangular array by producing a [..., num_ant_v x num_ant_h] precoding vector
Input
------
precoding_mat : [..., num_antennas_vertical, num_antennas_horizontal], `tf.complex`
Precoding matrix. The element :math:`(i,j)` contains the precoding
coefficient of the antenna element located at row :math:`i` and column
:math:`j` of a rectangular antenna array.
by_column : `bool`, (default `True`)
If `True`, then flattening occurs on a per-column basis, i.e., the first
column is appended to the second, and so on. Else, flattening is performed on
a per-row basis.
Output
-------
: [..., num_antennas_vertical x num_antennas_horizontal], `tf.complex`
Flattened precoding matrix
"""
# Transpose the last two dimensions
if by_column:
precoding_mat = tf.linalg.matrix_transpose(precoding_mat)
# Flatten the last two dimensions
precoding_vec = tf.reshape(
precoding_mat, precoding_mat.shape[:-2] + [math.prod(precoding_mat.shape[2:])])
return precoding_vec
[docs]
def normalize_precoding_power(precoding_vec,
tx_power_list=None,
precision=None):
# pylint: disable=line-too-long
r""" Normalizes the beam coefficient power to 1 by default, or to
``tx_power_list`` if provided as input
Input
------
precoding_vec : [N,M], `tf.complex`
Each row contains a set of antenna coefficients whose power is to be normalized.
tx_power_list : [N], float
The :math:`i`-th element defines the power of the :math:`i`-th precoding
vector.
precision : `None` (default) | "single" | "double"
Precision used for internal calculations and outputs.
If set to `None`,
:attr:`~sionna.phy.config.Config.precision` is used.
Output
-------
: [N,M] `tf.complex`
Normalized antenna coefficients
"""
if precision is None:
cdtype = config.tf_cdtype
else:
cdtype = dtypes[precision]["tf"]["cdtype"]
precoding_vec = tf.cast(precoding_vec, dtype=cdtype)
if len(precoding_vec.shape)==1:
precoding_vec = precoding_vec[tf.newaxis, :]
if tx_power_list is None:
# By default, power is normalized to 1
tx_power_list = [1] * precoding_vec.shape[0]
precoding_vec_norm = tf.norm(precoding_vec, axis=1)[:, tf.newaxis]
tx_power = tf.constant(tx_power_list, dtype=cdtype)[:, tf.newaxis]
# Normalize the power of each row
precoding_vec = tf.math.multiply(tf.math.divide(
precoding_vec, precoding_vec_norm), tx_power)
return precoding_vec