stress_operator/utils/viz.py

"""Shared stress-field renderer (used by the Gate-0 sanity plot and the Gradio demo)."""
from __future__ import annotations

from typing import Optional

import numpy as np


def render_stress(
    coords: np.ndarray,
    sigma: np.ndarray,
    title: str = "von Mises stress",
    ax=None,
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
    cmap: str = "coolwarm",
):
    """Scatter the per-node stress over the point cloud. Returns the matplotlib Figure.

    coords: (N, 2) node coordinates; sigma: (N,) per-node stress.
    """
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots(figsize=(5, 5))
    else:
        fig = ax.figure
    sc = ax.scatter(
        coords[:, 0], coords[:, 1], c=sigma, s=18, cmap=cmap,
        vmin=vmin, vmax=vmax, edgecolor="w", lw=0.1,
    )
    ax.set_aspect("equal")
    ax.set_title(title)
    fig.colorbar(sc, ax=ax, shrink=0.8, label="sigma")
    return fig