bode.py

"""Bode plot: magnitude (dB) and phase (deg) vs frequency.

One function:

* ``bode_plot`` — dual-panel Bode diagram (magnitude + phase).

Accepts either a transfer-function callable or pre-computed arrays.

Example
-------
>>> import numpy as np
>>> from academic_plot import bode_plot, savefig
>>> # 2nd-order low-pass: H(s) = wn^2 / (s^2 + 2*z*wn*s + wn^2)
>>> wn, zeta = 1000.0, 0.3
>>> def tf(f):
...     s = 1j * 2 * np.pi * f
...     return wn**2 / (s**2 + 2*zeta*wn*s + wn**2)
>>> fig = bode_plot(tf, title="Low-Pass Filter")
>>> savefig(fig, "bode_demo")
"""

from __future__ import annotations

from typing import Callable

import matplotlib.pyplot as plt
import numpy as np

from .style import COLORS, FIGSIZES, GRID_ALPHA, GRID_LINEWIDTH, GRID_COLOR, savefig


def bode_plot(
    tf: Callable[[np.ndarray], complex] | None = None,
    *,
    freq: np.ndarray | None = None,
    mag: np.ndarray | None = None,
    phase: np.ndarray | None = None,
    title: str | None = None,
    mag_ylabel: str = "Magnitude (dB)",
    phase_ylabel: str = "Phase (deg)",
    freq_xlabel: str = "Frequency (Hz)",
    mag_color: str | None = None,
    phase_color: str | None = None,
    linewidth: float = 1.2,
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    figsize: tuple[float, float] = FIGSIZES["tall"],
) -> plt.Figure:
    """Bode plot with magnitude (top) and phase (bottom).

    Provide data in one of two ways:

    1. **Transfer function** — pass ``tf(f) -> complex`` where *f* is a
       1-D frequency array in Hz.  Magnitude (dB) and phase (deg) are
       computed automatically.
    2. **Pre-computed** — pass ``freq``, ``mag`` (dB), and ``phase`` (deg)
       arrays directly.

    Parameters
    ----------
    tf : callable or None
        ``tf(freq) -> complex`` transfer function.
    freq : np.ndarray or None
        Frequency points in Hz.  Defaults to ``np.logspace(-1, 5, 500)``.
    mag : np.ndarray or None
        Pre-computed magnitude in dB (only used when *tf* is ``None``).
    phase : np.ndarray or None
        Pre-computed phase in degrees (only used when *tf* is ``None``).
    title : str or None
        Super-title for the figure.
    mag_ylabel : str
        y-axis label for the magnitude subplot.
    phase_ylabel : str
        y-axis label for the phase subplot.
    freq_xlabel : str
        x-axis label (only on the bottom subplot).
    mag_color : str or None
        Colour of the magnitude curve.  Falls back to ``COLORS["blue"]``.
    phase_color : str or None
        Colour of the phase curve.  Falls back to ``COLORS["red"]``.
    linewidth : float
        Line thickness for both curves.
    show_grid : bool
        Show grid on both subplots.
    grid_alpha : float
        Grid opacity.
    grid_linewidth : float
        Grid thickness.
    grid_color : str
        Grid colour.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure

    Raises
    ------
    ValueError
        If neither *tf* nor all of *(freq, mag, phase)* are provided.
    """
    # ── Compute magnitude and phase from transfer function ─────────────
    if tf is not None:
        if freq is None:
            freq = np.logspace(-1, 5, 500)
        H = tf(freq)
        mag = 20.0 * np.log10(np.abs(H) + 1e-30)
        phase = np.degrees(np.unwrap(np.angle(H)))
    else:
        if freq is None or mag is None or phase is None:
            raise ValueError("Provide either (tf) or (freq, mag, phase).")

    mc = mag_color or COLORS["blue"]
    pc = phase_color or COLORS["red"]

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize, sharex=True)

    # ── Top panel: magnitude ───────────────────────────────────────────
    ax1.semilogx(freq, mag, color=mc, linewidth=linewidth)
    ax1.set_ylabel(mag_ylabel)
    if show_grid:
        ax1.grid(True, which="both", linewidth=grid_linewidth,
                 alpha=grid_alpha, color=grid_color)
    ax1.minorticks_on()

    # ── Bottom panel: phase ────────────────────────────────────────────
    ax2.semilogx(freq, phase, color=pc, linewidth=linewidth)
    ax2.set_ylabel(phase_ylabel)
    ax2.set_xlabel(freq_xlabel)
    if show_grid:
        ax2.grid(True, which="both", linewidth=grid_linewidth,
                 alpha=grid_alpha, color=grid_color)
    ax2.minorticks_on()

    if title:
        fig.suptitle(title, fontsize=10, y=0.98)

    fig.align_ylabels([ax1, ax2])
    fig.tight_layout(pad=0.3, rect=[0, 0, 1, 0.96] if title else None)
    return fig