LucasPlant/controlsim

1"""Basic re used plotting utils"""2 3import plotly.graph_objs as go4import numpy as np5 6MAX_PLOT_POINTS = 10007 8 9def get_plot_sample_indices(num_points: int, max_points: int = MAX_PLOT_POINTS) -> np.ndarray:10    """Return evenly spaced indices capped at max_points."""11    if num_points <= 0:12        return np.array([], dtype=int)13    if num_points <= max_points:14        return np.arange(num_points, dtype=int)15    return np.linspace(0, num_points - 1, max_points, dtype=int)16 17 18def multivar_plot(19    vars: np.ndarray, t: np.ndarray, state_info: list[str], title: str20) -> go.Figure:21    """22    A utility to easily plot multiple vars on the same axes good for state plots23    TODO Look into replacing this with dataframe based plotting for ease24 25    Args:26    vars: numpy array containing the vars over time to plot (time, var)27    t: the numpy array of timestamps28    state_info: list containing labels for the variables for the legend29    title: the title of the plot30    """31    num_points = min(len(t), vars.shape[0])32    sample_indices = get_plot_sample_indices(num_points)33    t_plot = t[:num_points][sample_indices]34    vars_plot = vars[:num_points][sample_indices]35 36    state_plot = go.Figure()37    plot_colors = [38        "#636EFA",39        "#EF553B",40        "#00CC96",41        "#AB63FA",42        "#FFA15A",43        "#19D3F3",44        "#FF6692",45        "#B6E880",46    ]47    num_vars = vars_plot.shape[1]48    for i in range(num_vars):49        state_name = state_info[i]50        axis_name = "y" if i == 0 else f"y{i + 1}"51        state_plot.add_trace(52            go.Scatter(53                x=t_plot,54                y=vars_plot[:, i],55                mode="lines",56                name=state_name,57                yaxis=axis_name,58                line={"color": plot_colors[i % len(plot_colors)]},59            )60        )61 62    layout_update = {63        "title": title,64        "xaxis_title": "Time (s)",65        "legend_title": "Variables",66        "legend": {67            "orientation": "h",68            "yanchor": "bottom",69            "y": 1.02,70            "xanchor": "left",71            "x": 0.0,72        },73        "yaxis": {74            "showline": True,75            "linecolor": plot_colors[0],76            "tickfont": {"color": plot_colors[0]},77            "zeroline": False,78            "nticks": 5,79            "automargin": True,80        },81        "margin": {"t": 100, "l": 80, "r": 80},82    }83 84    # Give each line its own overlaid y-axis so traces with different magnitudes85    # remain readable while sharing the same time axis.86    if num_vars > 1:87        left_extra = sum(1 for i in range(1, num_vars) if i % 2 == 0)88        right_extra = sum(1 for i in range(1, num_vars) if i % 2 == 1)89        left_margin = 80 + 45 * left_extra90        right_margin = 80 + 45 * right_extra91        x_start = 0.08 + 0.06 * left_extra92        x_end = 0.92 - 0.06 * right_extra93        if x_end - x_start < 0.5:94            x_start, x_end = 0.25, 0.7595 96        left_axis_idx = 097        right_axis_idx = 098        for i in range(1, num_vars):99            axis_key = f"yaxis{i + 1}"100            is_right = (i % 2) == 1101            if is_right:102                right_axis_idx += 1103                axis_position = min(0.99, x_end + 0.05 * right_axis_idx)104            else:105                left_axis_idx += 1106                axis_position = max(0.01, x_start - 0.05 * left_axis_idx)107 108            layout_update[axis_key] = {109                "overlaying": "y",110                "side": "right" if is_right else "left",111                "anchor": "free",112                "position": axis_position,113                "showgrid": False,114                "showline": True,115                "linecolor": plot_colors[i % len(plot_colors)],116                "tickfont": {"color": plot_colors[i % len(plot_colors)]},117                "zeroline": False,118                "nticks": 5,119                "automargin": True,120            }121 122        state_plot.update_layout(123            xaxis={124                "domain": [x_start, x_end],125                "title": "Time (s)",126            },127            margin={"t": 100, "l": left_margin, "r": right_margin},128        )129 130    state_plot.update_layout(**layout_update)131    return state_plot132 133 134def mode_plot(eigenvalues: np.ndarray, title: str) -> go.Figure:135    """136    Make a plot of the modes given the eigenvalues137 138    Args:139    eigenvalues: a ndarray containing the complex eigenvalues of the system140    title: the title of the plot141 142    Returns:143    the figure containing the plot144    """145    real_vals = np.real(eigenvalues)146    imag_vals = np.imag(eigenvalues)147 148    fig = go.Figure()149    fig.add_trace(150        go.Scatter(151            x=real_vals,152            y=imag_vals,153            mode="markers",154            marker=dict(155                size=12,156                color="blue",157                symbol="x",  # Use 'x' marker to match controls standards158            ),159            name="Eigenvalues",160        )161    )162    fig.update_layout(163        title=title,164        xaxis_title="Real Part",165        yaxis_title="Imaginary Part",166        showlegend=True,167        width=600,168        height=400,169    )170    # Add some padding to the axis limits for better visualization171    pad_x = (172        (real_vals.max() - real_vals.min()) * 0.1173        if real_vals.max() != real_vals.min()174        else 1175    )176    pad_y = (177        (imag_vals.max() - imag_vals.min()) * 0.1178        if imag_vals.max() != imag_vals.min()179        else 1180    )181    fig.update_xaxes(range=[real_vals.min() - pad_x, real_vals.max() + pad_x])182    fig.update_yaxes(range=[imag_vals.min() - pad_y, imag_vals.max() + pad_y])183    return fig184