commkit 1.0.0__py3-none-any.whl

commkit/plotting/theme.py

@@ -0,0 +1,203 @@
+"""Shared plotting theme, layout, and formatting helpers."""
+
+import matplotlib as mpl
+import matplotlib.font_manager as fm
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ..backend import to_device
+from ..logger import logger
+
+
+def apply_default_theme() -> None:
+    """
+    Applies the library's default visual theme to Matplotlib.
+
+    This theme configures publication-quality defaults, including:
+    - Roboto or standard sans-serif typography.
+    - Optimized figure dimensions and DPI.
+    - Consistent grid styling and axis formatting.
+    - Improved LaTeX math rendering.
+
+    Notes
+    -----
+    This function modifies `matplotlib.rcParams` globally. It is recommended
+    to call this at the start of a script or notebook for consistent styling.
+    """
+    logger.debug("Applying default plotting theme.")
+    try:
+        font_prop = fm.FontProperties(family="Roboto", weight="regular")
+        fm.findfont(font_prop, fallback_to_default=False)
+        font_name = "Roboto"
+    except ValueError:
+        font_name = "sans"
+        logger.warning("Roboto font not found, falling back to default sans-serif.")
+
+    plt.style.use("dark_background")
+
+    mpl.rcParams.update(
+        {
+            "figure.figsize": (7, 4.9),
+            "font.family": font_name,
+            "font.size": 12,
+            "legend.fontsize": 10,
+            "lines.linewidth": 2,
+            "axes.linewidth": 1,
+            "axes.grid": True,
+            "grid.alpha": 0.4,
+            "axes.titleweight": "bold",
+            "figure.constrained_layout.use": True,
+            "savefig.dpi": 300,
+            "xtick.direction": "in",
+            "ytick.direction": "in",
+            "xtick.major.width": 1,
+            "ytick.major.width": 1,
+            "xtick.major.size": 4,
+            "ytick.major.size": 4,
+            "xtick.minor.width": 1,
+            "ytick.minor.width": 1,
+            "xtick.minor.size": 2,
+            "ytick.minor.size": 2,
+            "xtick.top": True,
+            "ytick.right": True,
+        }
+    )
+
+    plt.rcParams["mathtext.fontset"] = "custom"
+    plt.rcParams["mathtext.rm"] = font_name
+    plt.rcParams["mathtext.it"] = f"{font_name}:italic"
+    plt.rcParams["mathtext.bf"] = f"{font_name}:bold"
+
+
+def _create_subplot_grid(num_axes: int, max_cols: int = 2) -> tuple[int, int]:
+    """
+    Computes a grid layout (rows, cols) for a given number of axes.
+
+    Limits the maximum number of columns to prevent excessively wide figures.
+
+    Parameters
+    ----------
+    num_axes : int
+        Total number of subplots required.
+    max_cols : int, default 2
+        Maximum allowed number of columns.
+
+    Returns
+    -------
+    nrows : int
+        Number of rows in the grid.
+    ncols : int
+        Number of columns in the grid.
+    """
+    if num_axes <= max_cols:
+        return 1, num_axes
+    ncols = max_cols
+    nrows = (num_axes + ncols - 1) // ncols  # Ceiling division
+    return nrows, ncols
+
+
+def _grid_figsize(
+    nrows: int, ncols: int, panel: tuple[float, float] | None = None
+) -> tuple[float, float]:
+    """
+    Figsize for a subplot grid that preserves the default single-panel size.
+
+    Tiles ``panel`` (default: ``rcParams["figure.figsize"]``) across the
+    grid, so a single row keeps the default height, a single column keeps
+    the default width, and an NxN grid keeps the default size per panel.
+
+    Parameters
+    ----------
+    nrows, ncols : int
+        Grid shape.
+    panel : (float, float), optional
+        Per-panel ``(width, height)`` in inches. Defaults to the theme's
+        ``figure.figsize``.
+
+    Returns
+    -------
+    (width, height) : tuple of float
+    """
+    w, h = panel if panel is not None else mpl.rcParams["figure.figsize"]
+    return (w * ncols, h * nrows)
+
+
+def _square_figsize(n: int = 1) -> tuple[float, float]:
+    """
+    Square panel size derived from the default figure height.
+
+    Used for constellation-style plots that require an equal aspect ratio,
+    where the default (non-square) ``figure.figsize`` doesn't apply directly.
+    Tiling ``n`` panels preserves the same per-panel size as ``_grid_figsize``.
+    """
+    _, h = mpl.rcParams["figure.figsize"]
+    return (h * n, h * n)
+
+
+def _decimate_minmax(
+    x: np.ndarray, y: np.ndarray, max_points: int = 4000
+) -> tuple[np.ndarray, np.ndarray]:
+    """Down-sample a line for fast plotting while preserving its envelope.
+
+    Buckets the data and keeps the per-bucket min **and** max (emitted in
+    x-order), so sharp features - spectral peaks, window edges, glitches -
+    survive, unlike plain striding.  Matplotlib renders every vertex, so for
+    long oversampled records (N ≳ 10⁵) plotting the raw trace is the dominant
+    cost; reducing to a few thousand points is visually identical but orders of
+    magnitude faster.
+
+    Parameters
+    ----------
+    x, y : np.ndarray
+        1-D arrays of equal length.
+    max_points : int, default 4000
+        Approximate cap on the number of plotted points.  Returned length is
+        ``≈ max_points`` (``2`` per bucket).  Pass ``<= 0`` to disable.
+
+    Returns
+    -------
+    (x_dec, y_dec) : tuple of np.ndarray
+        Decimated arrays, or the inputs unchanged when already short enough.
+    """
+    n = len(y)
+    if max_points <= 0 or n <= max_points:
+        return x, y
+
+    n_buckets = max(1, max_points // 2)
+    bucket = n // n_buckets
+    trimmed = n_buckets * bucket
+    yr = y[:trimmed].reshape(n_buckets, bucket)
+    xr = x[:trimmed].reshape(n_buckets, bucket)
+
+    rows = np.arange(n_buckets)
+    i_min = yr.argmin(axis=1)
+    i_max = yr.argmax(axis=1)
+    # Emit the two extrema per bucket in their original x-order so the line
+    # does not zig-zag backwards.
+    first_is_min = i_min <= i_max
+    x_out = np.empty(n_buckets * 2, dtype=x.dtype)
+    y_out = np.empty(n_buckets * 2, dtype=y.dtype)
+    x_out[0::2] = np.where(first_is_min, xr[rows, i_min], xr[rows, i_max])
+    y_out[0::2] = np.where(first_is_min, yr[rows, i_min], yr[rows, i_max])
+    x_out[1::2] = np.where(first_is_min, xr[rows, i_max], xr[rows, i_min])
+    y_out[1::2] = np.where(first_is_min, yr[rows, i_max], yr[rows, i_min])
+
+    if trimmed < n:  # keep the tail the reshape dropped
+        x_out = np.concatenate([x_out, x[trimmed:]])
+        y_out = np.concatenate([y_out, y[trimmed:]])
+    return x_out, y_out
+
+
+def _set_eng_formatter(ax, which: str, unit: str) -> None:
+    """Apply an engineering (SI-prefix) tick formatter to an axis."""
+    fmt = mpl.ticker.EngFormatter(unit=unit, sep=" ")
+    if which in ("x", "both"):
+        ax.xaxis.set_major_formatter(fmt)
+    if which in ("y", "both"):
+        ax.yaxis.set_major_formatter(fmt)
+
+
+def _as_channels(arr) -> np.ndarray:
+    """Bring to CPU and promote ``(N,)`` -> ``(1, N)``."""
+    arr = np.atleast_2d(np.asarray(to_device(arr, "cpu")))
+    return arr

commkit/plotting/waveform.py

@@ -0,0 +1,200 @@
+"""Time-domain waveform plots."""
+
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ..backend import dispatch, to_device
+from ..core.signal import Signal
+from ..logger import logger
+from .theme import (
+    _create_subplot_grid,
+    _grid_figsize,
+)
+
+
+def plot_time_domain(
+    samples: Any,
+    sampling_rate: float = 1.0,
+    start_symbol: int = 0,
+    num_symbols: int | None = None,
+    sps: float | None = None,
+    ax: Any | None = None,
+    title: str | None = "Waveform",
+    show: bool = False,
+    **kwargs: Any,
+) -> tuple[Any, Any] | None:
+    """
+    Plots the time-domain representation of the signal.
+
+    For complex signals, both In-Phase (I) and Quadrature (Q) components
+    are plotted. Handles SI scaling (s, ms, us, etc.) for the time axis.
+
+    Parameters
+    ----------
+    samples : array_like or Signal
+        Input signal samples. Shape: (..., N_samples).
+    sampling_rate : float, default 1.0
+        Sampling rate in Hz.
+    start_symbol : int, default 0
+        The starting symbol to plot.
+    num_symbols : int, optional
+        Limit plot to a specific number of symbol periods. Requires `sps`.
+    sps : float, optional
+        Samples per symbol (required if `num_symbols` is used).
+    ax : matplotlib.axes.Axes, optional
+        Existing axis to plot on.
+    title : str, optional
+        Plot title. Defaults to "Waveform".
+    show : bool, default False
+        If True, calls `plt.show()`.
+    **kwargs : Any
+        Additional keyword arguments passed to `ax.plot`.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The figure object.
+    ax : matplotlib.axes.Axes or ndarray
+        The axis or array of axes used for the plot.
+    """
+    if isinstance(samples, Signal):
+        sig = samples
+        return plot_time_domain(
+            sig.samples,
+            sampling_rate=sig.sampling_rate,
+            start_symbol=start_symbol,
+            num_symbols=num_symbols,
+            sps=sig.sps,
+            ax=ax,
+            title=title,
+            show=show,
+            **kwargs,
+        )
+
+    logger.debug("Generating time-domain plot.")
+
+    samples, xp, _ = dispatch(samples)
+
+    # Handle Multichannel
+    # Convention: (Channels, Time)
+    if samples.ndim > 1:
+        num_channels = samples.shape[0]
+
+        if ax is None:
+            nrows, ncols = _create_subplot_grid(num_channels)
+            fig, axes = plt.subplots(
+                nrows, ncols, figsize=_grid_figsize(nrows, ncols), squeeze=False
+            )
+        else:
+            if not isinstance(ax, (list, tuple, np.ndarray)):
+                logger.warning(
+                    "Multiple channels detected but single axis provided. Overlaying plots."
+                )
+                axes = np.array([[ax] * num_channels])
+                fig = ax.figure
+            else:
+                axes = np.atleast_2d(ax)
+                fig = axes.flat[0].figure
+
+        for i in range(num_channels):
+            channel_samples = samples[i]
+
+            # Determine target axis using 2D indexing
+            row, col = divmod(i, axes.shape[1])
+            target_ax = axes[row, col] if row < axes.shape[0] else axes.flat[-1]
+
+            ch_title = f"{title} (Ch {i})" if title else f"Channel {i}"
+
+            plot_time_domain(
+                channel_samples,
+                sampling_rate=sampling_rate,
+                start_symbol=start_symbol,
+                num_symbols=num_symbols,
+                sps=sps,
+                ax=target_ax,
+                title=ch_title,
+                show=False,
+                **kwargs,
+            )
+
+        if show:
+            plt.show()
+            return None
+        return fig, axes
+
+    # --- 1D Logic ---
+
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.figure
+
+    samples = to_device(samples, "cpu")
+
+    start_idx = int(start_symbol * sps) if sps is not None else int(start_symbol)
+
+    if num_symbols is not None and sps is not None:
+        limit = start_idx + int(num_symbols * sps)
+        if limit > len(samples):
+            limit = len(samples)
+            logger.warning(
+                "Limit exceeds number of symbols. Plotting up to last symbol."
+            )
+        plot_samples = samples[start_idx:limit]
+    else:
+        plot_samples = samples[start_idx:]
+
+    time_axis = np.arange(len(plot_samples)) / sampling_rate
+
+    # Auto-scale time axis
+    max_time = time_axis[-1] if len(time_axis) > 0 else 0
+    if max_time < 1e-9:
+        scale_factor = 1e12
+        unit = "ps"
+    elif max_time < 1e-6:
+        scale_factor = 1e9
+        unit = "ns"
+    elif max_time < 1e-3:
+        scale_factor = 1e6
+        unit = "µs"
+    elif max_time < 1:
+        scale_factor = 1e3
+        unit = "ms"
+    else:
+        scale_factor = 1.0
+        unit = "s"
+
+    time_axis = time_axis * scale_factor
+    xlabel = f"Time [{unit}]"
+
+    if np.iscomplexobj(plot_samples):
+        ax.plot(
+            time_axis,
+            plot_samples.real,
+            label="I",
+            **kwargs,
+        )
+        ax.plot(
+            time_axis,
+            plot_samples.imag,
+            label="Q",
+            **kwargs,
+        )
+        ax.legend()
+    else:
+        ax.plot(
+            time_axis,
+            plot_samples,
+            **kwargs,
+        )
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel("Amplitude")
+    if title is not None:
+        ax.set_title(title)
+
+    if show:
+        plt.show()
+        return None
+    return fig, ax

commkit/recovery/__init__.py

@@ -0,0 +1,51 @@
+"""
+Carrier phase recovery utilities.
+
+This package provides routines for carrier phase recovery (CPR), including
+streaming decision-directed PLL, block-based Viterbi-Viterbi, Blind Phase
+Search, MAP Tikhonov-RTS, and pilot-aided methods, along with cycle-slip
+correction and phase ambiguity resolution.
+
+The public API is unchanged from when this was a single module:
+``from commkit.recovery import recover_carrier_phase_bps, ...`` continues to
+work.
+"""
+
+from __future__ import annotations
+
+# Re-exported so ``patch("commkit.recovery.logger...")`` and similar
+# attribute access on the package namespace keep working.
+from ..logger import logger
+from .bps import recover_carrier_phase_bps
+from .corrections import (
+    correct_carrier_phase,
+    correct_cycle_slips,
+    correct_phase_rotation,
+    resolve_channel_permutation,
+    resolve_phase_ambiguity,
+    smooth_phase_wiener,
+)
+from .pilots import (
+    recover_carrier_phase_pilot_symbols,
+    recover_carrier_phase_pilot_tone,
+    recover_carrier_phase_pilot_tones,
+)
+from .pll import recover_carrier_phase_pll
+from .tikhonov import recover_carrier_phase_tikhonov
+from .viterbi_viterbi import recover_carrier_phase_viterbi_viterbi
+
+__all__ = [
+    "correct_carrier_phase",
+    "correct_cycle_slips",
+    "correct_phase_rotation",
+    "recover_carrier_phase_bps",
+    "recover_carrier_phase_pilot_tone",
+    "recover_carrier_phase_pilot_tones",
+    "recover_carrier_phase_pilot_symbols",
+    "recover_carrier_phase_pll",
+    "recover_carrier_phase_tikhonov",
+    "recover_carrier_phase_viterbi_viterbi",
+    "resolve_channel_permutation",
+    "resolve_phase_ambiguity",
+    "smooth_phase_wiener",
+]