70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
from plotly.basedatatypes import BaseTraceType
|
|
from skfem import Mesh
|
|
|
|
|
|
# 将网格与标量场转为 plotly 三角形 traces + 布局，供 RL 环境实时渲染
|
|
def get_plotly_mesh_traces_and_layout(
|
|
mesh: Mesh,
|
|
scalars: np.ndarray,
|
|
title: str = "Mesh",
|
|
mesh_dimension: int = 2,
|
|
boundary: Optional[np.ndarray] = None,
|
|
) -> Tuple[List[BaseTraceType], Dict[str, Any]]:
|
|
vertices = mesh.p
|
|
triangles = mesh.t
|
|
n_elements = triangles.shape[1]
|
|
s = np.asarray(scalars, dtype=np.float64).flatten()
|
|
|
|
x_tri = vertices[0, triangles].T
|
|
y_tri = vertices[1, triangles].T
|
|
intensity_tri = s[triangles].T
|
|
|
|
vmin, vmax = s.min(), s.max()
|
|
|
|
traces = []
|
|
for elem_idx in range(n_elements):
|
|
x_e, y_e, s_e = x_tri[elem_idx], y_tri[elem_idx], intensity_tri[elem_idx]
|
|
traces.append(go.Scatter(
|
|
x=x_e.tolist() + [x_e[0]],
|
|
y=y_e.tolist() + [y_e[0]],
|
|
mode="lines",
|
|
fill="toself",
|
|
fillcolor=_get_color(float(np.mean(s_e)), vmin, vmax),
|
|
line=dict(color="black", width=0.5),
|
|
showlegend=False,
|
|
hoverinfo="skip",
|
|
))
|
|
|
|
if traces:
|
|
traces[0].marker = dict(
|
|
color=s.min(), colorscale="RdBu_r", showscale=True,
|
|
colorbar=dict(title="Solution"),
|
|
)
|
|
|
|
layout = {
|
|
"title": title,
|
|
"xaxis": {"title": "x", "scaleanchor": "y"},
|
|
"yaxis": {"title": "y"},
|
|
"showlegend": False,
|
|
}
|
|
if boundary is not None:
|
|
layout["xaxis"]["range"] = [boundary[0], boundary[2]]
|
|
layout["yaxis"]["range"] = [boundary[1], boundary[3]]
|
|
|
|
return traces, layout
|
|
|
|
|
|
# 标量值 → matplotlib RdBu_r 色表映射的 RGBA 字符串
|
|
def _get_color(value: float, vmin: float, vmax: float) -> str:
|
|
import matplotlib.cm as cm
|
|
import matplotlib.colors as mcolors
|
|
|
|
norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
|
|
rgba = cm.RdBu_r(norm(value))
|
|
r, g, b, a = rgba
|
|
return f"rgba({int(r * 255)},{int(g * 255)},{int(b * 255)},{a:.2f})"
|