commkit 1.0.0__py3-none-any.whl

commkit/impairments/source.py

@@ -0,0 +1,219 @@
+"""Optical/electronic source impairments (laser/oscillator phase noise)."""
+
+import math
+
+import numpy as np
+
+from ..backend import ArrayType, dispatch, is_cupy_available, to_device
+from ..logger import logger
+
+__all__ = ["apply_phase_noise", "generate_phase_noise"]
+
+
+def _phase_trajectory(
+    shape: tuple[int, int],
+    sampling_rate: float,
+    linewidth: float,
+    flicker: float,
+    flicker_f_min: float | None,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """
+    NumPy float64 phase trajectories with one-sided FM-noise PSD
+
+        S_f(f) = linewidth / pi  +  flicker / f      [Hz^2/Hz].
+
+    The white-FM part is generated exactly as a discrete Wiener walk
+    (per-sample increments N(0, 2*pi*linewidth/f_s)); the flicker part by
+    spectral shaping of white frequency noise.  Generated on the CPU so a
+    given seed yields the identical trajectory on every backend.
+    """
+    num_samples = shape[-1]
+    phi = np.zeros(shape, dtype=np.float64)
+
+    if linewidth > 0.0:
+        std = math.sqrt(2.0 * math.pi * linewidth / sampling_rate)
+        phi += np.cumsum(rng.normal(0.0, std, shape), axis=-1)
+
+    if flicker > 0.0:
+        f = np.fft.rfftfreq(num_samples, 1.0 / sampling_rate)
+        f_min = (
+            flicker_f_min if flicker_f_min is not None else sampling_rate / num_samples
+        )
+        # A unit-variance white input has one-sided PSD 2/f_s, so shaping to
+        # S_f = flicker/f requires the amplitude gain sqrt(flicker/f * f_s/2).
+        gain = np.sqrt(flicker / np.maximum(f, f_min)) * math.sqrt(sampling_rate / 2.0)
+        spec = np.fft.rfft(rng.normal(0.0, 1.0, shape), axis=-1)
+        df = np.fft.irfft(spec * gain, num_samples, axis=-1)
+        phi += 2.0 * math.pi * np.cumsum(df, axis=-1) / sampling_rate
+
+    return phi
+
+
+def generate_phase_noise(
+    num_samples: int,
+    sampling_rate: float,
+    linewidth: float = 0.0,
+    flicker: float = 0.0,
+    flicker_f_min: float | None = None,
+    num_streams: int = 1,
+    seed: int | None = None,
+) -> ArrayType:
+    """
+    Generates laser/oscillator phase-noise trajectories phi[n] in radians.
+
+    The instantaneous-frequency (FM) noise follows the standard power-law
+    model with a white and a flicker component:
+
+        S_f(f) = linewidth / pi  +  flicker / f      [Hz^2/Hz, one-sided]
+
+    * ``linewidth`` is the Lorentzian (white-FM / Wiener) linewidth
+      delta_nu: the phase performs a random walk with per-sample increments
+      N(0, 2*pi*delta_nu / f_s) and the field spectrum is a Lorentzian of
+      FWHM delta_nu.
+    * ``flicker`` is the 1/f FM coefficient h_-1: technical noise (current
+      source, temperature, acoustics) that dominates below the corner
+      frequency f_c = pi * h_-1 / delta_nu where the two terms cross.
+
+    Returning the trajectory itself (rather than a rotated signal) makes the
+    ground truth available for estimator validation; apply it with
+    ``samples * xp.exp(1j * phi)`` or via :func:`apply_phase_noise`.
+
+    Parameters
+    ----------
+    num_samples : int
+        Trajectory length per stream.
+    sampling_rate : float
+        Sampling rate in Hz.
+    linewidth : float, default 0.0
+        White-FM (Lorentzian) linewidth delta_nu in Hz.
+    flicker : float, default 0.0
+        Flicker-FM coefficient h_-1 in Hz^2 (one-sided ``S_f = h_-1 / f``).
+    flicker_f_min : float, optional
+        Frequency below which the flicker shaping is held flat (the 1/f
+        divergence must be capped).  Defaults to the record resolution
+        ``sampling_rate / num_samples``.
+    num_streams : int, default 1
+        Number of independent trajectories.
+    seed : int, optional
+        Random seed for reproducible trajectories.
+
+    Returns
+    -------
+    array_like
+        Phase in radians, ``float64``, on the active device (GPU when CuPy
+        is available).  Shape ``(num_samples,)`` for ``num_streams=1``,
+        else ``(num_streams, num_samples)``.
+
+    Notes
+    -----
+    The trajectory is always generated with NumPy's ``default_rng`` and then
+    transferred, so a given seed produces the identical trajectory on CPU
+    and GPU (same convention as :func:`~commkit.helpers.generate_bits`).
+    """
+    logger.info(
+        "Generating phase noise (linewidth=%.3g Hz, flicker=%.3g Hz², %s stream(s)).",
+        linewidth,
+        flicker,
+        num_streams,
+    )
+
+    rng = np.random.default_rng(seed)
+    phi = _phase_trajectory(
+        (num_streams, num_samples),
+        sampling_rate,
+        linewidth,
+        flicker,
+        flicker_f_min,
+        rng,
+    )
+    if num_streams == 1:
+        phi = phi[0]
+    if is_cupy_available():
+        phi = to_device(phi, "gpu")
+    return phi
+
+
+def apply_phase_noise(
+    samples: ArrayType,
+    sampling_rate: float,
+    linewidth: float,
+    flicker: float = 0.0,
+    flicker_f_min: float | None = None,
+    seed: int | None = None,
+    shared_lo: bool = False,
+) -> ArrayType:
+    """
+    Adds laser / oscillator phase noise to a signal.
+
+    Each sample is rotated by an accumulated phase drawn from the power-law
+    FM-noise model of :func:`generate_phase_noise` (white-FM Wiener walk
+    plus optional 1/f flicker):
+
+        r[n] = s[n] * exp(j * phi[n])
+
+    Parameters
+    ----------
+    samples : array_like
+        Complex baseband signal. Shape: ``(N,)`` (SISO) or ``(C, N)`` (MIMO).
+    sampling_rate : float
+        Sampling rate in Hz.
+    linewidth : float
+        Combined transmitter + receiver laser linewidth delta_nu in Hz.
+        Typical values: 100 kHz (narrow-linewidth laser) to 10 MHz (DFB).
+    flicker : float, default 0.0
+        Flicker-FM coefficient h_-1 in Hz^2 (one-sided ``S_f = h_-1 / f``).
+    flicker_f_min : float, optional
+        Low-frequency cap for the flicker shaping; see
+        :func:`generate_phase_noise`.
+    seed : int, optional
+        Random seed for reproducible noise.
+    shared_lo : bool, default False
+        When ``False`` (default), each channel receives independent phase noise
+        (separate oscillators / lasers per TX-RX path).
+        When ``True``, a single phase noise trajectory is shared across all
+        channels (common local oscillator in a coherent system).
+
+    Returns
+    -------
+    array_like
+        Phase-noise-impaired signal, same shape, dtype, and backend as input.
+
+    Examples
+    --------
+    >>> noisy = apply_phase_noise(sig.samples, linewidth=100e3,
+    ...                           sampling_rate=sig.sampling_rate)
+    """
+    logger.info(
+        "Applying phase noise (linewidth=%.3g Hz, flicker=%.3g Hz², shared_lo=%s).",
+        linewidth,
+        flicker,
+        shared_lo,
+    )
+
+    samples, xp, _ = dispatch(samples)
+    was_1d = samples.ndim == 1
+    if was_1d:
+        samples = samples[None, :]  # (1, N)
+    C, N = samples.shape
+
+    rng = np.random.default_rng(seed)
+    num_trajectories = 1 if shared_lo else C
+    phase = xp.asarray(
+        _phase_trajectory(
+            (num_trajectories, N),
+            sampling_rate,
+            linewidth,
+            flicker,
+            flicker_f_min,
+            rng,
+        )
+    )
+    result = samples * xp.exp(1j * phase)  # (1, N) broadcasts across channels
+
+    if result.dtype != samples.dtype:
+        result = result.astype(samples.dtype)
+
+    if was_1d:
+        return result[0]
+    return result

commkit/io.py

@@ -0,0 +1,308 @@
+"""
+Signal I/O and file management utilities.
+
+File layout
+-----------
+The .npz file contains the following named entries:
+
+  ``samples``              - IQ sample array  (always present)
+  ``source_bits``          - source bit array  (omitted if None)
+  ``source_symbols``       - source symbol array  (omitted if None)
+  ``resolved_symbols``     - cached symbol array  (only with include_cache=True)
+  ``resolved_bits``        - cached bit array     (only with include_cache=True)
+  ``__metadata__``         - zero-d object array holding a YAML string with all
+                             scalar fields.
+  ``__frame_metadata__``   - zero-d object array holding a YAML string with the
+                             serialised SingleCarrierFrame fields (omitted when
+                             the signal was not generated from a frame).
+  ``frame_payload_symbols`` - frame payload symbols array  (omitted if no frame)
+  ``frame_pilot_symbols``   - frame pilot symbols array    (omitted if no frame/pilots)
+  ``frame_payload_bits``    - frame payload bits array     (omitted if no frame)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import yaml
+
+from . import backend as _backend
+
+if TYPE_CHECKING:
+    from .core import Signal
+
+# -----------------------------------------------------------------------------
+# Internal constants
+# -----------------------------------------------------------------------------
+
+# Scalar / primitive metadata fields to round-trip through YAML
+_META_FIELDS: tuple[str, ...] = (
+    "sampling_rate",
+    "symbol_rate",
+    "signal_type",
+    "mod_scheme",
+    "mod_order",
+    "mod_unipolar",
+    "mod_rz",
+    "pulse_shape",
+    "filter_span",
+    "rrc_rolloff",
+    "rc_rolloff",
+    "duty_cycle",
+    "rise_time",
+    "spectral_domain",
+    "physical_domain",
+    "center_frequency",
+    "digital_frequency_offset",
+    "ps_nu",
+)
+
+# Optional array fields (not always present).  Pilot metadata is stored here -
+# as native npz arrays, exactly like the sample/symbol arrays - rather than in
+# the YAML meta block, so no array-to-list conversion is ever needed.
+_OPTIONAL_ARRAY_FIELDS: tuple[str, ...] = (
+    "source_bits",
+    "source_symbols",
+    "ps_pmf",
+    "pilot_tone_frequency",
+    "pilot_tone_power_ratio_db",
+)
+
+# Derived / cached array fields (only written when include_cache=True)
+_CACHE_FIELDS: tuple[str, ...] = ("resolved_symbols", "resolved_bits")
+
+
+# -----------------------------------------------------------------------------
+# Public API
+# -----------------------------------------------------------------------------
+
+
+def save_npz(
+    signal: Signal,
+    path: str | Path,
+    *,
+    compressed: bool = True,
+    include_cache: bool = False,
+) -> None:
+    """
+    Save a ``Signal`` to a NumPy archive (.npz).
+
+    Parameters
+    ----------
+    signal : Signal
+        The signal to persist.
+    path : str or Path
+        Destination path.  A ``.npz`` extension is appended automatically
+        if absent.
+    compressed : bool, default True
+        Use ``savez_compressed`` (zlib).  Set to ``False`` to use
+        the uncompressed ``savez`` (faster write, larger file).
+    include_cache : bool, default False
+        Also save ``resolved_symbols`` and ``resolved_bits`` if present.
+        These can be recomputed from the signal, so they are omitted by
+        default to keep file sizes small.
+
+    Notes
+    -----
+    When the signal was generated by ``to_signal``,
+    the originating frame is serialised into the archive
+    (``__frame_metadata__`` + ``frame_payload_symbols`` etc.).  On load,
+    ``load_npz`` reconstructs the frame and re-attaches it so that
+    ``correct_timing`` and frame-aware equalizer workflows work
+    identically to a freshly generated signal.
+
+    Examples
+    --------
+    >>> save_npz(sig, "capture.npz")
+    >>> save_npz(sig, "capture", compressed=False, include_cache=True)
+    """
+    path = Path(path)
+    if path.suffix != ".npz":
+        path = path.with_suffix(".npz")
+
+    # -------------------------------------------------------------------------
+    # Collect arrays
+    # -------------------------------------------------------------------------
+    arrays: dict[str, Any] = {"samples": _backend.to_device(signal.samples, "CPU")}
+
+    for field in _OPTIONAL_ARRAY_FIELDS:
+        arr = getattr(signal, field, None)
+        if arr is not None:
+            arrays[field] = _backend.to_device(arr, "CPU")
+
+    if include_cache:
+        for field in _CACHE_FIELDS:
+            arr = getattr(signal, field, None)
+            if arr is not None:
+                arrays[field] = _backend.to_device(arr, "CPU")
+
+    # -------------------------------------------------------------------------
+    # Serialise originating SingleCarrierFrame (if present)
+    # -------------------------------------------------------------------------
+    # The frame holds the complete transmit-side description (structure map,
+    # pilot masks, preamble sequence, generated payload/pilot symbols).
+    # Saving it ensures that correct_timing() and pilot extraction all work
+    # after a save/load round-trip without requiring the caller to keep a
+    # reference to the original frame object.
+    frame = signal.frame
+    if frame is not None:
+        # All public fields are JSON-serializable primitives; nested Preamble
+        # is a Pydantic model and is also captured by model_dump().
+        # _frame_type stores the class name so load_npz can reconstruct the
+        # correct type when multiple frame classes exist (SingleCarrierFrame,
+        # future OFDMFrame, etc.) without hardcoding the class.
+        frame_dict = frame.model_dump(mode="json")
+        frame_dict["_frame_type"] = type(frame).__name__
+        yaml_frame = yaml.dump(frame_dict, default_flow_style=False, allow_unicode=True)
+        arrays["__frame_metadata__"] = np.array(yaml_frame, dtype=object)
+
+        # Save the generated symbol/bit arrays that live in PrivateAttrs and
+        # are NOT reproduced by model_dump().  Payload symbols and bits are
+        # random; pilot symbols are deterministic but cheap to cache anyway.
+        for npz_key, frame_attr in (
+            ("frame_payload_symbols", "payload_symbols"),
+            ("frame_pilot_symbols", "pilot_symbols"),
+            ("frame_payload_bits", "payload_bits"),
+        ):
+            arr = getattr(frame, frame_attr, None)
+            if arr is not None:
+                arrays[npz_key] = _backend.to_device(arr, "CPU")
+
+    # -------------------------------------------------------------------------
+    # Build metadata dict and serialise to YAML
+    # -------------------------------------------------------------------------
+    meta: dict = {f: getattr(signal, f) for f in _META_FIELDS}
+    yaml_str = yaml.dump(meta, default_flow_style=False, allow_unicode=True)
+
+    # Store as a zero-d object array so np.savez treats it as a single entry
+    arrays["__metadata__"] = np.array(yaml_str, dtype=object)
+
+    # -------------------------------------------------------------------------
+    # Write
+    # -------------------------------------------------------------------------
+    if compressed:
+        np.savez_compressed(path, **arrays)  # type: ignore[arg-type]
+    else:
+        np.savez(path, **arrays)  # type: ignore[arg-type]
+
+
+def load_npz(
+    path: str | Path,
+    *,
+    device: str = "auto",
+) -> Signal:
+    """
+    Load a ``Signal`` from a .npz archive.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to the ``.npz`` file.  A ``.npz`` extension is appended
+        automatically if absent.
+    device : {"auto", "cpu", "gpu"}, default "auto"
+        Target device after loading.  ``"auto"`` moves to GPU when CuPy is
+        available, otherwise stays on CPU.
+
+    Returns
+    -------
+    Signal
+
+    Notes
+    -----
+    If the archive contains ``__frame_metadata__`` (written by
+    ``save_npz`` when the signal originated from a frame), the
+    originating frame is reconstructed (class determined by the
+    ``_frame_type`` key embedded in the metadata) and re-attached as
+    ``sig.frame``.  This makes ``correct_timing`` and
+    frame-aware equalizer workflows work without any extra arguments.
+
+    Examples
+    --------
+    >>> sig = load_npz("capture.npz")           # auto: GPU if available
+    >>> sig_cpu = load_npz("capture.npz", device="cpu")
+    >>> sig_gpu = load_npz("capture.npz", device="gpu")
+    """
+    from .core import Signal
+
+    path = Path(path)
+    if path.suffix != ".npz":
+        path = path.with_suffix(".npz")
+
+    # allow_pickle=True is required to read the zero-d object array that
+    # holds the YAML string; no arbitrary Python objects are loaded.
+    data = np.load(path, allow_pickle=True)
+
+    # -------------------------------------------------------------------------
+    # Parse YAML metadata
+    # -------------------------------------------------------------------------
+    yaml_str = str(data["__metadata__"])
+    meta: dict = yaml.safe_load(yaml_str)
+
+    # Build Signal constructor kwargs
+    # -------------------------------------------------------------------------
+    kwargs: dict = {f: meta.get(f) for f in _META_FIELDS}
+    kwargs["samples"] = data["samples"]
+
+    for field in _OPTIONAL_ARRAY_FIELDS:
+        if field in data:
+            kwargs[field] = data[field]
+
+    sig = Signal(**kwargs)
+
+    # -------------------------------------------------------------------------
+    # Restore cached arrays (bypass re-computation if present in file)
+    # -------------------------------------------------------------------------
+    for field in _CACHE_FIELDS:
+        if field in data:
+            setattr(sig, field, data[field])
+
+    # -------------------------------------------------------------------------
+    # Reconstruct originating frame (if serialised)
+    # -------------------------------------------------------------------------
+    if "__frame_metadata__" in data:
+        from . import core as _core
+
+        frame_dict = yaml.safe_load(str(data["__frame_metadata__"]))
+        frame_type_name = frame_dict.pop("_frame_type", "SingleCarrierFrame")
+
+        # Registry of known frame classes - extend here as new frame types land.
+        _FRAME_CLASSES = {
+            "SingleCarrierFrame": _core.SingleCarrierFrame,
+        }
+        frame_cls = _FRAME_CLASSES.get(frame_type_name)
+        if frame_cls is None:
+            raise ValueError(
+                f"Cannot reconstruct frame of type {frame_type_name!r}: "
+                "unknown frame class. Extend _FRAME_CLASSES in io.py."
+            )
+        frame = frame_cls(**frame_dict)
+
+        # Inject the cached symbol/bit arrays back into the frame's PrivateAttrs
+        # so that frame.payload_symbols, frame.pilot_symbols, frame.payload_bits
+        # return the original generated data without re-randomising.
+        if "frame_payload_symbols" in data:
+            frame._payload_symbols = data["frame_payload_symbols"]
+        if "frame_pilot_symbols" in data:
+            frame._pilot_symbols = data["frame_pilot_symbols"]
+        if "frame_payload_bits" in data:
+            frame._payload_bits = data["frame_payload_bits"]
+
+        sig.frame = frame
+
+        # _payload_ps_pmf is a PrivateAttr set during _ensure_payload_generated().
+        # When _payload_bits is restored above, that method returns early and never
+        # sets _payload_ps_pmf.  sig.ps_pmf was saved via _OPTIONAL_ARRAY_FIELDS and
+        # is already loaded, so restore from it directly.
+        if sig.ps_pmf is not None:
+            frame._payload_ps_pmf = sig.ps_pmf
+
+    # -------------------------------------------------------------------------
+    # Move to target device
+    # -------------------------------------------------------------------------
+    target = device.lower()
+    if target == "auto":
+        target = "gpu" if _backend.is_cupy_available() else "cpu"
+    sig = sig.to(target)
+    return sig

commkit/logger.py

@@ -0,0 +1,103 @@
+"""
+Logging utilities for the CommKit library.
+
+This module provides a unified, colorized logging interface for monitoring
+signal processing workflows and debugging complex system failures.
+"""
+
+import logging
+import sys
+
+
+class ColorFormatter(logging.Formatter):
+    """
+    Custom logging formatter providing ANSI-colored output based on log levels.
+
+    This formatter enhances readability by using distinct colors for different
+    severities (e.g., Cyan for DEBUG, Red for ERROR).
+    """
+
+    GREY = "\x1b[38;20m"
+    CYAN = "\x1b[36;20m"
+    GREEN = "\x1b[32;20m"
+    YELLOW = "\x1b[33;20m"
+    RED = "\x1b[31;20m"
+    BOLD_RED = "\x1b[31;1m"
+    RESET = "\x1b[0m"
+    FORMAT = "%(asctime)s [%(levelname)s] [%(name)s/%(filename)s] %(message)s"
+
+    LEVEL_COLORS = {
+        logging.DEBUG: CYAN,
+        logging.INFO: GREEN,
+        logging.WARNING: YELLOW,
+        logging.ERROR: RED,
+        logging.CRITICAL: BOLD_RED,
+    }
+
+    def format(self, record: logging.LogRecord) -> str:
+        """
+        Formats the log record with ANSI color codes.
+
+        Parameters
+        ----------
+        record : logging.LogRecord
+            The log record containing the message and metadata.
+
+        Returns
+        -------
+        str
+            The formatted log message with embedded ANSI escape sequences.
+        """
+        log_color = self.LEVEL_COLORS.get(record.levelno, self.RESET)
+        formatter = logging.Formatter(
+            f"{log_color}{self.FORMAT}{self.RESET}",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        return formatter.format(record)
+
+
+def get_logger(name: str = "commkit") -> logging.Logger:
+    """
+    Retrieves and configures a logger instance for the library.
+
+    If the requested logger has no handlers, a `StreamHandler` with the
+    `ColorFormatter` is automatically attached to ensure immediate visibility.
+
+    Parameters
+    ----------
+    name : str, default "commkit"
+        The namespace for the logger.
+
+    Returns
+    -------
+    logging.Logger
+        A configured logger instance.
+    """
+    logger = logging.getLogger(name)
+
+    if not logger.handlers:
+        logger.setLevel(logging.INFO)
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(ColorFormatter())
+        logger.addHandler(handler)
+
+    return logger
+
+
+# Create a default logger for the package
+logger = get_logger()
+
+
+def set_log_level(level):
+    """
+    Sets the global log level for the CommKit library.
+
+    Parameters
+    ----------
+    level : int or str
+        The logging level to apply. Accepts standard `logging` constants
+        (e.g., `logging.DEBUG`) or string identifiers (e.g., "DEBUG", "INFO").
+    """
+    if isinstance(level, str):
+        level = getattr(logging, level.upper())
+    logger.setLevel(level)

commkit/mapping/__init__.py

@@ -0,0 +1,46 @@
+"""
+Symbol mapping, demapping, and constellation management.
+
+This package provides high-performance routines for the transition between
+digital bits and physical IQ symbols.  It is organised by mathematical concern:
+
+- :mod:`~commkit.mapping.gray` - constellation geometry and Gray labelling.
+- :mod:`~commkit.mapping.bits` - hard bit mapping / demapping.
+- :mod:`~commkit.mapping.llr` - soft-decision (LLR) demapping.
+- :mod:`~commkit.mapping.shaping` - probabilistic shaping (PS-QAM).
+- :mod:`~commkit.mapping.constellation` - the :class:`Constellation` value
+  object bundling points + Gray labels + optional shaping pmf.
+
+The public import surface is stable: every name previously importable from the
+flat ``commkit.mapping`` module is re-exported here.  ``Constellation`` is an
+additive convenience over the existing loose-array free functions.
+
+Note: codes and constellations are generated using NumPy (host-side).
+"""
+
+from .bits import demap_symbols_hard, map_bits
+from .constellation import Constellation
+from .gray import gray_code, gray_constellation, gray_to_binary
+from .llr import compute_llr
+from .shaping import (
+    constellation_power,
+    maxwell_boltzmann,
+    optimal_nu,
+    ps_entropy,
+    sample_ps_symbols,
+)
+
+__all__ = [
+    "Constellation",
+    "compute_llr",
+    "constellation_power",
+    "demap_symbols_hard",
+    "gray_code",
+    "gray_constellation",
+    "gray_to_binary",
+    "map_bits",
+    "maxwell_boltzmann",
+    "optimal_nu",
+    "ps_entropy",
+    "sample_ps_symbols",
+]