commkit 1.0.0__py3-none-any.whl

commkit/mapping/bits.py

@@ -0,0 +1,240 @@
+"""
+Hard bit mapping and demapping.
+
+Maps a flat bit sequence to Gray-coded symbols (``map_bits``) and recovers the
+most-likely bits from received symbols via minimum-distance decoding
+(``demap_symbols_hard``).
+"""
+
+import numpy as np
+
+from ..backend import ArrayType, dispatch
+from ..core.signal import Signal
+from ..logger import logger
+from .gray import gray_code, gray_constellation
+from .shaping import constellation_power
+
+__all__ = ["demap_symbols_hard", "map_bits"]
+
+
+def map_bits(
+    bits: ArrayType,
+    modulation: str,
+    order: int,
+    unipolar: bool = False,
+) -> ArrayType:
+    """
+    Maps a bit sequence to complex or real symbols.
+
+    This function follows a bit-first architecture: it takes a flat sequence
+    of bits and packs them into symbols according to the modulation scheme
+    and Gray mapping.
+
+    Output dtype is ``complex64`` for PSK/QAM and ``float32`` for ASK/PAM.
+
+    Parameters
+    ----------
+    bits : array_like
+        Input binary sequence (0s and 1s).
+    modulation : {"psk", "qam", "ask", "pam"}
+        Modulation scheme.
+    order : int
+        Modulation order (number of symbols).
+    unipolar : bool, default False
+        Trigger unipolar mapping for ASK/PAM.
+
+    Returns
+    -------
+    array_like
+        Array of symbols. Shape: (N_bits / log2(order),).
+        Backend (NumPy/CuPy) matches the input `bits`.
+    """
+    logger.debug("Mapping bits to %s %s-level symbols.", modulation.upper(), order)
+    bits, xp, _ = dispatch(bits)
+
+    if order < 2:
+        raise ValueError("Order must be at least 2 for modulation")
+
+    k = int(np.log2(order))
+    if 2**k != order:
+        raise ValueError(f"Order must be a power of 2, got {order}")
+
+    if len(bits) % k != 0:
+        raise ValueError(
+            f"Number of bits ({len(bits)}) must be divisible by bits per symbol ({k})"
+        )
+
+    # Reshape bits into symbols (N/k, k)
+    num_symbols = len(bits) // k
+
+    # Pack bits into integer indices
+    bits_reshaped = bits.reshape((num_symbols, k))
+    powers = 2 ** xp.arange(k - 1, -1, -1, dtype=int)
+    indices = xp.sum(bits_reshaped * powers, axis=1)
+
+    # Get constellation (returns NumPy)
+    constellation = gray_constellation(modulation, order, unipolar=unipolar)
+
+    # Ensure constellation is on the same backend and dtype
+    constellation = xp.asarray(constellation)
+
+    # ASK/PAM constellations are real-valued; PSK/QAM are complex.
+    mod_lower = modulation.lower()
+
+    if "ask" in mod_lower or "pam" in mod_lower:
+        constellation = constellation.astype(xp.float32)
+    else:
+        constellation = constellation.astype(xp.complex64)
+
+    # Map indices to points
+    return constellation[indices]
+
+
+def demap_symbols_hard(
+    symbols: ArrayType | Signal,
+    modulation: str | None = None,
+    order: int | None = None,
+    unipolar: bool = False,
+    pmf: np.ndarray | None = None,
+) -> ArrayType | Signal:
+    """
+    Maps complex symbols back to a sequence of bits (hard decisions).
+
+    This function performs minimum Euclidean distance decoding to find the
+    most likely bit pattern for each received symbol.
+
+    Parameters
+    ----------
+    symbols : array_like
+        Input array of received symbols. Shape: (..., N_symbols).
+    modulation : {"psk", "qam", "ask"}
+        Modulation type.
+    order : int
+        Modulation order.
+    unipolar : bool, default False
+        Trigger unipolar demapping for ASK/PAM.
+    pmf : np.ndarray, optional
+        Symbol PMF of shape ``(order,)`` for PS-QAM.  When provided, the
+        input symbols are scaled by ``sqrt(E_PS)`` (where
+        ``E_PS = Σ P(s_m) |s_m|²`` on the normalised grid) before the
+        nearest-neighbour search, mapping unit-avg-power resolved symbols
+        back to the ``{s_m}`` grid used by ``gray_constellation``.
+        Use this when ``symbols`` comes from
+        ``resolved_symbols`` of a PS-QAM signal.
+        Has no effect for uniform modulations.
+
+    Returns
+    -------
+    array_like
+        Sequence of bits (0s and 1s). Shape: (..., N_symbols * log2(order)).
+        The backend (NumPy/CuPy) matches the input `symbols`.
+    """
+    if isinstance(symbols, Signal):
+        sig = symbols
+        if sig.signal_type is not None:
+            logger.warning(
+                "demap_symbols_hard() called on a frame-generated signal - skipping. "
+                "Extract the payload segment via frame.get_structure_map() and build "
+                "a plain Signal before demapping."
+            )
+            return sig.copy()
+        if sig.mod_scheme is None or sig.mod_order is None:
+            raise ValueError("Modulation scheme and order required for demapping.")
+        if sig.resolved_symbols is None:
+            raise ValueError(
+                "No resolved symbols available. Call resolve_symbols(sig) first."
+            )
+        new = sig.copy()
+        new.resolved_bits = demap_symbols_hard(
+            sig.resolved_symbols,
+            sig.mod_scheme,
+            sig.mod_order,
+            unipolar=sig.mod_unipolar or False,
+            pmf=sig.ps_pmf if pmf is None else pmf,
+        )
+        return new
+
+    if modulation is None or order is None:
+        raise ValueError("demap_symbols_hard() requires modulation and order.")
+
+    logger.debug("Demapping %s %s-level symbols to bits.", modulation.upper(), order)
+    symbols, xp, _ = dispatch(symbols)
+
+    # Capture original shape to restore structure later
+    # If input is (N,), output (N*k,)
+    # If input is (C, N), output (C, N*k)
+    original_shape = symbols.shape
+
+    # Ensure 1D for vectorized distance broadcasting
+    symbols_flat = symbols.flatten()
+
+    # Get constellation
+    constellation = gray_constellation(modulation, order, unipolar=unipolar)
+
+    # PS-QAM: receive-path symbols at unit average power live on the
+    # ``{s_m/sqrt(E_PS)}`` grid.  Rescale by ``sqrt(E_PS)`` to bring them back
+    # to the ``{s_m}`` grid that ``gray_constellation`` returns, so the
+    # nearest-neighbour search is exact.  No-op for uniform modulations.
+    if pmf is not None:
+        e_ps = constellation_power(constellation, pmf)
+        if e_ps < 1.0 - 1e-6:
+            symbols_flat = symbols_flat * xp.asarray(
+                np.sqrt(e_ps), dtype=symbols_flat.real.dtype
+            )
+
+    constellation = xp.asarray(constellation)
+
+    # 1. Find nearest constellation point (Hard Decision)
+    k = int(np.log2(order))
+    is_sq_qam = (modulation == "qam") and (order != 8) and (k % 2 == 0)
+
+    if is_sq_qam:
+        # O(1) component-wise rounding - no (N, M) distance matrix.
+        # Levels are evenly spaced; round each axis independently.
+        n_ax = k // 2  # bits per axis
+        side = 2**n_ax  # points per axis
+        levels = xp.sort(xp.unique(constellation.real))  # (side,)
+        lev_min = float(levels[0])
+        d_grid = float(levels[1] - levels[0])
+        # Gray LUT: sorted-level index (geometric) -> natural-binary symbol index
+        gray_lut = xp.asarray(gray_code(n_ax))  # (side,)
+        g_i = xp.clip(
+            xp.round((symbols_flat.real - lev_min) / d_grid).astype(xp.int64),
+            0,
+            side - 1,
+        )
+        g_q = xp.clip(
+            xp.round((symbols_flat.imag - lev_min) / d_grid).astype(xp.int64),
+            0,
+            side - 1,
+        )
+        indices = (gray_lut[g_i] << n_ax) | gray_lut[g_q]  # (N_flat,)
+    else:
+        # General path: chunk N to bound peak memory at (CHUNK_N, M_const).
+        CHUNK_N = 4096
+        indices = xp.empty(len(symbols_flat), dtype=xp.int64)
+        for n0 in range(0, len(symbols_flat), CHUNK_N):
+            n1 = min(n0 + CHUNK_N, len(symbols_flat))
+            d = xp.abs(symbols_flat[n0:n1, xp.newaxis] - constellation[xp.newaxis, :])
+            indices[n0:n1] = xp.argmin(d, axis=1)
+
+    # 2. Convert indices to bits
+    # Extract bits from indices: (N, k)
+    # We use bit shifting: (index >> shift) & 1
+    shifts = xp.arange(k - 1, -1, -1, dtype="int32")
+    bits = ((indices[:, xp.newaxis] >> shifts) & 1).astype(xp.int8)
+
+    # 3. Reshape to restore original structure
+    # bits is currently (Total_Symbols, k)
+    # We flatten it to (Total_Symbols * k)
+    flat_bits = bits.flatten()
+
+    # Calculate new shape: replace last dimension D with D*k
+    # If original shape was (D1, D2, ..., Dn), new shape is (D1, D2, ..., Dn * k)
+    if len(original_shape) > 0:
+        new_shape = list(original_shape)
+        new_shape[-1] = new_shape[-1] * k
+        return flat_bits.reshape(new_shape)
+    else:
+        # Scalar input case (if supported), returns 1D array of bits
+        return flat_bits

commkit/mapping/constellation.py

@@ -0,0 +1,153 @@
+"""
+The :class:`Constellation` value object.
+
+A small immutable container bundling the three things that are otherwise
+passed around as loose tuples/arrays at every CMA/RDE radius and GMI/MI call
+site: the constellation ``points``, their natural-binary ``bit_labels``, and an
+optional probabilistic-shaping ``pmf``.  ``Constellation.gray(...)`` is the
+canonical constructor (delegating to :func:`gray_constellation`), and
+``power()`` / ``map()`` / ``demap()`` / ``llr()`` are thin wrappers over the
+existing free functions so a single object carries geometry + labels + prior
+together.
+
+The loose-array forms (``gray_constellation``, ``map_bits``,
+``demap_symbols_hard``, ``compute_llr``, ``constellation_power``) remain the
+public surface; this object is an additive convenience layered on top.
+"""
+
+from dataclasses import dataclass, replace
+from functools import lru_cache
+
+import numpy as np
+
+from ..backend import ArrayType
+from .bits import demap_symbols_hard, map_bits
+from .gray import gray_constellation
+from .llr import compute_llr
+from .shaping import constellation_power
+
+__all__ = ["Constellation"]
+
+
+@dataclass(frozen=True, eq=False)
+class Constellation:
+    """Immutable constellation: points, Gray bit-labels, and optional PS pmf.
+
+    Attributes
+    ----------
+    points : np.ndarray
+        Constellation points, shape ``(M,)``, indexed by natural-binary symbol
+        value (``points[i]`` is the symbol for integer ``i``).  Complex for
+        PSK/QAM, real for ASK/PAM.
+    bit_labels : np.ndarray
+        Natural-binary bit pattern per symbol index, shape ``(M, k)`` where
+        ``k = log2(M)``, dtype ``int8`` (MSB first).
+    modulation : str
+        Modulation scheme (``"psk"``, ``"qam"``, ``"ask"``, ``"pam"``).
+    order : int
+        Modulation order ``M``.
+    unipolar : bool
+        Whether the grid is the unipolar ASK/PAM variant.
+    pmf : np.ndarray or None
+        Optional Maxwell-Boltzmann shaping prior, shape ``(M,)``, aligned with
+        ``points``.  ``None`` means uniform.
+    """
+
+    points: np.ndarray
+    bit_labels: np.ndarray
+    modulation: str
+    order: int
+    unipolar: bool = False
+    pmf: np.ndarray | None = None
+
+    @property
+    def bits_per_symbol(self) -> int:
+        """Number of bits per symbol, ``log2(order)``."""
+        return int(self.bit_labels.shape[1])
+
+    @classmethod
+    def gray(
+        cls,
+        modulation: str,
+        order: int,
+        *,
+        normalize: bool = True,
+        unipolar: bool = False,
+        pmf: np.ndarray | None = None,
+    ) -> "Constellation":
+        """Build a Gray-mapped constellation (cached for the ``pmf=None`` base).
+
+        Parameters mirror :func:`gray_constellation`.  ``pmf`` attaches a
+        probabilistic-shaping prior without rebuilding the geometry.
+        """
+        base = _gray_base(modulation, order, normalize, unipolar)
+        if pmf is None:
+            return base
+        return replace(base, pmf=np.asarray(pmf, dtype=np.float64))
+
+    def power(self) -> float:
+        """pmf-weighted average symbol power ``E[|s|^2]`` (uniform if no pmf)."""
+        return constellation_power(self.points, self.pmf)
+
+    def map(self, bits: ArrayType) -> ArrayType:
+        """Map a flat bit sequence to symbols (see :func:`map_bits`)."""
+        return map_bits(bits, self.modulation, self.order, unipolar=self.unipolar)
+
+    def demap(self, symbols: ArrayType) -> ArrayType:
+        """Hard-decision demap symbols to bits (see :func:`demap_symbols_hard`).
+
+        Carries this constellation's ``pmf`` through for PS-QAM rescaling.
+        """
+        return demap_symbols_hard(
+            symbols,
+            self.modulation,
+            self.order,
+            unipolar=self.unipolar,
+            pmf=self.pmf,
+        )
+
+    def llr(
+        self,
+        symbols: ArrayType,
+        noise_var: float,
+        *,
+        method: str = "maxlog",
+        output: str = "jax",
+    ) -> ArrayType:
+        """Soft-decision LLRs (see :func:`compute_llr`), carrying this ``pmf``."""
+        return compute_llr(
+            symbols,
+            self.modulation,
+            self.order,
+            noise_var,
+            method=method,
+            unipolar=self.unipolar,
+            output=output,
+            pmf=self.pmf,
+        )
+
+
+@lru_cache(maxsize=128)
+def _gray_base(
+    modulation: str, order: int, normalize: bool, unipolar: bool
+) -> Constellation:
+    """Cached Gray constellation (geometry + labels), without a shaping pmf."""
+    points = gray_constellation(
+        modulation, order, normalize=normalize, unipolar=unipolar
+    )
+    k = int(np.log2(order))
+    bit_labels = (
+        (
+            np.arange(order, dtype="int32")[:, None]
+            >> np.arange(k - 1, -1, -1, dtype="int32")
+        )
+        & 1
+    ).astype(np.int8)
+    return Constellation(
+        points=points,
+        bit_labels=bit_labels,
+        modulation=modulation,
+        order=order,
+        unipolar=unipolar,
+        pmf=None,
+    )