kaira.channels.NonlinearChannel

Inheritance diagram for NonlinearChannel

class kaira.channels.NonlinearChannel(nonlinear_fn, add_noise=False, avg_noise_power: float | None = None, snr_db: float | None = None, complex_mode='direct', *args: Any, **kwargs: Any)

Bases: BaseChannel

General nonlinear channel with configurable transfer function.

Models various nonlinear effects by applying a user-specified nonlinear function to the input signal, optionally followed by additive noise. Handles both real and complex-valued signals. Common nonlinear models include the Saleh model for traveling-wave tube amplifiers [Saleh, 1981].

Mathematical Model:

y = f(x) + n where f is a nonlinear function and n is optional noise

Parameters:

• nonlinear_fn (callable) – A function that implements the nonlinear transformation

• add_noise (bool) – Whether to add noise after the nonlinear operation

• avg_noise_power (float, optional) – The average noise power if add_noise is True

• snr_db (float, optional) – SNR in dB (alternative to avg_noise_power)

• complex_mode (str, optional) – How to handle complex inputs: ‘direct’ (default) passes the complex signal directly to nonlinear_fn, ‘cartesian’ applies the function separately to real and imaginary parts, ‘polar’ applies to magnitude and preserves phase

Example

>>> # Create a channel with cubic nonlinearity for real signals
>>> channel = NonlinearChannel(lambda x: x + 0.2 * x**3)
>>> x = torch.linspace(-1, 1, 100)
>>> y = channel(x)  # Output with cubic distortion

>>> # For complex signals, using polar mode (apply nonlinearity to magnitude only)
>>> def mag_distortion(x): return x * (1 - 0.1 * x)  # compression
>>> channel = NonlinearChannel(mag_distortion, complex_mode='polar')
>>> x = torch.complex(torch.randn(100), torch.randn(100))
>>> y = channel(x)  # Output with magnitude distortion, phase preserved

Methods

__init__	Initialize the Nonlinear channel.
forward	Apply the nonlinear function and optional noise to the input signal.
get_config	Get a dictionary of the channel's configuration.

Attributes

avg_noise_power
snr_db

get_config() → Dict[str, Any]

Get a dictionary of the channel’s configuration.

This method returns a dictionary containing the channel’s parameters, which can be used to recreate the channel instance.

Returns:

Dictionary of parameter names and values

Return type:

Dict[str, Any]

__init__(nonlinear_fn, add_noise=False, avg_noise_power: float | None = None, snr_db: float | None = None, complex_mode='direct', *args: Any, **kwargs: Any)

Initialize the Nonlinear channel.

Parameters:

• nonlinear_fn (callable) – The nonlinear transformation function.

• add_noise (bool) – Whether to add noise after nonlinearity.

• avg_noise_power (float, optional) – Average noise power if add_noise=True.

• snr_db (float, optional) – SNR in dB (alternative if add_noise=True).

• complex_mode (str) – How to handle complex inputs (‘direct’, ‘cartesian’, ‘polar’).

• *args – Variable length argument list passed to the base class.

• **kwargs – Arbitrary keyword arguments passed to the base class.

avg_noise_power: float | None

snr_db: float | None

forward(x: Tensor, *args: Any, **kwargs: Any) → Tensor

Apply the nonlinear function and optional noise to the input signal.

Parameters:

• x (torch.Tensor) – The input tensor.

• *args – Additional positional arguments (unused).

• **kwargs – Additional keyword arguments (unused).

Returns:

The output tensor after applying nonlinearity and noise.

Return type:

torch.Tensor