#
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
"""Class definition for the OFDM Demodulator"""
from typing import Optional, Union
import numpy as np
import torch
from sionna.phy import Block, PI
from sionna.phy.config import Precision
from sionna.phy.signal import fft
from sionna.phy.utils import expand_to_rank
__all__ = ["OFDMDemodulator"]
[docs]
class OFDMDemodulator(Block):
r"""Computes the frequency-domain representation of an OFDM waveform
with cyclic prefix removal.
The demodulator assumes that the input sequence is generated by the
:class:`~sionna.phy.channel.TimeChannel`. For a single pair of antennas,
the received signal sequence is given as:
.. math::
y_b = \sum_{\ell =L_\text{min}}^{L_\text{max}} \bar{h}_\ell x_{b-\ell} + w_b, \quad b \in[L_\text{min}, N_B+L_\text{max}-1]
where :math:`\bar{h}_\ell` are the discrete-time channel taps,
:math:`x_{b}` is the transmitted signal, and :math:`w_\ell` Gaussian noise.
Starting from the first symbol, the demodulator cuts the input
sequence into pieces of size ``cyclic_prefix_length + fft_size``,
and throws away any trailing symbols. For each piece, the cyclic
prefix is removed and the ``fft_size``-point discrete Fourier
transform is computed. It is also possible that every OFDM symbol
has a cyclic prefix of different length.
Since the input sequence starts at time :math:`L_\text{min}`,
the FFT-window has a timing offset of :math:`L_\text{min}` symbols,
which leads to a subcarrier-dependent phase shift of
:math:`e^{\frac{j2\pi k L_\text{min}}{N}}`, where :math:`k`
is the subcarrier index, :math:`N` is the FFT size,
and :math:`L_\text{min} \le 0` is the largest negative time lag of
the discrete-time channel impulse response. This phase shift
is removed in this layer, by explicitly multiplying
each subcarrier by :math:`e^{\frac{-j2\pi k L_\text{min}}{N}}`.
This is a very important step to enable channel estimation with
sparse pilot patterns that needs to interpolate the channel frequency
response across subcarriers. It also ensures that the
channel frequency response `seen` by the time-domain channel
is close to the :class:`~sionna.phy.channel.OFDMChannel`.
:param fft_size: FFT size (i.e., the number of subcarriers).
:param l_min: The largest negative time lag of the discrete-time channel
impulse response. It should be the same value as that used by the
``cir_to_time_channel`` function.
:param cyclic_prefix_length: Integer or vector of integers indicating the
length of the cyclic prefix that is prepended to each OFDM symbol.
None of its elements can be larger than the FFT size. Defaults to `0`.
: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 tensor operations. If `None`,
:attr:`~sionna.phy.config.Config.device` is used.
:input inputs: [..., num_ofdm_symbols*(fft_size+cyclic_prefix_length)+n] or [..., num_ofdm_symbols*fft_size+sum(cyclic_prefix_length)+n], `torch.complex`.
Tensor containing the time-domain signal along the last dimension.
`n` is a nonnegative integer.
:output x: [..., num_ofdm_symbols, fft_size], `torch.complex`.
Tensor containing the OFDM resource grid along the last
two dimensions.
.. rubric:: Examples
.. code-block:: python
import torch
from sionna.phy.ofdm import OFDMModulator, OFDMDemodulator
modulator = OFDMModulator(cyclic_prefix_length=16)
demodulator = OFDMDemodulator(fft_size=72, l_min=0, cyclic_prefix_length=16)
x_freq = torch.randn(64, 14, 72, dtype=torch.complex64)
x_time = modulator(x_freq)
x_hat = demodulator(x_time)
print(torch.allclose(x_freq, x_hat, atol=1e-5))
# True
"""
def __init__(
self,
fft_size: int,
l_min: int,
cyclic_prefix_length: Union[int, np.ndarray, torch.Tensor] = 0,
precision: Optional[Precision] = None,
device: Optional[str] = None,
**kwargs,
) -> None:
super().__init__(precision=precision, device=device, **kwargs)
self._fft_size: Optional[int] = None
self._l_min: Optional[int] = None
self._cp_length_scalar: Optional[int] = None # Cached scalar for call()
# Register tensors as buffers for CUDA graph compatibility
self.register_buffer("_cyclic_prefix_length", None)
self.register_buffer("_phase_compensation", None)
# Store symbol_length as buffer for torch.compile compatibility
self.register_buffer("_symbol_length", None)
self.register_buffer("_ind", None)
self.fft_size = fft_size
self.l_min = l_min
self.cyclic_prefix_length = cyclic_prefix_length
# Compute phase compensation at init time to avoid CUDA graph issues
self._init_phase_compensation()
def _init_phase_compensation(self) -> None:
"""Compute and register phase compensation buffer."""
tmp = (
-2
* PI
* float(self.l_min)
/ float(self.fft_size)
* torch.arange(self.fft_size, dtype=self.dtype, device=self.device)
)
self.register_buffer(
"_phase_compensation", torch.exp(torch.complex(torch.zeros_like(tmp), tmp))
)
@property
def fft_size(self) -> int:
"""FFT size (number of subcarriers)"""
return self._fft_size
@fft_size.setter
def fft_size(self, value: int) -> None:
assert value > 0, "`fft_size` must be positive."
self._fft_size = int(value)
@property
def l_min(self) -> int:
"""Largest negative time lag of the discrete-time channel"""
return self._l_min
@l_min.setter
def l_min(self, value: int) -> None:
assert value <= 0, "`l_min` must be nonpositive."
self._l_min = int(value)
@property
def cyclic_prefix_length(self) -> torch.Tensor:
"""Cyclic prefix length (scalar or per-symbol)"""
return self._cyclic_prefix_length
@cyclic_prefix_length.setter
def cyclic_prefix_length(self, value: Union[int, np.ndarray, torch.Tensor]) -> None:
if isinstance(value, (int, float)):
value = torch.tensor([value], dtype=torch.int32, device=self.device)
elif isinstance(value, np.ndarray):
value = torch.tensor(value, dtype=torch.int32, device=self.device)
else:
value = value.to(dtype=torch.int32, device=self.device)
if not torch.all(value >= 0):
raise ValueError("`cyclic_prefix_length` must be nonnegative.")
if not 0 <= value.dim() <= 1:
raise ValueError("`cyclic_prefix_length` must be of rank 0 or 1.")
# Store as 0D if scalar, 1D otherwise
if value.numel() == 1:
# Register as buffer for CUDA graph compatibility
self.register_buffer("_cyclic_prefix_length", value.squeeze())
# Cache scalar value to avoid .item() during tracing
self._cp_length_scalar = int(value.item())
# Store symbol_length as buffer for torch.compile compatibility
symbol_length = self.fft_size + self._cp_length_scalar
self.register_buffer(
"_symbol_length",
torch.tensor(symbol_length, dtype=torch.int64, device=self.device),
)
else:
self.register_buffer("_cyclic_prefix_length", value)
self._cp_length_scalar = None
self.register_buffer("_symbol_length", None)
[docs]
def build(self, input_shape: tuple) -> None:
"""Build the demodulator based on input shape."""
cp_len = self._cyclic_prefix_length
if cp_len.dim() == 0:
# Same CP length for all OFDM symbols
# No pre-computation needed - done dynamically in call() for
# torch.compile compatibility
pass
else:
# Individual CP lengths for OFDM symbols
num_ofdm_symbols = cp_len.shape[0]
row_lengths = cp_len + self.fft_size
offsets = torch.cumsum(
torch.cat(
[
torch.zeros(1, dtype=torch.int32, device=self.device),
row_lengths[:-1],
]
),
dim=0,
)
# Build indices for each OFDM symbol
# Convert to lists once to avoid .item() calls during tracing
cp_lengths = cp_len.tolist()
offsets_list = offsets.tolist()
indices_list = []
for i in range(num_ofdm_symbols):
cp_length_i = cp_lengths[i]
start = offsets_list[i] + cp_length_i
indices = torch.arange(
start, start + self.fft_size, dtype=torch.int64, device=self.device
)
indices_list.append(indices)
# [num_ofdm_symbols, fft_size]
self.register_buffer("_ind", torch.stack(indices_list, dim=0))
def call(self, inputs: torch.Tensor) -> torch.Tensor:
"""Demodulate OFDM waveform onto a resource grid."""
cp_len = self._cyclic_prefix_length
if cp_len.dim() == 0:
# Same CP length for all OFDM symbols
# Use cached scalar for CP length
cp_length = self._cp_length_scalar
symbol_length = self._symbol_length # fft_size + cp_length as tensor
# Compute number of full OFDM symbols dynamically for torch.compile
# Use integer division on tensor to make it traceable
input_length = inputs.shape[-1]
num_ofdm_symbols = input_length // symbol_length
rest = input_length % symbol_length
# Cut last samples that do not fit into an OFDM symbol
if rest > 0:
inputs = inputs[..., :-rest]
# Reshape input to separate OFDM symbols
# Use view with -1 for the num_ofdm_symbols dimension to avoid
# shape mismatch issues during torch.compile tracing
batch_shape = inputs.shape[:-1]
x = inputs.view(*batch_shape, -1, symbol_length)
# Remove cyclic prefix
if cp_length > 0:
x = x[..., cp_length:]
else:
# Individual CP length for OFDM symbols
x = inputs[..., self._ind]
# Compute FFT
x = fft(x, precision=self.precision)
# Apply phase shift compensation to all subcarriers
rot = self._phase_compensation.to(x.dtype)
rot = expand_to_rank(rot, x.dim(), axis=0)
x = x * rot
# Shift DC subcarrier to the middle
x = torch.fft.fftshift(x, dim=-1)
return x