commkit 1.0.0__py3-none-any.whl

commkit/mapping/gray.py

@@ -0,0 +1,429 @@
+"""
+Constellation geometry and Gray labelling.
+
+Gray code sequences, their inverse, and the Gray-mapped constellation
+generators for PSK / ASK / square-, 8-rect-, and cross-QAM.  These are the
+geometric foundation consumed by the bit mapping, soft demapping, and
+probabilistic-shaping layers.
+
+Note: codes and constellations are generated using NumPy (host-side).
+"""
+
+from functools import lru_cache
+
+import numpy as np
+
+from ..logger import logger
+
+__all__ = [
+    "gray_code",
+    "gray_constellation",
+    "gray_to_binary",
+]
+
+
+@lru_cache(maxsize=128)
+def gray_code(n: int) -> np.ndarray:
+    """
+    Generates a Gray code sequence for `n` bits.
+
+    Gray coding ensures that adjacent symbols in the constellation differ
+    by only one bit, minimizing bit error rate (BER) for a given symbol
+    error rate (SER).
+
+    Parameters
+    ----------
+    n : int
+        Number of bits.
+
+    Returns
+    -------
+    np.ndarray
+        Array of integers representing the Gray code sequence.
+        Shape: (2^n,).
+
+    Examples
+    --------
+    >>> gray_code(2)
+    array([0, 1, 3, 2])
+    """
+    if n < 0:
+        raise ValueError("n must be non-negative")
+    if n == 0:
+        return np.array([0], dtype=int)
+
+    # Vectorized implementation: s ^ (s >> 1)
+    i = np.arange(1 << n, dtype=int)
+
+    return i ^ (i >> 1)
+
+
+@lru_cache(maxsize=128)
+def gray_to_binary(n: int) -> np.ndarray:
+    """
+    Computes the inverse Gray code mapping: symbol index to bit pattern.
+
+    In soft demapping, it is necessary to identify which bits correspond to
+    each constellation point. Since `gray_code(n)[i]` provides the Gray
+    code for natural binary `i`, this inverse mapping finds the natural
+    binary representation corresponding to a specific Gray-coded index.
+
+    Parameters
+    ----------
+    n : int
+        Number of bits per symbol.
+
+    Returns
+    -------
+    np.ndarray
+        Array where `result[s]` provides the natural binary integer
+        representing the bit pattern for symbol `s`. Shape: (2^n,).
+    """
+    if n < 0:
+        raise ValueError("n must be non-negative")
+    if n == 0:
+        return np.array([0], dtype=int)
+
+    # gray_code gives: natural_binary -> gray_symbol
+    # We want: gray_symbol -> natural_binary (inverse)
+    gray = gray_code(n)
+    inverse = np.zeros(1 << n, dtype=int)
+    inverse[gray] = np.arange(1 << n)
+
+    return inverse
+
+
+@lru_cache(maxsize=128)
+def gray_constellation(
+    modulation: str,
+    order: int,
+    normalize: bool = True,
+    unipolar: bool = False,
+) -> np.ndarray:
+    """
+    Generate constellation points with Gray mapping.
+
+    The returned array is indexed by the symbol's natural binary value.
+    `constellation[i]` is the complex/float value for the symbol representing
+    the bit pattern for integer `i`.
+
+    Parameters
+    ----------
+    modulation : {"psk", "qam", "ask", "pam"}
+        Modulation type. If string contains 'unipolar', unipolar=True is triggered.
+    order : int
+        Modulation order (number of symbols). Must be a power of 2.
+    normalize : bool, default True
+        If True, scales the constellation to unit average power (E_s = 1).
+    unipolar : bool, default False
+        If True, shifts the constellation to be strictly non-negative (for 'ask'/'pam').
+        If the modulation string contains 'unipol', this is automatically forced to True.
+
+    Returns
+    -------
+    np.ndarray
+        Array of constellation points (NumPy). Shape: (order,).
+        Complex for 'psk' and 'qam', Float for 'ask'.
+
+    Notes
+    -----
+    By default, constellations are normalized to **unit average power (E_s = 1)**.
+    This ensures that when different modulation schemes are mixed (e.g., PSK
+    pilots with QAM payload), they start from a consistent power baseline.
+    """
+    logger.debug(
+        "Generating Gray-coded constellation: modulation=%s, order=%s, normalize=%s",
+        modulation,
+        order,
+        normalize,
+    )
+    modulation = modulation.lower()
+
+    # Extract core modulation scheme
+    if "psk" in modulation:
+        modulation = "psk"
+    elif "qam" in modulation:
+        modulation = "qam"
+    elif "ask" in modulation or "pam" in modulation:
+        modulation = "ask"
+
+    if order < 2:
+        raise ValueError("Order must be at least 2 for modulation")
+
+    if modulation == "psk":
+        result = _gray_psk(order)
+    elif modulation == "ask":
+        result = _gray_ask(order, unipolar=unipolar)
+    elif modulation == "qam":
+        k = int(np.log2(order))
+        if 2**k != order:
+            raise ValueError(f"Order must be power of 2 for {modulation}")
+
+        # Check for 8-QAM
+        if order == 8:
+            result = _gray_qam_8_rect()
+        elif k % 2 == 0:
+            result = _gray_qam_square(order)
+        else:
+            result = _gray_qam_cross(order)
+    else:
+        raise ValueError(f"Unsupported modulation type: {modulation}")
+
+    if normalize:
+        if modulation == "psk":
+            # PSK is already on unit circle (E_s = 1)
+            pass
+        else:
+            from .. import helpers
+
+            result = helpers.normalize(result, mode="average_power")
+
+    return result
+
+
+def _gray_psk(order: int) -> np.ndarray:
+    """
+    Generates an M-PSK constellation with Gray mapping.
+
+    Parameters
+    ----------
+    order : int
+        Modulation order (must be a power of 2).
+
+    Returns
+    -------
+    np.ndarray
+        Complex array of PSK constellation points.
+    """
+    # Bits per symbol
+    k = int(np.log2(order))
+    if 2**k != order:
+        raise ValueError("Order must be power of 2 for psk")
+
+    # Gray codes for k bits
+    gray = gray_code(k)
+
+    # BPSK Special Case: strictly real values
+    if order == 2:
+        points = np.array([-1.0, 1.0], dtype=float)
+        constellation = np.zeros(order, dtype=float)
+        constellation[gray] = points
+        return constellation
+
+    # Geometric points (phases)
+    # Phases are 0, 2pi/M, ..., 2pi(M-1)/M
+    phases = np.arange(order) * 2 * np.pi / order
+    points = np.exp(1j * phases)
+
+    constellation = np.zeros(order, dtype=complex)
+    constellation[gray] = points
+
+    return constellation
+
+
+def _gray_ask(order: int, unipolar: bool = False) -> np.ndarray:
+    """
+    Generates an M-ASK constellation with Gray mapping.
+
+    By default, returns bipolar values centered at zero. If unipolar=True,
+    shifts values to start from zero.
+
+    Parameters
+    ----------
+    order : int
+        Modulation order (must be a power of 2).
+
+    Returns
+    -------
+    np.ndarray
+        Real-valued array of ASK constellation points.
+    """
+    k = int(np.log2(order))
+    if 2**k != order:
+        raise ValueError("Order must be power of 2 for ask")
+
+    gray = gray_code(k)
+    # Points are symmetric: -M+1, -M+3, ..., M-3, M-1
+    points = np.linspace(-order + 1, order - 1, order)
+
+    if unipolar:
+        points = points - np.min(points)
+
+    constellation = np.zeros(order, dtype=float)
+    constellation[gray] = points
+
+    return constellation
+
+
+def _gray_qam_square(order: int) -> np.ndarray:
+    """
+    Generates a Square M-QAM constellation with Gray mapping.
+
+    Constructed as the Cartesian product of two Gray-coded ASK
+    constellations. Valid for orders where log2(M) is even.
+
+    Parameters
+    ----------
+    order : int
+        Modulation order (must be an even power of 2, e.g., 16, 64).
+
+    Returns
+    -------
+    np.ndarray
+        Complex array of square QAM constellation points.
+    """
+    # M = 2^(2k). I bits = k, Q bits = k.
+    k_total = int(np.log2(order))
+    n = k_total // 2  # bits per axis
+
+    # Generate ASK for each axis
+    m_axis = 2**n
+    pam = _gray_ask(m_axis)
+
+    # Cartesian product
+    # Symbol s (2n bits). High n bits -> I, Low n bits -> Q.
+    s = np.arange(order)
+    mask_q = (1 << n) - 1
+
+    idx_i = s >> n
+    idx_q = s & mask_q
+
+    i_vals = pam[idx_i]
+    q_vals = pam[idx_q]
+
+    constellation = i_vals + 1j * q_vals
+
+    return constellation
+
+
+def _gray_qam_8_rect() -> np.ndarray:
+    """
+    Generates an 8-QAM constellation with Rectangular Gray mapping.
+
+    Returns
+    -------
+    np.ndarray
+        Complex array of 8-QAM constellation points.
+    """
+
+    points = np.zeros(8, dtype=complex)
+
+    # 000 (I=-3, Q=-1)
+    points[0] = -3.0 - 1.0j
+
+    # 001 (I=-3, Q=+1)
+    points[1] = -3.0 + 1.0j
+
+    # 010 (I=-1, Q=-1)
+    points[2] = -1.0 - 1.0j
+
+    # 011 (I=-1, Q=+1)
+    points[3] = -1.0 + 1.0j
+
+    # 100 (I=+3, Q=-1)
+    points[4] = 3.0 - 1.0j
+
+    # 101 (I=+3, Q=+1)
+    points[5] = 3.0 + 1.0j
+
+    # 110 (I=+1, Q=-1)
+    points[6] = 1.0 - 1.0j
+
+    # 111 (I=+1, Q=+1)
+    points[7] = 1.0 + 1.0j
+
+    return points
+
+
+def _gray_qam_cross(order: int) -> np.ndarray:
+    """
+    Generates a Cross-QAM constellation using a Quasi-Gray mapping.
+
+    Cross constellations are used for QAM orders where log2(M) is odd (e.g., 32, 128)
+    to maintain a more circular/compact shape than a purely rectangular grid.
+
+    Parameters
+    ----------
+    order : int
+        Modulation order (must be an odd power of 2).
+
+    Returns
+    -------
+    np.ndarray
+        Complex array of cross QAM constellation points.
+    """
+    # M = 2^(2k+1).
+    k_total = int(np.log2(order))
+
+    # I bits n = ceil(k_total/2) = (k+1)//2
+    # Q bits m = floor(k_total/2) = k // 2
+    n = (k_total + 1) // 2
+    m = k_total // 2
+
+    width = 2**n
+    height = 2**m
+
+    # 1. Generate Gray sequences for underlying Rectangular Grid
+    gray_i_seq = gray_code(n)
+    gray_q_seq = gray_code(m)
+
+    # 2. Build Inverse Lookup Tables (Symbol -> Geometric Index)
+    inv_gray_i = np.zeros(width, dtype=int)
+    inv_gray_i[gray_i_seq] = np.arange(width)
+
+    inv_gray_q = np.zeros(height, dtype=int)
+    inv_gray_q[gray_q_seq] = np.arange(height)
+
+    # 3. Decompose input Symbols into I and Q components
+    s = np.arange(order)
+    mask_q = (1 << m) - 1
+    idx_i_sym = s >> m
+    idx_q_sym = s & mask_q  # This is the Gray-coded Symbol Value for Q
+
+    # 4. Map Symbols to Geometric Indices (0..Width-1, 0..Height-1)
+    geo_i = inv_gray_i[idx_i_sym]
+    geo_q = inv_gray_q[idx_q_sym]
+
+    # 5. Determine Folding Parameters
+    n_shift = 0
+    if n >= 3:
+        n_shift = 2 ** (n - 3)
+
+    if n_shift == 0:
+        # Fallback
+        val_i = (-width + 1 + 2 * geo_i).astype(float)
+        val_q = (-height + 1 + 2 * geo_q).astype(float)
+        return val_i + 1j * val_q
+
+    # 6. Identify Wings
+    mask_left = geo_i < n_shift
+    mask_right = geo_i >= (width - n_shift)
+
+    # 7. Calculate Rotation/Translation
+    center_offset = (width - height) // 2
+
+    # Left Wing -> Top Cap
+    rotated_i_left = geo_q + center_offset
+    rotated_q_left = height + (n_shift - 1 - geo_i)
+
+    # Right Wing -> Bottom Cap
+    rotated_i_right = geo_q + center_offset
+    rotated_q_right = -1 - (geo_i - (width - n_shift))
+
+    # Apply Updates
+    geo_i_final = geo_i.copy()
+    geo_q_final = geo_q.copy()
+
+    geo_i_final = np.where(mask_left, rotated_i_left, geo_i_final)
+    geo_i_final = np.where(mask_right, rotated_i_right, geo_i_final)
+
+    geo_q_final = np.where(mask_left, rotated_q_left, geo_q_final)
+    geo_q_final = np.where(mask_right, rotated_q_right, geo_q_final)
+
+    # Convert Final Geometric Indices to Values
+    final_i_vals = (-width + 1 + 2 * geo_i_final).astype(float)
+    final_q_vals = (-height + 1 + 2 * geo_q_final).astype(float)
+
+    constellation = final_i_vals + 1j * final_q_vals
+
+    return constellation

commkit/mapping/llr.py

@@ -0,0 +1,253 @@
+"""
+Soft-decision demapping (log-likelihood ratios).
+
+JAX JIT-compiled max-log and exact LLR kernels, with optional probabilistic-
+shaping prior.  Sign convention: positive LLR -> bit 0 more likely.
+"""
+
+from typing import Any
+
+import numpy as np
+
+from ..backend import ArrayType, _get_jax, dispatch, is_jax_array, to_jax
+from ..logger import logger
+from .gray import gray_constellation
+
+__all__ = ["compute_llr"]
+
+# Lazy cache for JIT-compiled soft demapping kernels
+_JITTED_SOFT_DEMAP: dict[str, Any] = {}
+
+
+def _get_jitted_soft_demap():
+    """
+    Returns JIT-compiled maxlog and exact LLR computation functions.
+
+    Functions are defined and compiled lazily on first call to avoid
+    importing JAX at module load time.
+    """
+    if not _JITTED_SOFT_DEMAP:
+        jax, jnp, _ = _get_jax()
+        if jax is None:
+            raise ImportError(
+                "JAX is required for soft demapping. Install with: pip install jax"
+            )
+
+        @jax.jit
+        def maxlog(symbols, constellation, bits_table_t, sigma_sq, log_pmf):
+            """Max-log LLR with PS prior.
+
+            symbols (N,), constellation (M,), bits_table_t (k, M), log_pmf (M,).
+            Uniform case: pass log_pmf = jnp.zeros(M) - constant offset cancels.
+            PS case: log_pmf = log P(sₘ).
+
+            Effective metric: eff_m = d_m/σ² - log P(sₘ)
+            LLR_k = min_{b=1} eff - min_{b=0} eff
+            """
+            distances_sq = (
+                jnp.abs(symbols[:, None] - constellation[None, :]) ** 2
+            )  # (N, M)
+            eff = distances_sq / sigma_sq - log_pmf[None, :]  # (N, M)
+
+            def bit_llr(bit_row):  # (M,)
+                d0 = jnp.where(bit_row == 0, eff, jnp.inf)  # (N, M)
+                d1 = jnp.where(bit_row == 1, eff, jnp.inf)  # (N, M)
+                return jnp.min(d1, axis=1) - jnp.min(d0, axis=1)  # (N,)
+
+            return jax.vmap(bit_llr)(bits_table_t).T  # (k, N) -> (N, k)
+
+        @jax.jit
+        def exact(symbols, constellation, bits_table_t, sigma_sq, log_pmf):
+            """Exact LLR with PS prior via log-sum-exp.
+
+            symbols (N,), constellation (M,), bits_table_t (k, M), log_pmf (M,).
+            Uniform case: pass log_pmf = jnp.zeros(M).
+            PS case: log_pmf = log P(sₘ).
+
+            log_terms_m = log P(sₘ) - d_m/σ²
+            LLR_k = LSE_{b=0}(log_terms) - LSE_{b=1}(log_terms)
+            """
+            distances_sq = (
+                jnp.abs(symbols[:, None] - constellation[None, :]) ** 2
+            )  # (N, M)
+            log_terms = log_pmf[None, :] - distances_sq / sigma_sq  # (N, M)
+
+            def bit_llr(bit_row):  # (M,)
+                e0 = jnp.where(bit_row == 0, log_terms, -jnp.inf)  # (N, M)
+                e1 = jnp.where(bit_row == 1, log_terms, -jnp.inf)  # (N, M)
+                return jax.scipy.special.logsumexp(
+                    e0, axis=1
+                ) - jax.scipy.special.logsumexp(e1, axis=1)  # (N,)
+
+            return jax.vmap(bit_llr)(bits_table_t).T  # (k, N) -> (N, k)
+
+        _JITTED_SOFT_DEMAP["maxlog"] = maxlog
+        _JITTED_SOFT_DEMAP["exact"] = exact
+
+    return _JITTED_SOFT_DEMAP["maxlog"], _JITTED_SOFT_DEMAP["exact"]
+
+
+def compute_llr(
+    symbols: ArrayType,
+    modulation: str,
+    order: int,
+    noise_var: float,
+    method: str = "maxlog",
+    unipolar: bool = False,
+    output: str = "jax",
+    pmf: np.ndarray | None = None,
+) -> ArrayType:
+    """
+    Compute Log-Likelihood Ratios (LLRs) for soft-decision decoding.
+
+    Positive LLR -> bit 0 more likely; negative -> bit 1; magnitude = confidence.
+    JAX JIT-compiled with ``jax.vmap`` over bit positions; fully differentiable.
+
+    Parameters
+    ----------
+    symbols : array_like
+        Received noisy symbols. Shape: (..., N_symbols). NumPy, CuPy, or JAX.
+    modulation : {"psk", "qam", "ask"}
+        Modulation type.
+    order : int
+        Modulation order.
+    noise_var : float
+        Complex noise variance sigma^2 referenced to the normalised
+        constellation (unit avg power).  sigma^2 = 10^(-EsN0_dB / 10).
+    method : {"maxlog", "exact"}, default "maxlog"
+        LLR algorithm. ``"maxlog"`` is faster; ``"exact"`` uses log-sum-exp.
+    unipolar : bool, default False
+        Use unipolar constellation for ASK/PAM.
+    output : {"jax", "input", "numpy"}, default "jax"
+        Output array type.  ``"jax"`` preserves differentiability;
+        ``"input"`` matches the input backend; ``"numpy"`` forces NumPy.
+    pmf : np.ndarray, optional
+        Symbol PMF of shape ``(order,)`` for PS-QAM.  Pass
+        ``maxwell_boltzmann(order, nu)`` to incorporate the non-uniform prior.
+        ``None`` assumes uniform prior.
+
+    Returns
+    -------
+    array_like
+        LLR values. Shape: (..., N_symbols * log2(order)).
+        Array type determined by ``output``.
+
+    Notes
+    -----
+    Max-Log: LLR_k ≈ (1/sigma^2) * (min_{S_1^k} |r-s|^2 - min_{S_0^k} |r-s|^2).
+    Exact: LLR_k = log sum_{S_0^k} exp(-|r-s|^2/sigma^2) - log sum_{S_1^k} ...
+
+    For PS-QAM, ``symbols`` must be on the same scale as
+    ``gray_constellation`` (unit avg power).  After
+    ``resolve_symbols`` the receiver renormalises;
+    use ``gmi`` instead for correct scale.
+    """
+    logger.debug(
+        "Computing LLRs for %s %s-level (method=%s, output=%s).",
+        modulation.upper(),
+        order,
+        method,
+        output,
+    )
+
+    k = int(np.log2(order))
+    if 2**k != order:
+        raise ValueError(f"Order must be a power of 2, got {order}")
+    if method not in ("maxlog", "exact"):
+        raise ValueError(f"Unknown method: {method}. Use 'maxlog' or 'exact'.")
+    if output not in ("jax", "input", "numpy"):
+        raise ValueError(f"Unknown output: {output!r}. Use 'jax', 'input', or 'numpy'.")
+
+    # Convert to JAX if not already
+    jax, jnp, _ = _get_jax()
+    if jax is None:
+        raise ImportError(
+            "JAX is required for LLR computation. Install with: pip install jax"
+        )
+
+    # Build constellation and bits table on CPU once - shared by both paths.
+    is_complex = symbols.dtype.kind == "c"
+    const = gray_constellation(modulation, order, unipolar=unipolar).astype(
+        "complex64" if is_complex else "float32"
+    )
+    bits_table_np = (
+        (
+            (
+                np.arange(order, dtype="int32")[:, None]
+                >> np.arange(k - 1, -1, -1, dtype="int32")
+            )
+            & 1
+        )
+        .astype("int32")
+        .T
+    )
+    sigma_np = np.float32(max(noise_var, 1e-20))
+
+    # Build log_pmf: zeros = uniform (constant offset cancels in LLR difference).
+    if pmf is not None:
+        log_pmf_np = np.log(np.clip(np.asarray(pmf, dtype=np.float32), 1e-40, None))
+    else:
+        log_pmf_np = np.zeros(order, dtype=np.float32)
+
+    # JAX path
+    jax_module, jnp, _ = _get_jax()
+    if is_jax_array(symbols):
+        assert jnp is not None
+        if hasattr(symbols, "shape"):
+            original_shape = symbols.shape
+        else:
+            # Fallback for JAX tracers or odd objects
+            original_shape = jnp.shape(symbols)
+
+        jax_symbols_flat = symbols.flatten()
+        constellation_jax = jnp.asarray(const)
+        bits_table_t_jax = jnp.asarray(bits_table_np)
+        sigma_sq = jnp.asarray(sigma_np)
+        log_pmf_jax = jnp.asarray(log_pmf_np)
+
+    # NumPy/CuPy path
+    else:
+        symbols, xp, _ = dispatch(symbols)
+        original_shape = symbols.shape
+
+        jax_symbols_flat = to_jax(
+            symbols, dtype="complex64" if is_complex else "float32"
+        ).flatten()
+        device = jax_symbols_flat.device
+
+        assert jax_module is not None
+        assert jnp is not None
+        # device_put accepts NumPy arrays directly - no intermediate jnp.asarray needed
+        constellation_jax = jax_module.device_put(const, device)
+        bits_table_t_jax = jax_module.device_put(bits_table_np, device)
+        sigma_sq = jax_module.device_put(jnp.asarray(sigma_np), device)
+        log_pmf_jax = jax_module.device_put(log_pmf_np, device)
+
+    # Compute LLRs via JIT-compiled kernels
+    maxlog_fn, exact_fn = _get_jitted_soft_demap()
+    if method == "maxlog":
+        llrs = maxlog_fn(
+            jax_symbols_flat, constellation_jax, bits_table_t_jax, sigma_sq, log_pmf_jax
+        )
+    else:
+        llrs = exact_fn(
+            jax_symbols_flat, constellation_jax, bits_table_t_jax, sigma_sq, log_pmf_jax
+        )
+
+    # Reshape to match input structure
+    flat_llrs = llrs.flatten()
+    if len(original_shape) > 1:
+        new_shape = list(original_shape)
+        new_shape[-1] = new_shape[-1] * k
+        flat_llrs = flat_llrs.reshape(new_shape)
+
+    # Convert output to the requested backend
+    if output == "jax":
+        return flat_llrs
+    elif output == "numpy":
+        return np.asarray(flat_llrs)
+    else:  # output == "input"
+        if is_jax_array(symbols):
+            return flat_llrs  # already JAX
+        # xp is NumPy or CuPy - convert via NumPy intermediate
+        return xp.asarray(np.asarray(flat_llrs))