Utils#

JAX utility functions for signal processing.

Overview#

The Utils module provides JAX implementations of:

  • AWGN - Additive White Gaussian Noise utilities

  • Complex Operations - Complex number operations

API Reference#

Complex arithmetic operations on stacked real/imag tensors.

ran.phy.jax.utils.awgn(rng: jax.Array, H: jax.Array, snr_db: float) → jax.Array[source]#

Add AWGN to channel (JAX version).

Parameters:

  • rng – JAX PRNG key

  • H – Channel with shape (n_sc, n_sym, n_ant) complex

  • snr_db – SNR in dB

Returns:

Noisy channel

ran.phy.jax.utils.complex_mul(a__ri: jax.Array, b__ri: jax.Array) → jax.Array[source]#

Complex multiplication: a * b.

Computes (a_real + i*a_imag) * (b_real + i*b_imag) for tensors where the first dimension contains [real, imag] components.

Parameters:

  • a__ri – Complex tensor with shape (2, …)

  • b__ri – Complex tensor with shape (2, …)

Returns:

Complex product with shape (2, …)

Return type:

result__ri

Notes

Formula: (a+bi) * (c+di) = (ac-bd) + (ad+bc)i - result[0] = real part = a[0]*b[0] - a[1]*b[1] - result[1] = imag part = a[0]*b[1] + a[1]*b[0]

ran.phy.jax.utils.complex_mul_conj(a__ri: jax.Array, b__ri: jax.Array) → jax.Array[source]#

Complex multiplication with conjugate: a * conj(b).

Computes (a_real + i*a_imag) * (b_real - i*b_imag) for tensors where the first dimension contains [real, imag] components.

Parameters:

  • a__ri – Complex tensor with shape (2, …)

  • b__ri – Complex tensor with shape (2, …)

Returns:

Complex product with shape (2, …)

Return type:

result__ri

Notes

Formula: (a+bi) * (c-di) = (ac+bd) + (bc-ad)i - result[0] = real part = a[0]*b[0] + a[1]*b[1] - result[1] = imag part = a[1]*b[0] - a[0]*b[1]