commkit 1.0.0__py3-none-any.whl

commkit/metrics.py

@@ -0,0 +1,949 @@
+"""
+Performance metrics for digital communications.
+
+This module provides high-performance routines for evaluating signal quality,
+error rates, and reliability across different computational backends.
+"""
+
+import logging
+
+import numpy as np
+
+from .backend import ArrayType, dispatch, to_device
+from .core.signal import Signal
+from .logger import logger
+
+# Modulation-aware frame signals (those generated by ``SingleCarrierFrame``)
+# carry mixed preamble/pilot/payload segments with potentially different
+# modulations or gains; whole-signal metrics on them are meaningless.  The
+# Signal dispatch branches below warn and return ``None`` for such inputs,
+# matching the previous ``Signal.evm``/``snr``/``ber``/``ser`` behaviour.
+
+
+def _ps_unit_power_rescale(rx, xp, modulation, order, pmf, noise_var):
+    """Rescale unit-avg-power PS-QAM symbols ``{c·s_m}`` back to ``{s_m}``.
+
+    ``resolve_symbols`` normalises the receive constellation to unit average
+    power, placing PS-QAM symbols on ``{s_m / sqrt(E_PS)}`` (``c = 1/sqrt(E_PS)``).
+    The MI/LLR machinery expects symbols on the ``gray_constellation`` grid
+    ``{s_m}``.  This helper applies the exact deterministic correction -
+    ``rx -> rx·sqrt(E_PS)`` and ``noise_var -> noise_var·E_PS`` - and is a no-op
+    for uniform modulations (``pmf is None``) or when ``E_PS ≈ 1``.
+
+    Returns ``(rx_scaled, noise_var_scaled)``.
+    """
+    if pmf is None:
+        return rx, noise_var
+    from .mapping import Constellation
+
+    e_ps = Constellation.gray(modulation, order, pmf=pmf).power()
+    if e_ps < 1.0 - 1e-6:
+        rx = rx * xp.asarray(np.sqrt(e_ps), dtype=rx.real.dtype)
+        noise_var = noise_var * e_ps
+    return rx, noise_var
+
+
+def evm(
+    rx_symbols: ArrayType | Signal,
+    tx_symbols: ArrayType | None = None,
+    *,
+    num_train_symbols: int | None = None,
+    mode: str = "data_aided",
+    modulation: str | None = None,
+    order: int | None = None,
+    pmf: np.ndarray | None = None,
+) -> tuple[float | ArrayType, float | ArrayType] | None:
+    """
+    Computes Error Vector Magnitude (EVM).
+
+    Two modes are supported:
+
+    * ``"data_aided"`` *(default)* - uses known transmitted symbols as the
+      reference (classical simulation metric, most accurate).
+    * ``"blind"`` - decision-directed: the nearest ideal constellation point
+      to each received sample is used as the reference.  No knowledge of the
+      transmitted sequence is required.  Equivalent to measuring the *blob
+      size* around each constellation cluster, as a spectrum analyser would.
+      At very low SNR the estimate is optimistically biased because erroneous
+      decisions pull the centroid toward the wrong cluster.  The caller is
+      responsible for passing a gain-corrected signal; no implicit gain
+      normalisation is applied (``gray_constellation`` always returns
+      unit-power constellations).
+
+    Calculation is performed per-channel (independent gain normalisation) to
+    properly handle MIMO/multichannel signals.
+
+    Parameters
+    ----------
+    rx_symbols : array_like
+        Received symbols. Shape: ``(..., N_symbols)``.
+    tx_symbols : array_like, optional
+        Ideal reference (transmitted) symbols. Shape: ``(..., N_symbols)``.
+        Required when ``mode="data_aided"``, ignored when ``mode="blind"``.
+    mode : {"data_aided", "blind"}, default "data_aided"
+        Estimation mode (keyword-only).
+    modulation : str, optional
+        Modulation type (``"psk"``, ``"qam"``, ``"ask"``).
+        Required when ``mode="blind"``.
+    order : int, optional
+        Modulation order *M*. Required when ``mode="blind"``.
+    pmf : np.ndarray, optional
+        Symbol PMF of shape ``(order,)`` for PS-QAM in ``mode="blind"``.
+        When provided, ``rx_symbols`` are rescaled by ``sqrt(E_PS)`` before
+        the nearest-neighbour search so unit-avg-power resolved symbols
+        line up with the ``{s_m}`` grid returned by
+        ``gray_constellation``.  Has no effect for uniform modulations
+        or for ``mode="data_aided"`` (which is scale-invariant by
+        per-channel power normalisation).
+
+    Returns
+    -------
+    evm_percent : float or ndarray
+        EVM as a percentage. Scalar for SISO, array ``(N_ch,)`` for MIMO.
+    evm_db : float or ndarray
+        EVM in decibels (dB).
+
+    Raises
+    ------
+    ValueError
+        If ``mode="data_aided"`` and ``tx_symbols`` is ``None``, or
+        ``mode="blind"`` and ``modulation``/``order`` are missing, or
+        an unknown mode string is given.
+
+    Notes
+    -----
+    For MIMO, metrics are calculated independently for each stream (row).
+
+    When ``rx_symbols`` is a :class:`Signal`, ``resolved_symbols`` is used as the
+    received sequence, the reference falls back to ``source_symbols`` (and
+    ``modulation``/``order``/``pmf`` to the signal's metadata), and
+    ``num_train_symbols`` leading symbols are discarded.  Frame-generated
+    signals return ``None`` with a warning.
+    """
+    if isinstance(rx_symbols, Signal):
+        sig = rx_symbols
+        if sig.signal_type is not None:
+            logger.warning(
+                "evm() called on a frame-generated signal. Extract the payload "
+                "(or pilot) segment via frame.get_structure_map() and build a "
+                "plain Signal first."
+            )
+            return None
+        if sig.resolved_symbols is None:
+            raise ValueError(
+                "No resolved symbols available. Call resolve_symbols(sig) first."
+            )
+        if mode == "blind":
+            return evm(
+                sig.resolved_symbols,
+                num_train_symbols=num_train_symbols,
+                mode="blind",
+                modulation=modulation or sig.mod_scheme,
+                order=order or sig.mod_order,
+                pmf=sig.ps_pmf,
+            )
+        ref = tx_symbols if tx_symbols is not None else sig.source_symbols
+        if ref is None:
+            raise ValueError(
+                "No reference available. Provide tx_symbols or ensure "
+                "source_symbols is set on the Signal."
+            )
+        return evm(
+            sig.resolved_symbols, ref, num_train_symbols=num_train_symbols, pmf=pmf
+        )
+
+    from . import helpers
+
+    _VALID_MODES = ("data_aided", "blind")
+    if mode not in _VALID_MODES:
+        raise ValueError(f"Unknown mode '{mode}'. Choose from {_VALID_MODES}.")
+
+    rx, xp, _ = dispatch(rx_symbols)
+    axis = -1
+
+    # Discard leading training symbols (e.g. equalizer warm-up).  Available to
+    # array callers directly; the Signal branch forwards num_train_symbols here.
+    trim = num_train_symbols or 0
+    if trim > 0:
+        if mode == "blind" or tx_symbols is None:
+            rx = rx[..., trim:]
+        else:
+            tx_symbols = xp.asarray(tx_symbols)
+            n = min(rx.shape[-1], tx_symbols.shape[-1])
+            rx = rx[..., trim:n]
+            tx_symbols = tx_symbols[..., trim:n]
+
+    def _is_normalized(arr, ax):
+        pwr = xp.mean(xp.abs(arr) ** 2, axis=ax)
+        return xp.allclose(pwr, 1.0, atol=1e-3)
+
+    # --- Build reference ---
+    if mode == "data_aided":
+        if tx_symbols is None:
+            raise ValueError(
+                "mode='data_aided' requires tx_symbols. "
+                "Pass tx_symbols or use mode='blind'."
+            )
+        tx = xp.asarray(tx_symbols)
+        if rx.shape != tx.shape:
+            raise ValueError(f"Shape mismatch: rx {rx.shape} != tx {tx.shape}")
+
+        if not _is_normalized(rx, axis):
+            rx = helpers.normalize(rx, axis=axis, mode="average_power")
+        if not _is_normalized(tx, axis):
+            tx = helpers.normalize(tx, axis=axis, mode="average_power")
+
+    else:  # blind
+        if modulation is None or order is None:
+            raise ValueError(
+                "mode='blind' requires modulation and order. "
+                "Example: evm(rx, mode='blind', modulation='qam', order=16)."
+            )
+        from .mapping import Constellation
+
+        # gray_constellation always returns unit-average-power constellations.
+        # No gain correction of rx is applied here - the caller is responsible
+        # for passing a gain-corrected signal at the expected constellation power.
+        c = Constellation.gray(modulation, order, pmf=pmf)
+        constellation_np = c.points
+        constellation = xp.asarray(constellation_np)  # (M,) unit-avg-power
+
+        # PS-QAM: receive-path symbols at unit average power live on the
+        # ``{s_m/sqrt(E_PS)}`` grid.  Rescale rx by ``sqrt(E_PS)`` so the
+        # nearest-neighbour decision against ``{s_m}`` is exact.
+        if pmf is not None:
+            e_ps = c.power()
+            if e_ps < 1.0 - 1e-6:
+                rx = rx * xp.asarray(np.sqrt(e_ps), dtype=rx.real.dtype)
+
+        # ML hard decision: nearest constellation point per symbol.
+        # rx shape (..., N) -> (..., N, 1) vs (M,) -> (..., N, M)
+        dist = xp.abs(rx[..., None] - constellation) ** 2
+        tx = constellation[xp.argmin(dist, axis=-1)]  # same shape as rx
+
+    # --- EVM computation (shared) ---
+    ref_pwr = xp.mean(xp.abs(tx) ** 2, axis=axis)
+    low_pwr_mask = ref_pwr < 1e-20
+    if xp.any(low_pwr_mask):
+        logger.warning("Reference signal power near zero in one or more channels.")
+
+    error = rx - tx
+    error_power = xp.mean(xp.abs(error) ** 2, axis=axis)
+    evm_ratio = xp.sqrt(error_power)  # ref power is 1.0 after normalisation
+
+    if evm_ratio.ndim == 0:
+        if low_pwr_mask:
+            return float("inf"), float("inf")
+        evm_percent = float(evm_ratio * 100.0)
+        with np.errstate(divide="ignore"):
+            evm_db = (
+                float(20.0 * xp.log10(evm_ratio)) if evm_ratio > 0 else float("-inf")
+            )
+        logger.info("EVM [%s]: %.2f%% (%.2f dB)", mode, evm_percent, evm_db)
+        return evm_percent, evm_db
+
+    evm_percent = evm_ratio * 100.0
+    with np.errstate(divide="ignore"):
+        evm_db = 20.0 * xp.log10(evm_ratio)
+
+    evm_percent[low_pwr_mask] = float("inf")
+    evm_db[low_pwr_mask] = float("inf")
+
+    if logger.isEnabledFor(logging.INFO):
+        # One batched D2H per array (and none at all when INFO is off).
+        evm_pct_np = to_device(evm_percent, "cpu")
+        evm_db_np = to_device(evm_db, "cpu")
+        for ch in range(evm_pct_np.shape[0]):
+            logger.info(
+                "EVM [%s] Ch%s: %.2f%% (%.2f dB)",
+                mode,
+                ch,
+                float(evm_pct_np[ch]),
+                float(evm_db_np[ch]),
+            )
+
+    return evm_percent, evm_db
+
+
+def snr(
+    rx_symbols: ArrayType | Signal,
+    tx_symbols: ArrayType | None = None,
+    *,
+    num_train_symbols: int | None = None,
+) -> float | ArrayType | None:
+    """
+    Estimates SNR from received symbols using a known reference (Data-Aided).
+
+    SNR is computed as the ratio of average signal power to noise (error)
+    power:
+    SNR_linear ≈ 1.0 / E[|rx - tx|^2]
+
+    Calculation is performed per-channel (independent gain normalization).
+
+    Parameters
+    ----------
+    rx_symbols : array_like or Signal
+        Received noisy symbols. Shape: (..., N_symbols).
+    tx_symbols : array_like or Signal
+        Known transmitted symbols. Shape: (..., N_symbols).
+
+    Returns
+    -------
+    float or ndarray
+        Estimated SNR in dB. Returns array if input is multichannel.
+
+    When ``rx_symbols`` is a :class:`Signal`, ``resolved_symbols`` is used, the
+    reference falls back to ``source_symbols``, and ``num_train_symbols`` leading
+    symbols are discarded.  Frame-generated signals return ``None``.
+    """
+    if isinstance(rx_symbols, Signal):
+        sig = rx_symbols
+        if sig.signal_type is not None:
+            logger.warning(
+                "snr() called on a frame-generated signal. Extract the payload "
+                "segment via frame.get_structure_map() and build a plain Signal first."
+            )
+            return None
+        ref = tx_symbols if tx_symbols is not None else sig.source_symbols
+        if ref is None:
+            raise ValueError(
+                "No reference available. Provide tx_symbols or ensure "
+                "source_symbols is set on the Signal."
+            )
+        if sig.resolved_symbols is None:
+            raise ValueError(
+                "No resolved symbols available. Call resolve_symbols(sig) first."
+            )
+        return snr(sig.resolved_symbols, ref, num_train_symbols=num_train_symbols)
+
+    from . import helpers
+
+    if tx_symbols is None:
+        raise ValueError("snr() requires tx_symbols for array input.")
+
+    rx, xp, _ = dispatch(rx_symbols)
+    tx = xp.asarray(tx_symbols)
+
+    axis = -1
+
+    # Discard leading training symbols, clamping to the shorter sequence.
+    trim = num_train_symbols or 0
+    if trim > 0:
+        n = min(rx.shape[-1], tx.shape[-1])
+        rx = rx[..., trim:n]
+        tx = tx[..., trim:n]
+
+    if rx.shape != tx.shape:
+        raise ValueError(f"Shape mismatch: rx {rx.shape} != tx {tx.shape}")
+
+    def _is_normalized(arr, ax):
+        pwr = xp.mean(xp.abs(arr) ** 2, axis=ax)
+        return xp.allclose(pwr, 1.0, atol=1e-3)
+
+    # Normalize per channel if needed
+    if not _is_normalized(rx, axis):
+        rx = helpers.normalize(rx, axis=axis, mode="average_power")
+
+    if not _is_normalized(tx, axis):
+        tx = helpers.normalize(tx, axis=axis, mode="average_power")
+
+    # Check for zero power ref for SNR too
+    ref_pwr = xp.mean(xp.abs(tx) ** 2, axis=axis)
+    low_pwr_mask = ref_pwr < 1e-20
+
+    # Signal power (reference) is 1.0 per channel effectively, BUT
+    # if original ref was zero, SNR is 0 (-inf dB).
+    signal_power = 1.0
+
+    # Noise power (error)
+    error = rx - tx
+    noise_power = xp.mean(xp.abs(error) ** 2, axis=axis)
+
+    # Calculate SNR
+    # Handle zeros in noise_power
+    with np.errstate(divide="ignore"):
+        snr_linear = signal_power / noise_power
+        snr_db = 10.0 * xp.log10(snr_linear)
+
+    # Handle zero reference case
+    # If noise is zero, SNR is Inf (perfect signal), even if ref was zero.
+    # If noise is present but ref was zero, SNR is -Inf (0 linear).
+    zero_noise_mask = noise_power < 1e-20
+
+    if snr_db.ndim == 0:
+        if zero_noise_mask:
+            return float("inf")
+        if low_pwr_mask:
+            return float("-inf")
+        snr_val = float(snr_db)
+        logger.info("SNR: %.2f dB", snr_val)
+        return snr_val
+
+    # Array case: Overwrite with -inf where (low_pwr_mask AND NOT zero_noise_mask)
+    # Note: where zero_noise_mask is True, snr_db is already 'inf' from division
+    mask_to_neg_inf = low_pwr_mask & (~zero_noise_mask)
+    snr_db[mask_to_neg_inf] = float("-inf")
+
+    if logger.isEnabledFor(logging.INFO):
+        snr_db_np = to_device(snr_db, "cpu")
+        for ch in range(snr_db_np.shape[0]):
+            logger.info("SNR Ch%s: %.2f dB", ch, float(snr_db_np[ch]))
+
+    return snr_db
+
+
+def ber(
+    bits_rx: ArrayType | Signal,
+    bits_tx: ArrayType | None = None,
+    *,
+    num_train_symbols: int | None = None,
+) -> float | ArrayType | None:
+    """
+    Computes the Bit Error Rate (BER) between two bit sequences.
+
+    Calculation is performed per-channel (along ``axis=-1``) to properly
+    handle MIMO/multichannel signals.
+
+    Parameters
+    ----------
+    bits_rx : array_like
+        Received bit sequence (binary 0/1). Shape: (..., N).
+    bits_tx : array_like
+        Original transmitted bit sequence (binary 0/1). Shape: (..., N).
+
+    Returns
+    -------
+    float or ndarray
+        BER as a ratio in [0, 1]. Scalar for 1D input, array for multichannel.
+
+    Notes
+    -----
+    The calculation is backend-agnostic and supports comparison of massive
+    bit sequences on both CPU and GPU.
+
+    When ``bits_rx`` is a :class:`Signal`, ``resolved_bits`` is used, the
+    reference falls back to ``source_bits``, and ``num_train_symbols`` (in
+    *symbols*, converted to bits via ``bits_per_symbol``) leading bits are
+    discarded.  Frame-generated signals return ``None``.
+    """
+    if isinstance(bits_rx, Signal):
+        sig = bits_rx
+        if sig.signal_type is not None:
+            logger.warning(
+                "ber() called on a frame-generated signal. Extract the payload "
+                "segment via frame.get_structure_map() and build a plain Signal first."
+            )
+            return None
+        ref = bits_tx if bits_tx is not None else sig.source_bits
+        if ref is None:
+            raise ValueError(
+                "No reference bits available. Provide bits_tx or ensure "
+                "source_bits is set on the Signal."
+            )
+        if sig.resolved_bits is None:
+            raise ValueError(
+                "No resolved bits available. Call demap_symbols_hard(sig) first."
+            )
+        y = sig.resolved_bits
+        r = ref
+        trim = num_train_symbols or 0
+        if trim > 0 and sig.mod_order is not None:
+            bps = sig.bits_per_symbol
+            assert bps is not None
+            bit_trim = trim * bps
+            n = min(y.shape[-1], r.shape[-1])
+            y = y[..., bit_trim:n]
+            r = r[..., bit_trim:n]
+        return ber(y, r)
+
+    if bits_tx is None:
+        raise ValueError("ber() requires bits_tx for array input.")
+
+    bits_rx, xp, _ = dispatch(bits_rx)
+    bits_tx = xp.asarray(bits_tx)
+
+    if bits_rx.shape != bits_tx.shape:
+        raise ValueError(f"Shape mismatch: rx {bits_rx.shape} != tx {bits_tx.shape}")
+
+    if bits_rx.size == 0:
+        return 0.0
+
+    # Per-stream error count along time axis
+    errors = xp.sum(bits_rx != bits_tx, axis=-1)
+    total = bits_rx.shape[-1]
+    ber_values = errors / total
+
+    if ber_values.ndim == 0:
+        ber_value = float(ber_values)
+        logger.info("BER: %.2e (%s/%s errors)", ber_value, int(errors), total)
+        return ber_value
+
+    if logger.isEnabledFor(logging.INFO):
+        ber_np = to_device(ber_values, "cpu")
+        errors_np = to_device(errors, "cpu")
+        for ch in range(ber_np.shape[0]):
+            logger.info(
+                "BER Ch%s: %.2e (%s/%s errors)",
+                ch,
+                float(ber_np[ch]),
+                int(errors_np[ch]),
+                total,
+            )
+    return ber_values
+
+
+def ser(
+    rx_symbols: ArrayType | Signal,
+    tx_symbols: ArrayType | None = None,
+    modulation: str | None = None,
+    order: int | None = None,
+    *,
+    num_train_symbols: int | None = None,
+    pmf: np.ndarray | None = None,
+) -> float | ArrayType | None:
+    """
+    Computes the Symbol Error Rate (SER) using ML hard decisions.
+
+    Each received symbol is decided to its nearest constellation point
+    (minimum Euclidean distance).  The decision is then compared to the
+    corresponding transmitted symbol.  SER measures symbol-level errors
+    before bit demapping and does not assume Gray coding, making it
+    distinct from BER.
+
+    Computation is fully vectorised and backend-agnostic (NumPy / CuPy).
+
+    Parameters
+    ----------
+    rx_symbols : array_like
+        Received (noisy) symbols. Shape: ``(..., N_symbols)``.
+    tx_symbols : array_like
+        Transmitted (ideal) symbols. Shape: ``(..., N_symbols)``.
+    modulation : str
+        Modulation type: ``"psk"``, ``"qam"``, or ``"ask"``.
+    order : int
+        Modulation order *M*.
+    pmf : np.ndarray, optional
+        Symbol PMF of shape ``(order,)`` for PS-QAM.  When provided, the
+        comparison constellation is scaled by ``1/sqrt(E_PS)`` so the
+        nearest-neighbour search matches the scale of unit-avg-power
+        received symbols.  Both ``rx_symbols`` and ``tx_symbols`` are
+        decided against the same scaled constellation.  Has no effect for
+        uniform modulations.
+
+    Returns
+    -------
+    float or ndarray
+        SER as a ratio in ``[0, 1]``. Scalar for SISO input,
+        array of shape ``(N_ch,)`` for MIMO input.
+
+    Notes
+    -----
+    Both ``rx_symbols`` and ``tx_symbols`` are decided against the same
+    constellation so the function is robust to any floating-point rounding
+    in the reference values.
+
+    When ``rx_symbols`` is a :class:`Signal`, ``resolved_symbols`` is used, the
+    reference falls back to ``source_symbols``, ``modulation``/``order``/``pmf``
+    fall back to the signal's metadata, and ``num_train_symbols`` leading
+    symbols are discarded.  Frame-generated signals return ``None``.
+    """
+    if isinstance(rx_symbols, Signal):
+        sig = rx_symbols
+        if sig.signal_type is not None:
+            logger.warning(
+                "ser() called on a frame-generated signal. Extract the payload "
+                "segment via frame.get_structure_map() and build a plain Signal first."
+            )
+            return None
+        ref = tx_symbols if tx_symbols is not None else sig.source_symbols
+        if ref is None:
+            raise ValueError(
+                "No reference available. Provide tx_symbols or ensure "
+                "source_symbols is set on the Signal."
+            )
+        if sig.resolved_symbols is None:
+            raise ValueError(
+                "No resolved symbols available. Call resolve_symbols(sig) first."
+            )
+        mod = modulation or sig.mod_scheme
+        ord_ = order or sig.mod_order
+        if mod is None or ord_ is None:
+            raise ValueError(
+                "SER requires modulation and order. Pass them explicitly or ensure "
+                "mod_scheme/mod_order are set on the Signal."
+            )
+        return ser(
+            sig.resolved_symbols,
+            ref,
+            mod,
+            ord_,
+            num_train_symbols=num_train_symbols,
+            pmf=sig.ps_pmf if pmf is None else pmf,
+        )
+
+    if modulation is None or order is None:
+        raise ValueError("ser() requires modulation and order for array input.")
+    if tx_symbols is None:
+        raise ValueError("ser() requires tx_symbols for array input.")
+
+    from .mapping import Constellation
+
+    rx, xp, _ = dispatch(rx_symbols)
+    tx = xp.asarray(tx_symbols)
+
+    trim = num_train_symbols or 0
+    if trim > 0:
+        n = min(rx.shape[-1], tx.shape[-1])
+        rx = rx[..., trim:n]
+        tx = tx[..., trim:n]
+
+    if rx.shape != tx.shape:
+        raise ValueError(f"Shape mismatch: rx {rx.shape} != tx {tx.shape}")
+
+    c = Constellation.gray(modulation, order, pmf=pmf)
+    constellation_np = c.points
+    constellation = xp.asarray(constellation_np)  # (M,)
+
+    # PS-QAM: ``rx_symbols`` from ``resolved_symbols``
+    # are normalised to unit average power, placing them on the
+    # ``{s_m/sqrt(E_PS)}`` grid.  ``tx_symbols`` (``source_symbols``) live on
+    # the un-rescaled ``{s_m}`` grid (their average power is ``E_PS < 1``).
+    # Rescale ``rx`` by ``sqrt(E_PS)`` so the nearest-neighbour search against
+    # ``gray_constellation`` is correct for both rx and tx.  Has no effect on
+    # uniform modulations.
+    if pmf is not None:
+        e_ps = c.power()
+        if e_ps < 1.0 - 1e-6:
+            rx = rx * xp.asarray(np.sqrt(e_ps), dtype=rx.real.dtype)
+
+    # Broadcast: rx/tx (..., N) -> (..., N, 1) vs constellation (M,) -> (..., N, M)
+    dist_rx = xp.abs(rx[..., None] - constellation) ** 2  # (..., N, M)
+    dist_tx = xp.abs(tx[..., None] - constellation) ** 2  # (..., N, M)
+
+    dec_rx = xp.argmin(dist_rx, axis=-1)  # (..., N) - index into constellation
+    dec_tx = xp.argmin(dist_tx, axis=-1)  # (..., N)
+
+    errors = xp.sum(dec_rx != dec_tx, axis=-1)
+    total = rx.shape[-1]
+    ser_values = errors / total
+
+    if ser_values.ndim == 0:
+        ser_val = float(ser_values)
+        logger.info("SER: %.2e (%s/%s errors)", ser_val, int(errors), total)
+        return ser_val
+
+    if logger.isEnabledFor(logging.INFO):
+        ser_np = to_device(ser_values, "cpu")
+        errors_np = to_device(errors, "cpu")
+        for ch in range(ser_np.shape[0]):
+            logger.info(
+                "SER Ch%s: %.2e (%s/%s errors)",
+                ch,
+                float(ser_np[ch]),
+                int(errors_np[ch]),
+                total,
+            )
+    return ser_values
+
+
+def gmi(
+    llrs: ArrayType | Signal,
+    tx_bits: ArrayType | None = None,
+    *,
+    noise_var: float | None = None,
+    modulation: str | None = None,
+    order: int | None = None,
+    pmf: np.ndarray | None = None,
+    method: str = "maxlog",
+) -> float:
+    r"""
+    Computes the Generalised Mutual Information (GMI) from per-bit LLRs.
+
+    GMI is the achievable information rate under bit-interleaved coded
+    modulation (BICM) and is the standard capacity metric for coded optical /
+    wireless systems.  It uses pre-computed LLRs so the modulation format
+    does not need to be known by this function.
+
+    GMI = sum_{b=0..k-1} ( 1 - (1/N) * sum_{n=1..N} log2(1 + exp(-LLR_b[n] * (1 - 2 * c_b[n]))) )
+
+    where LLR_b[n] is the LLR for bit b of symbol n, and c_b[n] (0 or 1) is
+    the transmitted bit.
+
+    The per-term softplus is computed in a numerically stable form:
+
+    log2(1 + exp(-x)) = (1 / ln 2) * (log(1 + exp(-|x|)) + max(0, -x))
+
+    which avoids overflow for large |x|.
+
+    Parameters
+    ----------
+    llrs : array_like
+        Per-bit LLRs. Accepted shapes:
+
+        * ``(N * k,)`` - flat, all bits concatenated.
+        * ``(N, k)`` - symbols along axis 0, bits along axis 1.
+
+        where *N* is the number of symbols and *k = log₂(M)* bits per symbol.
+        Positive LLR -> bit 0 more likely; negative -> bit 1 more likely.
+        LLRs from ``compute_llr`` match this convention.
+    tx_bits : array_like
+        Transmitted bits (0/1 integers), same shape as ``llrs``.
+
+    Returns
+    -------
+    float
+        GMI in bits per channel use (b/cu). In the range ``[0, log₂(M)]``;
+        equals ``log₂(M)`` at infinite SNR and approaches 0 at very low SNR.
+
+    Examples
+    --------
+    >>> llrs = compute_llr(rx_symbols, "qam", 16, noise_var=0.1, output='numpy')
+    >>> gmi_value = gmi(llrs, tx_bits)
+    >>> gmi_value = gmi(sig, noise_var=0.1)        # Signal input: LLRs computed
+
+    When the first argument is a :class:`Signal`, LLRs are computed internally
+    from ``resolved_symbols`` (applying the PS-QAM scale correction, see
+    :func:`_ps_unit_power_rescale`) via ``compute_llr`` with ``noise_var``/
+    ``method``/``pmf``, and GMI is evaluated against ``source_bits``.
+    """
+    if isinstance(llrs, Signal):
+        sig = llrs
+        if sig.resolved_symbols is None:
+            raise ValueError(
+                "No resolved symbols available. Call resolve_symbols(sig) first."
+            )
+        if sig.source_bits is None:
+            raise ValueError(
+                "GMI requires source_bits. Ensure the Signal was created via a "
+                "factory (e.g. generate_qam(), generate_psqam())."
+            )
+        mod = modulation or sig.mod_scheme
+        ord_ = order or sig.mod_order
+        if mod is None or ord_ is None:
+            raise ValueError(
+                "GMI requires modulation and order. Ensure mod_scheme/mod_order "
+                "are set on the Signal."
+            )
+        if noise_var is None:
+            raise ValueError("gmi() requires noise_var.")
+
+        from .mapping import compute_llr
+
+        eff_pmf = sig.ps_pmf if pmf is None else pmf
+        rx, xp_, _ = dispatch(sig.resolved_symbols)
+        resolved, adj_noise_var = _ps_unit_power_rescale(
+            rx, xp_, mod, ord_, eff_pmf, noise_var
+        )
+        computed = compute_llr(
+            resolved,
+            mod,
+            ord_,
+            adj_noise_var,
+            method=method,
+            pmf=eff_pmf,
+            output="numpy",
+        )
+        src_bits = to_device(sig.source_bits, "cpu")
+        # compute_llr outputs (N*k,) / (C, N*k); reshape to (N_total, k) so the
+        # array core infers bits-per-symbol from the trailing axis.
+        k_bits = int(np.log2(ord_))
+        llrs_arr = np.asarray(computed)
+        bits_arr = np.asarray(src_bits)
+        n_bits = llrs_arr.size
+        return gmi(
+            llrs_arr.reshape(n_bits // k_bits, k_bits),
+            bits_arr.reshape(n_bits // k_bits, k_bits),
+        )
+
+    if tx_bits is None:
+        raise ValueError("gmi() requires tx_bits for array (LLR) input.")
+
+    llrs_arr, xp, _ = dispatch(llrs)
+    llrs_arr = llrs_arr.astype(xp.float64)
+    bits_arr = xp.asarray(tx_bits, dtype=xp.float64)
+
+    # Detect k (bits per symbol) from 2D shape before flattening.
+    # For (N, k) input: sum over k bit positions -> GMI ∈ [0, k].
+    # For 1D (N·k,) input: k=1, returns per-bit GMI ∈ [0, 1].
+    k = llrs_arr.shape[1] if llrs_arr.ndim == 2 else 1
+
+    llrs_flat = llrs_arr.ravel()
+    bits_flat = bits_arr.ravel()
+
+    if llrs_flat.shape != bits_flat.shape:
+        raise ValueError(
+            f"llrs and tx_bits must have the same number of elements. "
+            f"Got {llrs_flat.shape} vs {bits_flat.shape}."
+        )
+
+    # Signed LLR: positive when transmitted bit is 0, negative when bit is 1.
+    # The argument to log2(1 + exp(·)) is: -LLR * (1 - 2*bit)
+    # = -LLR when bit=0  (should be positive -> small loss)
+    # = +LLR when bit=1  (should be negative -> small loss at high SNR)
+    x = -llrs_flat * (1.0 - 2.0 * bits_flat)
+
+    # Numerically stable softplus: log2(1 + exp(x)) = log1p(exp(-|x|))/ln2 + max(0,x)/ln2
+    ln2 = xp.log(xp.asarray(2.0, dtype=xp.float64))
+    softplus = (
+        xp.log1p(xp.exp(-xp.abs(x))) + xp.maximum(xp.asarray(0.0, dtype=xp.float64), x)
+    ) / ln2
+
+    # GMI = Σ_{b=0}^{k-1} (1 - mean_n[softplus_b]) = k * (1 - mean_all[softplus])
+    gmi_value = float(k * (1.0 - xp.mean(softplus)))
+    logger.info("GMI: %.4f b/cu", gmi_value)
+    return gmi_value
+
+
+def mi(
+    symbols_rx: ArrayType | Signal,
+    modulation: str | None = None,
+    order: int | None = None,
+    noise_var: float | None = None,
+    *,
+    pmf: np.ndarray | None = None,
+    rescale_unit_power: bool = False,
+) -> float:
+    r"""
+    Estimates the Mutual Information (MI) under a Gaussian channel assumption.
+
+    MI is the Shannon capacity of the discrete-input continuous-output (DCMC)
+    AWGN channel.  It upper-bounds GMI and is achievable with optimal
+    (non-binary) coding.
+
+    Uses the Monte-Carlo estimator with Gaussian likelihoods:
+
+    p(s_m | r) is proportional to exp(-|r - s_m|^2 / sigma^2)
+
+    MI = log2(M) + (1/N) * sum_{n=1..N} sum_{m=1..M} ( p(s_m | r_n) * log2(p(s_m | r_n)) )
+
+    Computed via the log-sum-exp trick for numerical stability.
+
+    Parameters
+    ----------
+    symbols_rx : array_like
+        Received noisy symbols. Shape: ``(N,)``.
+    modulation : str
+        Modulation type: ``'psk'``, ``'qam'``, or ``'ask'``.
+    order : int
+        Modulation order M.
+    noise_var : float
+        Complex noise variance sigma^2 of the AWGN model
+        **referenced to the normalised constellation** (unit average power
+        under the uniform distribution, i.e. the same scale as
+        ``gray_constellation``).
+        For unit-power symbols at Es/N0 (dB):
+        sigma^2 = 10^(-Es/N0 / 10).
+    pmf : np.ndarray, optional
+        Symbol PMF of shape ``(M,)`` for PS-QAM. When provided, the
+        non-uniform prior ``P(sₘ)`` is incorporated into the likelihood
+        computation and the result is clipped to ``[0, H(X)]`` instead of
+        ``[0, log₂M]``. Pass ``signal.ps_pmf`` for PS-QAM signals.
+        ``None`` (default) assumes uniform prior.
+    rescale_unit_power : bool, default False
+        When ``True`` and ``pmf`` is given, rescale ``symbols_rx`` from the
+        unit-average-power grid ``{c·s_m}`` back to ``{s_m}`` (and
+        ``noise_var`` by ``E_PS``) before estimation - required when the input
+        comes from ``resolve_symbols``.  See :func:`_ps_unit_power_rescale`.
+        The Signal dispatch path sets this automatically.
+
+    Returns
+    -------
+    float
+        MI in bits per channel use (b/cu). In the range ``[0, H(X)]``
+        where ``H(X) = log₂M`` for uniform and ``H(X) < log₂M`` for PS.
+
+    Warning - PS-QAM scale convention: ``symbols_rx`` must be on the same
+    scale as ``gray_constellation`` (unit average power under the uniform
+    distribution). Samples from ``apply_awgn`` are already on the correct
+    scale. However, after ``resolve_symbols`` the receiver re-normalises,
+    shifting PS symbols from {s_m} to {c * s_m} (c = 1 / sqrt(E_PS)).  Either
+    pass a :class:`Signal` (the scale correction is applied automatically) or
+    set ``rescale_unit_power=True`` when passing ``resolved_symbols`` arrays.
+
+    Examples
+    --------
+    >>> mi_value = mi(rx_symbols, "qam", 64, noise_var=0.05)
+    >>> mi_value = mi(sig, noise_var=0.05)              # Signal input
+    """
+    if isinstance(symbols_rx, Signal):
+        sig = symbols_rx
+        if sig.resolved_symbols is None:
+            raise ValueError(
+                "No resolved symbols available. Call resolve_symbols(sig) first."
+            )
+        mod = modulation or sig.mod_scheme
+        ord_ = order or sig.mod_order
+        if mod is None or ord_ is None:
+            raise ValueError(
+                "MI requires modulation and order. Ensure mod_scheme/mod_order "
+                "are set on the Signal."
+            )
+        if noise_var is None:
+            raise ValueError("mi() requires noise_var.")
+        return mi(
+            sig.resolved_symbols,
+            mod,
+            ord_,
+            noise_var,
+            pmf=sig.ps_pmf if pmf is None else pmf,
+            rescale_unit_power=True,
+        )
+
+    if modulation is None or order is None:
+        raise ValueError("mi() requires modulation and order for array input.")
+    if noise_var is None:
+        raise ValueError("mi() requires noise_var.")
+
+    from .mapping import gray_constellation
+
+    rx, xp, _ = dispatch(symbols_rx)
+    rx = rx.ravel().astype(xp.complex128)  # (N,)
+    if rescale_unit_power:
+        rx, noise_var = _ps_unit_power_rescale(
+            rx, xp, modulation, order, pmf, noise_var
+        )
+    constellation = xp.asarray(
+        gray_constellation(modulation, order), dtype=xp.complex128
+    )  # (M,)
+    M = len(constellation)
+    ln2 = float(xp.log(xp.asarray(2.0, dtype=xp.float64)))
+    log2_M = float(xp.log2(xp.asarray(float(M), dtype=xp.float64)))
+
+    # Prior: uniform or PS-shaped
+    if pmf is not None:
+        pmf_arr = np.asarray(pmf, dtype=np.float64).clip(1e-300, None)
+        log_prior = xp.asarray(np.log(pmf_arr), dtype=xp.float64)  # (M,) nats
+        nz = pmf_arr > 0
+        h_x_bits = float(-np.sum(pmf_arr[nz] * np.log2(pmf_arr[nz])))
+    else:
+        log_prior = xp.full(
+            M, -xp.log(xp.asarray(float(M), dtype=xp.float64)), dtype=xp.float64
+        )  # log(1/M)
+        h_x_bits = log2_M
+
+    # Log-joint: log P(sₘ) + log p(r | sₘ)   (drop noise-normalisation constant)
+    # Chunk over N to bound the (chunk, M) complex128/float64 intermediates -
+    # a full-N allocation is ~4 GB at N=1e6, M=256.  The per-row math is
+    # unchanged; chunk sums accumulate on device (single sync at the end).
+    N_sym = rx.shape[0]
+    CHUNK_N = 65536
+    acc = xp.zeros((), dtype=xp.float64)
+    for n0 in range(0, N_sym, CHUNK_N):
+        rx_c = rx[n0 : n0 + CHUNK_N]
+        diff = rx_c[:, None] - constellation[None, :]  # (chunk, M)
+        log_liks = log_prior[None, :] - (diff.real**2 + diff.imag**2) / noise_var
+
+        # log p(s_m | r) via log-sum-exp for numerical stability
+        lse_shift = log_liks.max(axis=1, keepdims=True)  # (chunk, 1)
+        log_sum = xp.log(xp.sum(xp.exp(log_liks - lse_shift), axis=1)) + lse_shift[:, 0]
+
+        log_posterior = log_liks - log_sum[:, None]  # (chunk, M), log base e
+        posterior = xp.exp(log_posterior)  # (chunk, M), sum-to-1 per row
+
+        # MI = H(X) + (1/N) Σ_n Σ_m p(s_m|r_n) · log2 p(s_m|r_n)
+        acc += xp.sum(posterior * log_posterior)
+    mi_nats = float(acc) / N_sym
+    mi_value = float(np.clip(h_x_bits + mi_nats / ln2, 0.0, h_x_bits))
+
+    logger.info("MI: %.4f b/cu  (max %.2f b/cu)", mi_value, h_x_bits)
+    return mi_value