Source code for vane.viz.campbell_plot

"""Interactive Campbell-diagram plotting.

Renders a :class:`~vane.campbell.builder.CampbellDiagram` as an interactive Plotly
figure: one line per tracked mode (natural frequency versus the operating
parameter), the ``nP`` excitation lines, and markers at the detected resonance
crossings.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import plotly.graph_objects as go

from vane.campbell.excitation import DEFAULT_HARMONICS, find_resonances

if TYPE_CHECKING:
    from collections.abc import Sequence

    from vane.campbell.builder import CampbellDiagram

__all__ = ["plot_campbell"]

_RPM_PARAMETER = "rotor_speed_rpm"




def plot_campbell(
    diagram: CampbellDiagram,
    *,
    harmonics: Sequence[int] = DEFAULT_HARMONICS,
    show_excitation: bool = True,
    show_resonances: bool = True,
) -> go.Figure:
    """Plot a Campbell diagram as an interactive figure.

    Parameters
    ----------
    diagram : CampbellDiagram
        The diagram to plot.
    harmonics : Sequence[int], optional
        Excitation harmonics for the ``nP`` lines and resonance markers (only used
        for rotor-speed diagrams).
    show_excitation : bool, optional
        Draw the ``nP`` excitation lines (rotor-speed diagrams only).
    show_resonances : bool, optional
        Mark detected resonance crossings (rotor-speed diagrams only).

    Returns
    -------
    plotly.graph_objects.Figure
        The Campbell-diagram figure.
    """
    figure = go.Figure()
    for track in diagram.tracks:
        x, frequency, _ = diagram.track_curve(track)
        figure.add_trace(
            go.Scatter(
                x=x,
                y=frequency,
                mode="lines+markers",
                name=track.label.label,
            )
        )

    is_rpm = diagram.parameter_name == _RPM_PARAMETER
    if show_excitation and is_rpm:
        _add_excitation_lines(figure, diagram, harmonics)
    if show_resonances and is_rpm:
        _add_resonance_markers(figure, diagram, harmonics)

    figure.update_layout(
        title="Campbell Diagram",
        xaxis_title=diagram.parameter_name,
        yaxis_title="Natural frequency (Hz)",
    )
    return figure



def _add_excitation_lines(
    figure: go.Figure, diagram: CampbellDiagram, harmonics: Sequence[int]
) -> None:
    """Add the nP excitation lines spanning the operating-parameter range."""
    values = diagram.parameter_values
    if values.shape[0] == 0:
        return
    span = [float(values.min()), float(values.max())]
    for harmonic in harmonics:
        figure.add_trace(
            go.Scatter(
                x=span,
                y=[harmonic * value / 60.0 for value in span],
                mode="lines",
                line={"dash": "dash", "color": "gray"},
                name=f"{harmonic}P",
            )
        )


def _add_resonance_markers(
    figure: go.Figure, diagram: CampbellDiagram, harmonics: Sequence[int]
) -> None:
    """Add markers at the detected resonance crossings."""
    crossings = find_resonances(diagram, harmonics)
    if not crossings:
        return
    figure.add_trace(
        go.Scatter(
            x=[crossing.rotor_speed_rpm for crossing in crossings],
            y=[crossing.frequency_hz for crossing in crossings],
            mode="markers",
            marker={"symbol": "x", "size": 11, "color": "red"},
            name="Resonance",
        )
    )