oscura 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

oscura/visualization/signal_integrity.py

@@ -90,6 +90,31 @@ def plot_tdr(
         >>> fig = plot_tdr(z_profile, dist, z0=50, show=False)
         >>> fig.savefig("tdr_impedance.png")
     """
+    _validate_tdr_inputs(impedance, distance)
+    fig, ax = _setup_tdr_figure(ax, figsize)
+
+    distance_unit_final, dist_scaled = _scale_tdr_distance(distance, distance_unit)
+    impedance_display = np.clip(impedance, 0, 500)
+
+    _plot_tdr_impedance_profile(ax, dist_scaled, impedance_display)
+    _fill_tdr_impedance_regions(ax, dist_scaled, impedance_display, z0, discontinuity_threshold)
+
+    if show_reference:
+        _add_tdr_reference_line(ax, z0)
+
+    if show_discontinuities:
+        _annotate_tdr_discontinuities(ax, dist_scaled, impedance_display, discontinuity_threshold)
+
+    _format_tdr_axes(ax, dist_scaled, impedance_display, distance_unit_final, title)
+    _finalize_tdr_plot(fig, save_path, show)
+
+    return fig
+
+
+def _validate_tdr_inputs(
+    impedance: NDArray[np.floating[Any]], distance: NDArray[np.floating[Any]]
+) -> None:
+    """Validate TDR input arrays."""
     if not HAS_MATPLOTLIB:
         raise ImportError("matplotlib is required for visualization")
 
@@ -99,16 +124,23 @@ def plot_tdr(
             f"(got {len(impedance)} and {len(distance)})"
         )
 
-    # Create figure if needed
+
+def _setup_tdr_figure(ax: Axes | None, figsize: tuple[float, float]) -> tuple[Figure, Axes]:
+    """Setup TDR figure and axes."""
     if ax is None:
-        fig, ax = plt.subplots(figsize=figsize)
-    else:
-        fig_temp = ax.get_figure()
-        if fig_temp is None:
-            raise ValueError("Axes must have an associated figure")
-        fig = cast("Figure", fig_temp)
+        fig, ax_new = plt.subplots(figsize=figsize)
+        return fig, ax_new
+
+    fig_temp = ax.get_figure()
+    if fig_temp is None:
+        raise ValueError("Axes must have an associated figure")
+    return cast("Figure", fig_temp), ax
 
-    # Convert distance units
+
+def _scale_tdr_distance(
+    distance: NDArray[np.floating[Any]], distance_unit: str
+) -> tuple[str, NDArray[np.floating[Any]]]:
+    """Scale distance to appropriate unit."""
     if distance_unit == "auto":
         max_dist = np.max(distance)
         if max_dist < 0.01:
@@ -124,25 +156,29 @@ def plot_tdr(
         distance_mult = {"m": 1.0, "cm": 100.0, "mm": 1000.0}.get(distance_unit, 1.0)
 
     dist_scaled = distance * distance_mult
+    return distance_unit, dist_scaled
 
-    # Clip impedance for display (handle inf values)
-    impedance_display = np.clip(impedance, 0, 500)
 
-    # Plot impedance profile
+def _plot_tdr_impedance_profile(
+    ax: Axes,
+    dist_scaled: NDArray[np.floating[Any]],
+    impedance_display: NDArray[np.floating[Any]],
+) -> None:
+    """Plot main impedance profile line."""
     ax.plot(dist_scaled, impedance_display, "b-", linewidth=2, label="Impedance")
 
-    # Fill regions based on impedance deviation from z0
+
+def _fill_tdr_impedance_regions(
+    ax: Axes,
+    dist_scaled: NDArray[np.floating[Any]],
+    impedance_display: NDArray[np.floating[Any]],
+    z0: float,
+    discontinuity_threshold: float,
+) -> None:
+    """Fill colored regions based on impedance deviation."""
     for i in range(len(dist_scaled) - 1):
         z = impedance_display[i]
-        if z > z0 + discontinuity_threshold:
-            color = "#FFA500"  # Orange for high-Z
-            alpha = 0.3
-        elif z < z0 - discontinuity_threshold:
-            color = "#1E90FF"  # Blue for low-Z
-            alpha = 0.3
-        else:
-            color = "#90EE90"  # Light green for matched
-            alpha = 0.2
+        color, alpha = _get_tdr_region_color(z, z0, discontinuity_threshold)
 
         ax.fill_between(
             [dist_scaled[i], dist_scaled[i + 1]],
@@ -152,52 +188,85 @@ def plot_tdr(
             alpha=alpha,
         )
 
-    # Reference line
-    if show_reference:
-        ax.axhline(z0, color="gray", linestyle="--", linewidth=1.5, label=f"Z0 = {z0} Ω")
 
-    # Find and annotate discontinuities
-    if show_discontinuities:
-        # Find significant changes
-        z_diff = np.abs(np.diff(impedance_display))
-        discontinuities = np.where(z_diff > discontinuity_threshold)[0]
-
-        for idx in discontinuities:
-            z_before = impedance_display[idx]
-            z_after = impedance_display[idx + 1]
-            d = dist_scaled[idx]
-
-            # Determine discontinuity type
-            if z_after > z_before + discontinuity_threshold:
-                disc_type = "High-Z"
-                color = "orange"
-            elif z_after < z_before - discontinuity_threshold:
-                disc_type = "Low-Z"
-                color = "blue"
-            else:
-                continue
-
-            # Add marker
-            ax.plot(d, z_after, "o", color=color, markersize=8)
-
-            # Add annotation
-            z_str = f"{z_after:.0f}" if z_after < 500 else "Open"
-            ax.annotate(
-                f"{disc_type}\n{z_str} Ω",
-                xy=(d, z_after),
-                xytext=(10, 10),
-                textcoords="offset points",
-                fontsize=8,
-                ha="left",
-                bbox={"boxstyle": "round,pad=0.3", "facecolor": "white", "alpha": 0.8},
-            )
-
-    # Labels and formatting
+def _get_tdr_region_color(z: float, z0: float, threshold: float) -> tuple[str, float]:
+    """Get color and alpha for impedance region."""
+    if z > z0 + threshold:
+        return "#FFA500", 0.3  # Orange for high-Z
+    elif z < z0 - threshold:
+        return "#1E90FF", 0.3  # Blue for low-Z
+    else:
+        return "#90EE90", 0.2  # Light green for matched
+
+
+def _add_tdr_reference_line(ax: Axes, z0: float) -> None:
+    """Add reference impedance line."""
+    ax.axhline(z0, color="gray", linestyle="--", linewidth=1.5, label=f"Z0 = {z0} Ω")
+
+
+def _annotate_tdr_discontinuities(
+    ax: Axes,
+    dist_scaled: NDArray[np.floating[Any]],
+    impedance_display: NDArray[np.floating[Any]],
+    discontinuity_threshold: float,
+) -> None:
+    """Find and annotate impedance discontinuities."""
+    z_diff = np.abs(np.diff(impedance_display))
+    discontinuities = np.where(z_diff > discontinuity_threshold)[0]
+
+    for idx in discontinuities:
+        z_before = impedance_display[idx]
+        z_after = impedance_display[idx + 1]
+        d = dist_scaled[idx]
+
+        disc_type, color = _classify_tdr_discontinuity(z_before, z_after, discontinuity_threshold)
+        if disc_type is None:
+            continue
+
+        _add_tdr_discontinuity_marker(ax, d, z_after, disc_type, color)
+
+
+def _classify_tdr_discontinuity(
+    z_before: float, z_after: float, threshold: float
+) -> tuple[str | None, str]:
+    """Classify discontinuity type."""
+    if z_after > z_before + threshold:
+        return "High-Z", "orange"
+    elif z_after < z_before - threshold:
+        return "Low-Z", "blue"
+    return None, ""
+
+
+def _add_tdr_discontinuity_marker(
+    ax: Axes, d: float, z_after: float, disc_type: str, color: str
+) -> None:
+    """Add marker and annotation for discontinuity."""
+    ax.plot(d, z_after, "o", color=color, markersize=8)
+
+    z_str = f"{z_after:.0f}" if z_after < 500 else "Open"
+    ax.annotate(
+        f"{disc_type}\n{z_str} Ω",
+        xy=(d, z_after),
+        xytext=(10, 10),
+        textcoords="offset points",
+        fontsize=8,
+        ha="left",
+        bbox={"boxstyle": "round,pad=0.3", "facecolor": "white", "alpha": 0.8},
+    )
+
+
+def _format_tdr_axes(
+    ax: Axes,
+    dist_scaled: NDArray[np.floating[Any]],
+    impedance_display: NDArray[np.floating[Any]],
+    distance_unit: str,
+    title: str | None,
+) -> None:
+    """Format axes labels, limits, and title."""
     ax.set_xlabel(f"Distance ({distance_unit})", fontsize=11)
     ax.set_ylabel("Impedance (Ω)", fontsize=11)
     ax.set_xlim(0, dist_scaled[-1])
 
-    # Set y-axis limits with padding
     y_min = max(0, np.min(impedance_display) - 10)
     y_max = min(200, np.max(impedance_display) + 10)
     ax.set_ylim(y_min, y_max)
@@ -205,22 +274,73 @@ def plot_tdr(
     ax.grid(True, alpha=0.3)
     ax.legend(loc="upper right")
 
-    if title:
-        ax.set_title(title, fontsize=12, fontweight="bold")
-    else:
-        ax.set_title("TDR Impedance Profile", fontsize=12, fontweight="bold")
+    title_text = title if title else "TDR Impedance Profile"
+    ax.set_title(title_text, fontsize=12, fontweight="bold")
+
 
+def _finalize_tdr_plot(fig: Figure, save_path: str | Path | None, show: bool) -> None:
+    """Finalize plot layout, save, and show."""
     fig.tight_layout()
 
-    # Save if requested
     if save_path is not None:
         fig.savefig(save_path, dpi=300, bbox_inches="tight")
 
-    # Show if requested
     if show:
         plt.show()
 
-    return fig
+
+def _select_sparams_freq_unit(
+    frequencies: NDArray[np.floating[Any]], freq_unit: str
+) -> tuple[str, float]:
+    """Select frequency unit and divisor for S-parameter plots."""
+    if freq_unit == "auto":
+        max_freq = np.max(frequencies)
+        if max_freq >= 1e9:
+            return "GHz", 1e9
+        elif max_freq >= 1e6:
+            return "MHz", 1e6
+        elif max_freq >= 1e3:
+            return "kHz", 1e3
+        else:
+            return "Hz", 1.0
+    else:
+        freq_div = {"Hz": 1.0, "kHz": 1e3, "MHz": 1e6, "GHz": 1e9}.get(freq_unit, 1.0)
+        return freq_unit, freq_div
+
+
+def _convert_sparam_to_db(s: NDArray[Any]) -> NDArray[np.floating[Any]]:
+    """Convert S-parameter to dB."""
+    if np.iscomplexobj(s):
+        result: NDArray[np.floating[Any]] = 20 * np.log10(np.abs(s) + 1e-12)
+        return result
+    return np.asarray(s, dtype=np.float64)
+
+
+def _add_3db_marker(
+    ax: Axes,
+    freq_scaled: NDArray[np.floating[Any]],
+    s_db: NDArray[np.floating[Any]],
+    freq_unit: str,
+) -> None:
+    """Add -3dB bandwidth marker to S21 plot."""
+    max_db = np.max(s_db)
+    db_3_level = max_db - 3
+
+    crossings = np.where(np.diff(np.sign(s_db - db_3_level)))[0]
+    if len(crossings) > 0:
+        f_3db = float(freq_scaled[crossings[0]])
+        db_3_level_float = float(db_3_level)
+        ax.axhline(db_3_level_float, color="gray", linestyle=":", alpha=0.7, linewidth=1)
+        ax.axvline(f_3db, color="gray", linestyle=":", alpha=0.7, linewidth=1)
+        ax.plot(f_3db, db_3_level_float, "ko", markersize=6)
+        ax.annotate(
+            f"-3dB: {f_3db:.2f} {freq_unit}",
+            xy=(f_3db, db_3_level_float),
+            xytext=(10, -15),
+            textcoords="offset points",
+            fontsize=9,
+            ha="left",
+        )
 
 
 def plot_sparams_magnitude(
@@ -270,51 +390,17 @@ def plot_sparams_magnitude(
     if not HAS_MATPLOTLIB:
         raise ImportError("matplotlib is required for visualization")
 
-    # Create figure if needed
-    if ax is None:
-        fig, ax = plt.subplots(figsize=figsize)
-    else:
-        fig_temp = ax.get_figure()
-        if fig_temp is None:
-            raise ValueError("Axes must have an associated figure")
-        fig = cast("Figure", fig_temp)
-
-    # Select frequency unit
-    if freq_unit == "auto":
-        max_freq = np.max(frequencies)
-        if max_freq >= 1e9:
-            freq_unit = "GHz"
-            freq_div = 1e9
-        elif max_freq >= 1e6:
-            freq_unit = "MHz"
-            freq_div = 1e6
-        elif max_freq >= 1e3:
-            freq_unit = "kHz"
-            freq_div = 1e3
-        else:
-            freq_unit = "Hz"
-            freq_div = 1.0
-    else:
-        freq_div = {"Hz": 1.0, "kHz": 1e3, "MHz": 1e6, "GHz": 1e9}.get(freq_unit, 1.0)
-
+    fig, ax = _setup_tdr_figure(ax, figsize)
+    freq_unit, freq_div = _select_sparams_freq_unit(frequencies, freq_unit)
     freq_scaled = frequencies / freq_div
 
-    def to_db(s: NDArray[Any]) -> NDArray[np.floating[Any]]:
-        """Convert S-parameter to dB."""
-        if np.iscomplexobj(s):
-            result: NDArray[np.floating[Any]] = 20 * np.log10(np.abs(s) + 1e-12)
-            return result
-        return np.asarray(s, dtype=np.float64)
-
-    # Color scheme
     colors = {"S11": "#E74C3C", "S21": "#3498DB", "S12": "#2ECC71", "S22": "#9B59B6"}
     linestyles = {"S11": "-", "S21": "-", "S12": "--", "S22": "--"}
-
     params = [("S11", s11), ("S21", s21), ("S12", s12), ("S22", s22)]
 
     for name, s_param in params:
         if s_param is not None:
-            s_db = to_db(s_param)
+            s_db = _convert_sparam_to_db(s_param)
             ax.semilogx(
                 freq_scaled,
                 s_db,
@@ -324,46 +410,21 @@ def plot_sparams_magnitude(
                 label=name,
             )
 
-            # -3dB marker for S21
             if name == "S21" and db_3_marker:
-                max_db = np.max(s_db)
-                db_3_level = max_db - 3
-
-                # Find -3dB crossover
-                crossings = np.where(np.diff(np.sign(s_db - db_3_level)))[0]
-                if len(crossings) > 0:
-                    f_3db = float(freq_scaled[crossings[0]])
-                    db_3_level_float = float(db_3_level)
-                    ax.axhline(
-                        db_3_level_float, color="gray", linestyle=":", alpha=0.7, linewidth=1
-                    )
-                    ax.axvline(f_3db, color="gray", linestyle=":", alpha=0.7, linewidth=1)
-                    ax.plot(f_3db, db_3_level_float, "ko", markersize=6)
-                    ax.annotate(
-                        f"-3dB: {f_3db:.2f} {freq_unit}",
-                        xy=(f_3db, db_3_level_float),
-                        xytext=(10, -15),
-                        textcoords="offset points",
-                        fontsize=9,
-                        ha="left",
-                    )
-
-    # Labels and formatting
+                _add_3db_marker(ax, freq_scaled, s_db, freq_unit)
+
     ax.set_xlabel(f"Frequency ({freq_unit})", fontsize=11)
     ax.set_ylabel("Magnitude (dB)", fontsize=11)
     ax.grid(True, which="both", alpha=0.3)
     ax.legend(loc="best")
-
-    if title:
-        ax.set_title(title, fontsize=12, fontweight="bold")
-    else:
-        ax.set_title("S-Parameter Magnitude Response", fontsize=12, fontweight="bold")
+    ax.set_title(
+        title if title else "S-Parameter Magnitude Response", fontsize=12, fontweight="bold"
+    )
 
     fig.tight_layout()
 
     if save_path is not None:
         fig.savefig(save_path, dpi=300, bbox_inches="tight")
-
     if show:
         plt.show()
 
@@ -517,44 +578,127 @@ def plot_setup_hold_timing(
     if not HAS_MATPLOTLIB:
         raise ImportError("matplotlib is required for visualization")
 
-    # Create figure with multiple rows
+    # Create figure and axes
+    fig, axes = _create_timing_figure(ax, clock_data, figsize)
+
+    # Determine time scaling
+    time_unit_final, time_mult = _select_time_unit(time_unit, clock_edges, data_edges)
+    setup_scaled = setup_time * time_mult
+    hold_scaled = hold_time * time_mult
+
+    # Plot waveforms if provided
+    ax_timing = _plot_timing_waveforms(axes, clock_data, data_data, time_axis, time_mult)
+
+    # Setup timing annotation panel
+    _setup_timing_panel(ax_timing, clock_edges, data_edges, time_mult)
+
+    # Draw timing arrows
+    if show_margins:
+        _draw_timing_arrows(
+            ax_timing,
+            clock_edges,
+            data_edges,
+            setup_scaled,
+            hold_scaled,
+            time_mult,
+            time_unit_final,
+        )
+
+    # Add pass/fail status
+    _add_passfail_status(
+        ax_timing, setup_time, hold_time, setup_spec, hold_spec, time_mult, time_unit_final
+    )
+
+    # Finalize plot
+    axes[-1].set_xlabel(f"Time ({time_unit_final})", fontsize=11)
+    fig.suptitle(title if title else "Setup/Hold Timing Analysis", fontsize=14, fontweight="bold")
+    fig.tight_layout()
+
+    if save_path is not None:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    if show:
+        plt.show()
+
+    return fig
+
+
+def _create_timing_figure(
+    ax: Axes | None,
+    clock_data: NDArray[np.floating[Any]] | None,
+    figsize: tuple[float, float],
+) -> tuple[Figure, list[Any]]:
+    """Create figure and axes for timing diagram.
+
+    Args:
+        ax: Existing axes or None.
+        clock_data: Clock waveform data (determines row count).
+        figsize: Figure size.
+
+    Returns:
+        Tuple of (figure, axes_list).
+    """
     if ax is not None:
         fig_temp = ax.get_figure()
         if fig_temp is None:
             raise ValueError("Axes must have an associated figure")
-        fig = cast("Figure", fig_temp)
-        axes = [ax]
-        n_rows = 1
-    else:
-        n_rows = 3 if clock_data is not None else 1
-        fig, axes = plt.subplots(
-            n_rows, 1, figsize=figsize, sharex=True, gridspec_kw={"height_ratios": [1] * n_rows}
-        )
-        if n_rows == 1:
-            axes = [axes]
+        return cast("Figure", fig_temp), [ax]
+
+    n_rows = 3 if clock_data is not None else 1
+    fig, axes = plt.subplots(
+        n_rows, 1, figsize=figsize, sharex=True, gridspec_kw={"height_ratios": [1] * n_rows}
+    )
+    return fig, [axes] if n_rows == 1 else axes
+
+
+def _select_time_unit(
+    time_unit: str,
+    clock_edges: NDArray[np.floating[Any]],
+    data_edges: NDArray[np.floating[Any]],
+) -> tuple[str, float]:
+    """Select appropriate time unit and multiplier.
 
-    # Select time unit
+    Args:
+        time_unit: Requested unit or "auto".
+        clock_edges: Clock edge times.
+        data_edges: Data edge times.
+
+    Returns:
+        Tuple of (unit_string, multiplier).
+    """
     if time_unit == "auto":
         max_time = max(np.max(clock_edges), np.max(data_edges))
         if max_time < 1e-9:
-            time_unit = "ps"
-            time_mult = 1e12
+            return "ps", 1e12
         elif max_time < 1e-6:
-            time_unit = "ns"
-            time_mult = 1e9
+            return "ns", 1e9
         elif max_time < 1e-3:
-            time_unit = "us"
-            time_mult = 1e6
+            return "us", 1e6
         else:
-            time_unit = "ms"
-            time_mult = 1e3
+            return "ms", 1e3
     else:
-        time_mult = {"s": 1, "ms": 1e3, "us": 1e6, "ns": 1e9, "ps": 1e12}.get(time_unit, 1e9)
+        mult = {"s": 1, "ms": 1e3, "us": 1e6, "ns": 1e9, "ps": 1e12}.get(time_unit, 1e9)
+        return time_unit, mult
 
-    setup_scaled = setup_time * time_mult
-    hold_scaled = hold_time * time_mult
 
-    # If waveforms provided, plot them
+def _plot_timing_waveforms(
+    axes: list[Any],
+    clock_data: NDArray[np.floating[Any]] | None,
+    data_data: NDArray[np.floating[Any]] | None,
+    time_axis: NDArray[np.floating[Any]] | None,
+    time_mult: float,
+) -> Any:
+    """Plot clock and data waveforms.
+
+    Args:
+        axes: List of axes.
+        clock_data: Clock waveform.
+        data_data: Data waveform.
+        time_axis: Time axis.
+        time_mult: Time multiplier.
+
+    Returns:
+        Axes for timing annotations.
+    """
     if clock_data is not None and data_data is not None and time_axis is not None:
         time_scaled = time_axis * time_mult
 
@@ -578,88 +722,131 @@ def plot_setup_hold_timing(
         ax_data.set_yticks([0, 1])
         ax_data.grid(True, axis="x", alpha=0.3)
 
-        ax_timing = axes[2] if len(axes) > 2 else axes[-1]
+        return axes[2] if len(axes) > 2 else axes[-1]
     else:
-        ax_timing = axes[0]
-
-    # Timing annotation panel
-    ax_timing.set_ylim(0, 1)
-    ax_timing.set_xlim(0, max(clock_edges[-1], data_edges[-1]) * time_mult * 1.1)
-    ax_timing.axis("off")
-
-    # Draw timing arrows for first clock edge
-    if len(clock_edges) > 0 and len(data_edges) > 0:
-        clk_edge = clock_edges[0] * time_mult
-
-        # Find nearest data edge before clock
-        data_before = data_edges[data_edges < clock_edges[0]]
-        if len(data_before) > 0:
-            data_edge = data_before[-1] * time_mult
-
-            # Setup time arrow (data_edge to clk_edge)
-            if show_margins:
-                y_setup = 0.7
-                ax_timing.annotate(
-                    "",
-                    xy=(clk_edge, y_setup),
-                    xytext=(data_edge, y_setup),
-                    arrowprops={
-                        "arrowstyle": "<->",
-                        "color": "#27AE60",
-                        "lw": 2,
-                    },
-                )
-                ax_timing.text(
-                    (data_edge + clk_edge) / 2,
-                    y_setup + 0.1,
-                    f"Setup: {setup_scaled:.2f} {time_unit}",
-                    ha="center",
-                    va="bottom",
-                    fontsize=10,
-                    fontweight="bold",
-                    color="#27AE60",
-                )
-
-        # Find nearest data edge after clock
-        data_after = data_edges[data_edges > clock_edges[0]]
-        if len(data_after) > 0:
-            data_edge_after = data_after[0] * time_mult
-
-            # Hold time arrow (clk_edge to data_edge_after)
-            if show_margins:
-                y_hold = 0.3
-                ax_timing.annotate(
-                    "",
-                    xy=(data_edge_after, y_hold),
-                    xytext=(clk_edge, y_hold),
-                    arrowprops={
-                        "arrowstyle": "<->",
-                        "color": "#E67E22",
-                        "lw": 2,
-                    },
-                )
-                ax_timing.text(
-                    (clk_edge + data_edge_after) / 2,
-                    y_hold + 0.1,
-                    f"Hold: {hold_scaled:.2f} {time_unit}",
-                    ha="center",
-                    va="bottom",
-                    fontsize=10,
-                    fontweight="bold",
-                    color="#E67E22",
-                )
+        return axes[0]
 
-    # Add pass/fail status
+
+def _setup_timing_panel(
+    ax: Any,
+    clock_edges: NDArray[np.floating[Any]],
+    data_edges: NDArray[np.floating[Any]],
+    time_mult: float,
+) -> None:
+    """Setup timing annotation panel.
+
+    Args:
+        ax: Timing axes.
+        clock_edges: Clock edge times.
+        data_edges: Data edge times.
+        time_mult: Time multiplier.
+    """
+    ax.set_ylim(0, 1)
+    ax.set_xlim(0, max(clock_edges[-1], data_edges[-1]) * time_mult * 1.1)
+    ax.axis("off")
+
+
+def _draw_timing_arrows(
+    ax: Any,
+    clock_edges: NDArray[np.floating[Any]],
+    data_edges: NDArray[np.floating[Any]],
+    setup_scaled: float,
+    hold_scaled: float,
+    time_mult: float,
+    time_unit: str,
+) -> None:
+    """Draw setup and hold timing arrows.
+
+    Args:
+        ax: Timing axes.
+        clock_edges: Clock edge times.
+        data_edges: Data edge times.
+        setup_scaled: Scaled setup time.
+        hold_scaled: Scaled hold time.
+        time_mult: Time multiplier.
+        time_unit: Time unit string.
+    """
+    if len(clock_edges) == 0 or len(data_edges) == 0:
+        return
+
+    clk_edge = clock_edges[0] * time_mult
+
+    # Setup time arrow
+    data_before = data_edges[data_edges < clock_edges[0]]
+    if len(data_before) > 0:
+        data_edge = data_before[-1] * time_mult
+        y_setup = 0.7
+        ax.annotate(
+            "",
+            xy=(clk_edge, y_setup),
+            xytext=(data_edge, y_setup),
+            arrowprops={"arrowstyle": "<->", "color": "#27AE60", "lw": 2},
+        )
+        ax.text(
+            (data_edge + clk_edge) / 2,
+            y_setup + 0.1,
+            f"Setup: {setup_scaled:.2f} {time_unit}",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            fontweight="bold",
+            color="#27AE60",
+        )
+
+    # Hold time arrow
+    data_after = data_edges[data_edges > clock_edges[0]]
+    if len(data_after) > 0:
+        data_edge_after = data_after[0] * time_mult
+        y_hold = 0.3
+        ax.annotate(
+            "",
+            xy=(data_edge_after, y_hold),
+            xytext=(clk_edge, y_hold),
+            arrowprops={"arrowstyle": "<->", "color": "#E67E22", "lw": 2},
+        )
+        ax.text(
+            (clk_edge + data_edge_after) / 2,
+            y_hold + 0.1,
+            f"Hold: {hold_scaled:.2f} {time_unit}",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            fontweight="bold",
+            color="#E67E22",
+        )
+
+
+def _add_passfail_status(
+    ax: Any,
+    setup_time: float,
+    hold_time: float,
+    setup_spec: float | None,
+    hold_spec: float | None,
+    time_mult: float,
+    time_unit: str,
+) -> None:
+    """Add pass/fail status text.
+
+    Args:
+        ax: Timing axes.
+        setup_time: Measured setup time.
+        hold_time: Measured hold time.
+        setup_spec: Setup specification.
+        hold_spec: Hold specification.
+        time_mult: Time multiplier.
+        time_unit: Time unit string.
+    """
     status_y = 0.9
+
     if setup_spec is not None:
         setup_pass = setup_time >= setup_spec
         status = "PASS" if setup_pass else "FAIL"
         color = "#27AE60" if setup_pass else "#E74C3C"
-        ax_timing.text(
+        ax.text(
             0.02,
             status_y,
             f"Setup: {status} (spec: {setup_spec * time_mult:.2f} {time_unit})",
-            transform=ax_timing.transAxes,
+            transform=ax.transAxes,
             fontsize=10,
             color=color,
             fontweight="bold",
@@ -670,34 +857,16 @@ def plot_setup_hold_timing(
         hold_pass = hold_time >= hold_spec
         status = "PASS" if hold_pass else "FAIL"
         color = "#27AE60" if hold_pass else "#E74C3C"
-        ax_timing.text(
+        ax.text(
             0.02,
             status_y,
             f"Hold: {status} (spec: {hold_spec * time_mult:.2f} {time_unit})",
-            transform=ax_timing.transAxes,
+            transform=ax.transAxes,
             fontsize=10,
             color=color,
             fontweight="bold",
         )
 
-    # Set x-label on bottom axes
-    axes[-1].set_xlabel(f"Time ({time_unit})", fontsize=11)
-
-    if title:
-        fig.suptitle(title, fontsize=14, fontweight="bold")
-    else:
-        fig.suptitle("Setup/Hold Timing Analysis", fontsize=14, fontweight="bold")
-
-    fig.tight_layout()
-
-    if save_path is not None:
-        fig.savefig(save_path, dpi=300, bbox_inches="tight")
-
-    if show:
-        plt.show()
-
-    return fig
-
 
 def plot_timing_margin(
     setup_times: NDArray[np.floating[Any]],
@@ -735,6 +904,32 @@ def plot_timing_margin(
     if not HAS_MATPLOTLIB:
         raise ImportError("matplotlib is required for visualization")
 
+    fig, ax = _get_or_create_axes(ax, figsize)
+    time_mult = {"s": 1, "ms": 1e3, "us": 1e6, "ns": 1e9, "ps": 1e12}.get(time_unit, 1e9)
+
+    setup_scaled, hold_scaled = setup_times * time_mult, hold_times * time_mult
+
+    # Scatter plot
+    ax.scatter(setup_scaled, hold_scaled, c="#3498DB", alpha=0.6, s=50)
+
+    # Add specification lines and regions
+    _add_spec_lines(ax, setup_spec, hold_spec, time_mult, time_unit)
+    _add_pass_region(ax, setup_spec, hold_spec, time_mult)
+
+    # Configure axes
+    ax.set_xlabel(f"Setup Time ({time_unit})", fontsize=11)
+    ax.set_ylabel(f"Hold Time ({time_unit})", fontsize=11)
+    ax.grid(True, alpha=0.3)
+    ax.legend(loc="best")
+    ax.set_title(title or "Setup/Hold Timing Margin", fontsize=12, fontweight="bold")
+
+    fig.tight_layout()
+    _save_and_show_figure(fig, save_path, show)
+    return fig
+
+
+def _get_or_create_axes(ax: Axes | None, figsize: tuple[float, float]) -> tuple[Figure, Axes]:
+    """Get existing axes or create new figure with axes."""
     if ax is None:
         fig, ax = plt.subplots(figsize=figsize)
     else:
@@ -742,16 +937,13 @@ def plot_timing_margin(
         if fig_temp is None:
             raise ValueError("Axes must have an associated figure")
         fig = cast("Figure", fig_temp)
+    return fig, ax
 
-    time_mult = {"s": 1, "ms": 1e3, "us": 1e6, "ns": 1e9, "ps": 1e12}.get(time_unit, 1e9)
-
-    setup_scaled = setup_times * time_mult
-    hold_scaled = hold_times * time_mult
 
-    # Scatter plot
-    ax.scatter(setup_scaled, hold_scaled, c="#3498DB", alpha=0.6, s=50)
-
-    # Add specification lines if provided
+def _add_spec_lines(
+    ax: Axes, setup_spec: float | None, hold_spec: float | None, time_mult: float, time_unit: str
+) -> None:
+    """Add specification lines to timing margin plot."""
     if setup_spec is not None:
         spec_scaled = setup_spec * time_mult
         ax.axvline(
@@ -772,12 +964,13 @@ def plot_timing_margin(
             label=f"Hold Spec ({spec_scaled:.2f} {time_unit})",
         )
 
-    # Mark pass/fail regions
-    if setup_spec is not None and hold_spec is not None:
-        x_lim = ax.get_xlim()
-        y_lim = ax.get_ylim()
 
-        # Pass region (upper right)
+def _add_pass_region(
+    ax: Axes, setup_spec: float | None, hold_spec: float | None, time_mult: float
+) -> None:
+    """Add pass/fail region shading to timing margin plot."""
+    if setup_spec is not None and hold_spec is not None:
+        x_lim, y_lim = ax.get_xlim(), ax.get_ylim()
         ax.fill_between(
             [setup_spec * time_mult, x_lim[1]],
             [hold_spec * time_mult, hold_spec * time_mult],
@@ -787,22 +980,10 @@ def plot_timing_margin(
             label="Pass Region",
         )
 
-    ax.set_xlabel(f"Setup Time ({time_unit})", fontsize=11)
-    ax.set_ylabel(f"Hold Time ({time_unit})", fontsize=11)
-    ax.grid(True, alpha=0.3)
-    ax.legend(loc="best")
-
-    if title:
-        ax.set_title(title, fontsize=12, fontweight="bold")
-    else:
-        ax.set_title("Setup/Hold Timing Margin", fontsize=12, fontweight="bold")
-
-    fig.tight_layout()
 
+def _save_and_show_figure(fig: Figure, save_path: str | Path | None, show: bool) -> None:
+    """Save and optionally show figure."""
     if save_path is not None:
         fig.savefig(save_path, dpi=300, bbox_inches="tight")
-
     if show:
         plt.show()
-
-    return fig